Implement sorting of neighbor lists for PrecomputedNeighborhoodSearch (#138)

* Implement sorting of neighbor lists for `PrecomputedNeighborhoodSearch`

* Fix tests

* Fix

* Fix `freeze_neighborhood_search`

* Add comments

* Fix

Author: efaulhaber

src/PointNeighbors.jl

@@ -11,8 +11,8 @@ using LinearAlgebra: dot
 using Polyester: Polyester
 @reexport using StaticArrays: SVector
 
-include("vector_of_vectors.jl")
 include("util.jl")
+include("vector_of_vectors.jl")
 include("neighborhood_search.jl")
 include("nhs_trivial.jl")
 include("cell_lists/cell_lists.jl")

src/gpu.jl

@@ -30,7 +30,8 @@ function Adapt.adapt_structure(to, nhs::PrecomputedNeighborhoodSearch)
     neighborhood_search = Adapt.adapt_structure(to, nhs.neighborhood_search)
 
     return PrecomputedNeighborhoodSearch{ndims(nhs)}(neighbor_lists, search_radius,
-                                                     periodic_box, neighborhood_search)
+                                                     periodic_box, neighborhood_search,
+                                                     nhs.sort_neighbor_lists)
 end
 
 function Adapt.adapt_structure(to, cell_list::SpatialHashingCellList{NDIMS}) where {NDIMS}

src/nhs_precomputed.jl

@@ -3,7 +3,8 @@
                                          periodic_box = nothing, update_strategy = nothing,
                                          update_neighborhood_search = GridNeighborhoodSearch{NDIMS}(),
                                          backend = DynamicVectorOfVectors{Int32},
-                                         max_neighbors = max_neighbors(NDIMS))
+                                         max_neighbors = max_neighbors(NDIMS),
+                                         sort_neighbor_lists = true)
 
 Neighborhood search with precomputed neighbor lists. A list of all neighbors is computed
 for each point during initialization and update.
@@ -44,23 +45,29 @@ to strip the internal neighborhood search, which is not needed anymore.
 - `max_neighbors`: Maximum number of neighbors per particle. This will be used to
                    allocate the `DynamicVectorOfVectors`. It is not used with
                    other backends. The default is 64 in 2D and 324 in 3D.
+- `sort_neighbor_lists = true`: Whether to sort the neighbor lists after construction.
+                            This can improve cache hits on CPUs and improve coalesced
+                            memory access on GPUs.
 """
 struct PrecomputedNeighborhoodSearch{NDIMS, NL, ELTYPE, PB, NHS} <:
        AbstractNeighborhoodSearch
     neighbor_lists      :: NL
     search_radius       :: ELTYPE
     periodic_box        :: PB
     neighborhood_search :: NHS
+    sort_neighbor_lists :: Bool
 
     function PrecomputedNeighborhoodSearch{NDIMS}(neighbor_lists, search_radius,
                                                   periodic_box,
-                                                  update_neighborhood_search) where {NDIMS}
+                                                  update_neighborhood_search,
+                                                  sort_neighbor_lists) where {NDIMS}
         return new{NDIMS, typeof(neighbor_lists),
                    typeof(search_radius),
                    typeof(periodic_box),
                    typeof(update_neighborhood_search)}(neighbor_lists, search_radius,
                                                        periodic_box,
-                                                       update_neighborhood_search)
+                                                       update_neighborhood_search,
+                                                       sort_neighbor_lists)
     end
 end
 
@@ -73,11 +80,13 @@ function PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points =
                                                                                                          periodic_box,
                                                                                                          update_strategy),
                                               backend = DynamicVectorOfVectors{Int32},
-                                              max_neighbors = max_neighbors(NDIMS)) where {NDIMS}
+                                              max_neighbors = max_neighbors(NDIMS),
+                                              sort_neighbor_lists = true) where {NDIMS}
     neighbor_lists = construct_backend(backend, n_points, max_neighbors)
 
     PrecomputedNeighborhoodSearch{NDIMS}(neighbor_lists, search_radius,
-                                         periodic_box, update_neighborhood_search)
+                                         periodic_box, update_neighborhood_search,
+                                         sort_neighbor_lists)
 end
 
 # Default values for maximum neighbor count
@@ -111,7 +120,7 @@ function initialize!(search::PrecomputedNeighborhoodSearch,
     initialize!(neighborhood_search, x, y; parallelization_backend)
 
     initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                               parallelization_backend)
+                               parallelization_backend, search.sort_neighbor_lists)
 
     return search
 end
@@ -132,14 +141,14 @@ function update!(search::PrecomputedNeighborhoodSearch,
     # Skip update if both point sets are static
     if any(points_moving)
         initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                                   parallelization_backend)
+                                   parallelization_backend, search.sort_neighbor_lists)
     end
 
     return search
 end
 
 function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                                    parallelization_backend)
+                                    parallelization_backend, sort_neighbor_lists)
     # Initialize neighbor lists
     empty!(neighbor_lists)
     resize!(neighbor_lists, size(x, 2))
@@ -155,7 +164,8 @@ function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
 end
 
 function initialize_neighbor_lists!(neighbor_lists::DynamicVectorOfVectors,
-                                    neighborhood_search, x, y, parallelization_backend)
+                                    neighborhood_search, x, y, parallelization_backend,
+                                    sort_neighbor_lists)
     resize!(neighbor_lists, size(x, 2))
 
     # `Base.empty!.(neighbor_lists)`, but for all backends
@@ -168,6 +178,10 @@ function initialize_neighbor_lists!(neighbor_lists::DynamicVectorOfVectors,
                            parallelization_backend) do point, neighbor, _, _
         pushat!(neighbor_lists, point, neighbor)
     end
+
+    if sort_neighbor_lists
+        sorteach!(neighbor_lists)
+    end
 end
 
 @inline function foreach_neighbor(f, neighbor_system_coords,
@@ -225,5 +239,6 @@ end
     return PrecomputedNeighborhoodSearch{ndims(search)}(search.neighbor_lists,
                                                         search.search_radius,
                                                         search.periodic_box,
-                                                        nothing)
+                                                        nothing,
+                                                        search.sort_neighbor_lists)
 end

src/util.jl

@@ -80,7 +80,7 @@ For GPU arrays, the respective `KernelAbstractions.Backend` is returned.
 """
 @inline default_backend(::AbstractArray) = PolyesterBackend()
 @inline default_backend(x::AbstractGPUArray) = KernelAbstractions.get_backend(x)
