Add warning and separate aligned function for the vloada extract_smatrix

Author: efaulhaber

src/general/abstract_system.jl

@@ -84,7 +84,18 @@ end
 
 # Optimized version for 2D, which uses SIMD.jl to combine the 4 loads of the 2x2 matrix
 # into a single wide load. This is significantly faster on GPUs than 4 individual loads.
-@inline function extract_smatrix(A, ::Val{2}, particle)
+# WARNING:
+# 1. This only works if the matrix elements are stored contiguously in memory.
+#    The 4 elements of the 2x2 matrix for a particle must immediately follow the
+#    4 elements of the previous particle's matrix, with no padding in between.
+# 2. The pointer of `A` must be aligned to the size of `Vec{4, eltype(A)}`.
+#    This is guaranteed if `A` is allocated by Julia and has the correct size,
+#    but may not be true if `A` is a view or subarray.
+@propagate_inbounds function extract_smatrix_aligned(A, system, particle)
+    return extract_smatrix_aligned(A, Val(ndims(system)), particle)
+end
+
+@inline function extract_smatrix_aligned(A, ::Val{2}, particle)
     @boundscheck checkbounds(A, 2, 2, particle)
 
     # Note that this doesn't work in 3D because it requires a stride of 2^n.
@@ -93,6 +104,11 @@ end
     return SMatrix{2, 2}(Tuple(x))
 end
 
+# Fall back to the generic version when not in 2D.
+@propagate_inbounds function extract_smatrix_aligned(A, ::Val{NDIMS}, particle) where {NDIMS}
+    return extract_smatrix(A, Val(NDIMS), particle)
+end
+
 # Specifically get the current coordinates of a particle for all system types.
 @propagate_inbounds function current_coords(u, system, particle)
     return extract_svector(current_coordinates(u, system), system, particle)

src/general/corrections.jl

@@ -395,7 +395,9 @@ end
 
 function correction_matrix_inversion_step!(corr_matrix, system, semi)
     @threaded semi for particle in eachparticle(system)
-        L = extract_smatrix(corr_matrix, system, particle)
+        # `corr_matrix` is not a view, so we can use the fast `extract_smatrix_aligned`
+        # instead of the generic `extract_smatrix`.
+        L = extract_smatrix_aligned(corr_matrix, system, particle)
 
         # The matrix `L` only becomes singular when the particle and all neighbors
         # are collinear (in 2D) or lie all in the same plane (in 3D).

src/schemes/structure/total_lagrangian_sph/system.jl

@@ -356,9 +356,13 @@ end
 end
 
 @propagate_inbounds function deformation_gradient(system, particle)
+    # `deformation_grad` is not a view, so we can use the fast `extract_smatrix_aligned`
+    # instead of the generic `extract_smatrix`.
     extract_smatrix(system.deformation_grad, system, particle)
 end
 @propagate_inbounds function pk1_rho2(system, particle)
+    # `pk1_rho2` is not a view, so we can use the fast `extract_smatrix_aligned`
+    # instead of the generic `extract_smatrix`.
     extract_smatrix(system.pk1_rho2, system, particle)
 end