diff --git a/Project.toml b/Project.toml index c030c895..3143ac4b 100644 --- a/Project.toml +++ b/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" diff --git a/src/PointNeighbors.jl b/src/PointNeighbors.jl index faa8cefc..be176753 100644 --- a/src/PointNeighbors.jl +++ b/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 diff --git a/src/neighborhood_search.jl b/src/neighborhood_search.jl index 30356554..b1e7aee9 100644 --- a/src/neighborhood_search.jl +++ b/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 diff --git a/src/nhs_grid.jl b/src/nhs_grid.jl index 3b2f4602..e3eb0004 100644 --- a/src/nhs_grid.jl +++ b/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 diff --git a/src/nhs_precomputed.jl b/src/nhs_precomputed.jl index d228ab03..1b2c0c8a 100644 --- a/src/nhs_precomputed.jl +++ b/src/nhs_precomputed.jl @@ -205,12 +205,23 @@ 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 @@ -218,7 +229,7 @@ end 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.