-@inline default_backend(x::DynamicVectorOfVectors) = default_backend(x.backend)
+@inline default_backend(x::PermutedDimsArray) = default_backend(x.parent)
 
 """
     @threaded backend for ... end

src/vector_of_vectors.jl

@@ -23,6 +23,8 @@ end
 
 @inline Base.size(vov::DynamicVectorOfVectors) = (vov.length_[],)
 
+@inline default_backend(x::DynamicVectorOfVectors) = default_backend(x.backend)
+
 @inline function Base.getindex(vov::DynamicVectorOfVectors, i)
     (; backend, lengths) = vov
 
@@ -162,6 +164,40 @@ end
     return vov
 end
 
+# Sort each inner vector
+@inline function sorteach!(vov::DynamicVectorOfVectors)
+    # TODO remove this check when Metal supports sorting
+    if nameof(typeof(default_backend(vov.backend))) == :MetalBackend
+        @warn "sorting neighbor lists is not supported on Metal. Skipping sort."
+        return vov
+    end
+
+    # Note that we cannot just do `sort!(vov[i])` on GPUs because that would call `sort!`
+    # from within a GPU kernel, but this function is not GPU-compatible.
+    # We might be able to use a sorting function from AcceleratedKernels.jl,
+    # but for now the following workaround should be sufficient.
+
+    # Set all unused entries to `typemax` so that they are sorted to the end
+    @threaded default_backend(vov.backend) for i in axes(vov.backend, 2)
+        for j in (vov.lengths[i] + 1):size(vov.backend, 1)
+            @inbounds vov.backend[j, i] = typemax(eltype(vov.backend))
+        end
+    end
+
+    # Now we can sort full columns.
+    # Note that this will forward to highly optimized sorting functions on GPUs.
+    # It currently does not work on Metal.
+    sort!(vov.backend, dims = 1)
+
+    return vov
+end
+
+@inline function sorteach!(vov::Vector{<:Vector{T}}) where {T}
+    sort!.(vov)
+
+    return vov
+end
+
 function max_inner_length(cells::DynamicVectorOfVectors, fallback)
     return size(cells.backend, 1)
 end

test/vector_of_vectors.jl

@@ -8,9 +8,6 @@
         vov = PointNeighbors.DynamicVectorOfVectors{ELTYPE}(max_outer_length = 20,
                                                             max_inner_length = 4)
 
-        # Test internal size
-        @test size(vov.backend) == (4, 20)
-
         function verify(vov, vov_ref)
             @test length(vov) == length(vov_ref)
             @test eachindex(vov) == eachindex(vov_ref)
@@ -24,62 +21,92 @@
             end
         end
 
-        # Initial check
-        verify(vov, vov_ref)
+        @testset "Initial state" begin
+            # Test internal size
+            @test size(vov.backend) == (4, 20)
+
+            # Initial check
+            verify(vov, vov_ref)
+        end
+
+        @testset "First push!" begin
+            push!(vov_ref, type.([1, 3, 2]))
+            push!(vov, type.([1, 3, 2]))
+
+            verify(vov, vov_ref)
+        end
+
+        @testset "push! multiple items" begin
+            push!(vov_ref, type.([4]), type.([5, 6, 7, 8]))
+            push!(vov, type.([4]), type.([5, 6, 7, 8]))
+
+            verify(vov, vov_ref)
+        end
+
+        @testset "push! to an inner vector" begin
+            push!(vov_ref[1], type(12))
+            PointNeighbors.pushat!(vov, 1, type(12))
 
-        # First `push!`
-        push!(vov_ref, type.([1, 2, 3]))
-        push!(vov, type.([1, 2, 3]))
+            verify(vov, vov_ref)
+        end
 
-        verify(vov, vov_ref)
+        @testset "push! overflow" begin
+            error_ = ErrorException("cell list is full. Use a larger `max_points_per_cell`.")
+            @test_throws error_ PointNeighbors.pushat!(vov, 1, type(13))
 
-        # `push!` multiple items
-        push!(vov_ref, type.([4]), type.([5, 6, 7, 8]))
-        push!(vov, type.([4]), type.([5, 6, 7, 8]))
+            verify(vov, vov_ref)
+        end
 
-        verify(vov, vov_ref)
+        @testset "deleteat!" begin
+            # Delete entry of inner vector. Note that this changes the order of the elements.
+            deleteat!(vov_ref[3], 2)
+            PointNeighbors.deleteatat!(vov, 3, 2)
 
-        # `push!` to an inner vector
-        push!(vov_ref[1], type(12))
-        PointNeighbors.pushat!(vov, 1, type(12))
+            @test vov_ref[3] == type.([5, 7, 8])
+            @test vov[3] == type.([5, 8, 7])
 
-        # `push!` overflow
-        error_ = ErrorException("cell list is full. Use a larger `max_points_per_cell`.")
-        @test_throws error_ PointNeighbors.pushat!(vov, 1, type(13))
+            # Delete second to last entry
+            deleteat!(vov_ref[3], 2)
+            PointNeighbors.deleteatat!(vov, 3, 2)
 
-        verify(vov, vov_ref)
+            @test vov_ref[3] == type.([5, 8])
+            @test vov[3] == type.([5, 7])
 
-        # Delete entry of inner vector. Note that this changes the order of the elements.
-        deleteat!(vov_ref[3], 2)
-        PointNeighbors.deleteatat!(vov, 3, 2)
+            # Delete last entry
+            deleteat!(vov_ref[3], 2)
+            PointNeighbors.deleteatat!(vov, 3, 2)
 
-        @test vov_ref[3] == type.([5, 7, 8])
-        @test vov[3] == type.([5, 8, 7])
+            # Now they are identical again
+            verify(vov, vov_ref)
 
-        # Delete second to last entry
-        deleteat!(vov_ref[3], 2)
-        PointNeighbors.deleteatat!(vov, 3, 2)
+            # Delete the remaining entry of this vector
+            deleteat!(vov_ref[3], 1)
+            PointNeighbors.deleteatat!(vov, 3, 1)
 
-        @test vov_ref[3] == type.([5, 8])
-        @test vov[3] == type.([5, 7])
+            verify(vov, vov_ref)
+        end
 
-        # Delete last entry
-        deleteat!(vov_ref[3], 2)
-        PointNeighbors.deleteatat!(vov, 3, 2)
+        # Skip for Tuples
+        if ELTYPE <: Number
+            @testset "sorteach!" begin
+                # Make sure that the first inner vector is unsorted.
+                # If this fails, make sure the tests above don't yield a sorted vector.
+                @test vov[1] != sort(vov[1])
 
-        # Now they are identical again
-        verify(vov, vov_ref)
+                PointNeighbors.sorteach!(vov)
+                PointNeighbors.sorteach!(vov_ref)
 
-        # Delete the remaining entry of this vector
-        deleteat!(vov_ref[3], 1)
-        PointNeighbors.deleteatat!(vov, 3, 1)
+                @test vov[1] == sort(vov[1])
 
-        verify(vov, vov_ref)
+                verify(vov, vov_ref)
+            end
+        end
 
-        # `empty!`
-        empty!(vov_ref)
-        empty!(vov)
+        @testset "empty!" begin
+            empty!(vov_ref)
+            empty!(vov)
 
-        verify(vov, vov_ref)
+            verify(vov, vov_ref)
+        end
     end
 end