Add SIMD vectorization to mapreduce_neighbor_unsafe

Project.toml

@@ -8,6 +8,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+LLVMLoopInfo = "8b046642-f1f6-4319-8d3c-209ddc03c586"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -18,6 +19,7 @@ Adapt = "4"
 Atomix = "1"
 GPUArraysCore = "0.2"
 KernelAbstractions = "0.9"
+LLVMLoopInfo = "1.0.0"
 LinearAlgebra = "1"
 Polyester = "0.7.5"
 Reexport = "1"

src/PointNeighbors.jl

@@ -8,6 +8,7 @@ using Base: @propagate_inbounds
 using GPUArraysCore: AbstractGPUArray
 using KernelAbstractions: KernelAbstractions, @kernel, @index
 using LinearAlgebra: dot
+using LLVMLoopInfo: @loopinfo
 using Polyester: Polyester
 @reexport using StaticArrays: SVector
 

src/neighborhood_search.jl

@@ -271,7 +271,7 @@ end
                                   point, point_coords, search_radius)
     mapreduce_neighbor_inner(f, foreach_neighbor_op,
                              neighbor_coords, neighborhood_search,
-                             point, point_coords, search_radius, nothing)
+                             point, point_coords, search_radius, nothing, Val(false))
     return nothing
 end
 
@@ -318,7 +318,8 @@ Note that all these bounds checks are safe to skip if
 
     @inbounds mapreduce_neighbor_inner(f, foreach_neighbor_op,
                                        neighbor_coords, neighborhood_search,
-                                       point, point_coords, search_radius, nothing)
+                                       point, point_coords, search_radius, nothing,
+                                       Val(false))
     return nothing
 end
 
@@ -353,7 +354,7 @@ end
                                     neighborhood_search::AbstractNeighborhoodSearch,
                                     point, point_coords, search_radius, init)
     mapreduce_neighbor_inner(f, op, neighbor_coords, neighborhood_search,
-                             point, point_coords, search_radius, init)
+                             point, point_coords, search_radius, init, Val(false))
 end
 
 """
@@ -374,12 +375,13 @@ and when it is safe to skip them.
 @inline function mapreduce_neighbor_unsafe(f, op, system_coords, neighbor_coords,
                                            neighborhood_search::AbstractNeighborhoodSearch,
                                            point; init,
-                                           search_radius = search_radius(neighborhood_search))
+                                           search_radius = search_radius(neighborhood_search),
+                                           simd = Val(false))
     point_coords = @inbounds extract_svector(system_coords, Val(ndims(neighborhood_search)),
                                              point)
 
     @inbounds mapreduce_neighbor_inner(f, op, neighbor_coords, neighborhood_search,
-                                       point, point_coords, search_radius, init)
+                                       point, point_coords, search_radius, init, simd)
 end
 
 # This is the generic function that is called for `TrivialNeighborhoodSearch`.
@@ -390,7 +392,7 @@ end
 @propagate_inbounds function mapreduce_neighbor_inner(f, op, neighbor_coords,
                                                       neighborhood_search::AbstractNeighborhoodSearch,
                                                       point, point_coords,
-                                                      search_radius, init)
+                                                      search_radius, init, _)
     (; periodic_box) = neighborhood_search
 
     reduced = init

src/nhs_grid.jl

@@ -519,7 +519,7 @@ end
 @propagate_inbounds function mapreduce_neighbor_inner(f, op, neighbor_coords,
                                                       neighborhood_search::GridNeighborhoodSearch,
                                                       point, point_coords,
-                                                      search_radius, init)
+                                                      search_radius, init, _)
     (; cell_list, periodic_box) = neighborhood_search
     cell = cell_coords(point_coords, neighborhood_search)
     reduced = init

src/nhs_precomputed.jl

@@ -205,20 +205,31 @@ function initialize_neighbor_lists!(neighbor_lists::DynamicVectorOfVectors,
     end
 end
 
+macro optional_loopinfo(simd, loop)
+    return esc(quote
+                   if $simd
+                       @loopinfo vectorwidth=8 predicate $loop
+                   else
+                       $loop
+                   end
+               end)
+end
+
 # Note that calling this function with `@inbounds` is not safe.
 # See the comments in `foreach_neighbor_unsafe`.
 @propagate_inbounds function mapreduce_neighbor_inner(f, op, neighbor_coords,
                                                       neighborhood_search::PrecomputedNeighborhoodSearch,
                                                       point, point_coords,
-                                                      search_radius, init)
+                                                      search_radius, init,
+                                                      ::Val{SIMD}) where {SIMD}
     (; periodic_box, neighbor_lists) = neighborhood_search
 
     # Making the following `@inbounds` is not safe because the neighbor list
     # might not contain `point` if the NHS was not initialized correctly.
     neighbors = neighbor_lists[point]
     reduced = init
 
-    for neighbor_ in eachindex(neighbors)
+    @optional_loopinfo SIMD for neighbor_ in eachindex(neighbors)
         neighbor = @inbounds neighbors[neighbor_]
 
         # Making this `@inbounds` is not safe because
@@ -235,7 +246,8 @@ end
          distance2) = compute_periodic_distance(pos_diff, distance2, search_radius,
                                                 periodic_box)
 
-        distance = sqrt(distance2)
+        # We need `@fastmath` here for performance when SIMD-vectorizing this loop.
+        distance = @fastmath sqrt(distance2)
 
         # Inline to avoid loss of performance compared to not using this function
         # and unrolling everything.