Use contiguous memory layout with PrecomputedNeighborhoodSearch

Author: efaulhaber

src/cell_lists/cell_lists_util.jl

@@ -12,12 +12,12 @@ end
 
 # We need the prod() because FullGridCellList's size is a tuple of cells per dimension whereas
 # SpatialHashingCellList's size is an Integer for the number of cells in total.
-function construct_backend(::Type{<:AbstractCellList}, ::Type{Vector{Vector{T}}}, size,
+function construct_backend(::Any, ::Type{Vector{Vector{T}}}, size,
                            max_points_per_cell) where {T}
     return [T[] for _ in 1:prod(size)]
 end
 
-function construct_backend(::Type{<:AbstractCellList}, ::Type{DynamicVectorOfVectors{T}},
+function construct_backend(::Any, ::Type{DynamicVectorOfVectors{T}},
                            size,
                            max_points_per_cell) where {T}
     cells = DynamicVectorOfVectors{T}(max_outer_length = prod(size),
@@ -31,7 +31,7 @@ end
 # `DynamicVectorOfVectors{T}`, but a type `DynamicVectorOfVectors{T1, T2, T3, T4}`.
 # While `A{T} <: A{T1, T2}`, this doesn't hold for the types.
 # `Type{A{T}} <: Type{A{T1, T2}}` is NOT true.
-function construct_backend(cell_list::Type{<:AbstractCellList},
+function construct_backend(cell_list::Any,
                            ::Type{DynamicVectorOfVectors{T1, T2, T3, T4}}, size,
                            max_points_per_cell) where {T1, T2, T3, T4}
     return construct_backend(cell_list, DynamicVectorOfVectors{T1}, size,

src/cell_lists/full_grid.jl

@@ -24,7 +24,8 @@ See [`copy_neighborhood_search`](@ref) for more details.
 - `backend = DynamicVectorOfVectors{Int32}`: Type of the data structure to store the actual
     cell lists. Can be
     - `Vector{Vector{Int32}}`: Scattered memory, but very memory-efficient.
-    - `DynamicVectorOfVectors{Int32}`: Contiguous memory, optimizing cache-hits.
+    - `DynamicVectorOfVectors{Int32}`: Contiguous memory, optimizing cache-hits
+                                       and GPU compatible.
 - `max_points_per_cell = 100`: Maximum number of points per cell. This will be used to
                                allocate the `DynamicVectorOfVectors`. It is not used with
                                the `Vector{Vector{Int32}}` backend.

src/gpu.jl

@@ -10,6 +10,8 @@
 Adapt.@adapt_structure FullGridCellList
 Adapt.@adapt_structure SpatialHashingCellList
 Adapt.@adapt_structure DynamicVectorOfVectors
+Adapt.@adapt_structure GridNeighborhoodSearch
+Adapt.@adapt_structure PrecomputedNeighborhoodSearch
 
 # `adapt(CuArray, ::SVector)::SVector`, but `adapt(Array, ::SVector)::Vector`.
 # We don't want to change the type of the `SVector` here.
@@ -22,13 +24,3 @@ end
 function Adapt.adapt_structure(to::typeof(Array), range::UnitRange)
     return range
 end
-
-function Adapt.adapt_structure(to, nhs::GridNeighborhoodSearch)
-    (; search_radius, periodic_box, n_cells, cell_size) = nhs
-
-    cell_list = Adapt.adapt_structure(to, nhs.cell_list)
-    update_buffer = Adapt.adapt_structure(to, nhs.update_buffer)
-
-    return GridNeighborhoodSearch(cell_list, search_radius, periodic_box, n_cells,
-                                  cell_size, update_buffer, nhs.update_strategy)
-end

src/nhs_precomputed.jl

@@ -22,71 +22,93 @@ initialization and update.
                             [`PeriodicBox`](@ref).
 - `update_strategy`:        Strategy to parallelize `update!` of the internally used
                             `GridNeighborhoodSearch`. See [`GridNeighborhoodSearch`](@ref)
-                            for available options.
+                            for available options. This is only used for the default value
+                            of `update_neighborhood_search` below.
+- `update_neighborhood_search = GridNeighborhoodSearch{NDIMS}(; periodic_box, update_strategy)`:
+                            The neighborhood search used to compute the neighbor lists.
+                            By default, a [`GridNeighborhoodSearch`](@ref) is used.
+- `backend = DynamicVectorOfVectors{Int32}`: Type of the data structure to store
+    the neighbor lists. Can be
+    - `Vector{Vector{Int32}}`: Scattered memory, but very memory-efficient.
+    - `DynamicVectorOfVectors{Int32}`: Contiguous memory, optimizing cache-hits
+                                       and GPU-compatible.
+- `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.
 """
-struct PrecomputedNeighborhoodSearch{NDIMS, NHS, NL, PB} <: AbstractNeighborhoodSearch
-    neighborhood_search :: NHS
+struct PrecomputedNeighborhoodSearch{NDIMS, NL, ELTYPE, PB, NHS} <:
+       AbstractNeighborhoodSearch
     neighbor_lists      :: NL
+    search_radius       :: ELTYPE
     periodic_box        :: PB
+    neighborhood_search :: NHS
 
     function PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
                                                   periodic_box = nothing,
-                                                  update_strategy = nothing) where {NDIMS}
-        nhs = GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
-                                            periodic_box, update_strategy)
-
-        neighbor_lists = Vector{Vector{Int}}()
-
-        new{NDIMS, typeof(nhs),
-            typeof(neighbor_lists),
-            typeof(periodic_box)}(nhs, neighbor_lists, periodic_box)
+                                                  update_strategy = nothing,
+                                                  update_neighborhood_search = GridNeighborhoodSearch{NDIMS}(;
+                                                                                                             search_radius,
+                                                                                                             n_points,
+                                                                                                             periodic_box,
+                                                                                                             update_strategy),
+                                                  backend = DynamicVectorOfVectors{Int32},
+                                                  max_neighbors = 4 * NDIMS^4) where {NDIMS}
+        neighbor_lists = construct_backend(nothing, backend, n_points, max_neighbors)
+
+        new{NDIMS, typeof(neighbor_lists),
+            typeof(search_radius), typeof(periodic_box),
+            typeof(update_neighborhood_search)}(neighbor_lists, search_radius,
+                                                periodic_box, update_neighborhood_search)
     end
 end
 
 @inline Base.ndims(::PrecomputedNeighborhoodSearch{NDIMS}) where {NDIMS} = NDIMS
 
 @inline requires_update(::PrecomputedNeighborhoodSearch) = (true, true)
 
-@inline function search_radius(search::PrecomputedNeighborhoodSearch)
-    return search_radius(search.neighborhood_search)
-end
-
 function initialize!(search::PrecomputedNeighborhoodSearch,
                      x::AbstractMatrix, y::AbstractMatrix;
                      parallelization_backend = default_backend(x),
                      eachindex_y = axes(y, 2))
     (; neighborhood_search, neighbor_lists) = search
 
+    if eachindex_y != axes(y, 2)
+        error("this neighborhood search does not support inactive points")
+    end
+
     # Initialize grid NHS
-    initialize!(neighborhood_search, x, y; eachindex_y, parallelization_backend)
+    initialize!(neighborhood_search, x, y; parallelization_backend)
 
     initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                               parallelization_backend, eachindex_y)
+                               parallelization_backend)
+
+    return search
 end
 
-# WARNING! Experimental feature:
-# By default, determine the parallelization backend from the type of `x`.
-# Optionally, pass a `KernelAbstractions.Backend` to run the KernelAbstractions.jl code
-# on this backend. This can be useful to run GPU kernels on the CPU by passing
-# `parallelization_backend = KernelAbstractions.CPU()`, even though `x isa Array`.
 function update!(search::PrecomputedNeighborhoodSearch,
                  x::AbstractMatrix, y::AbstractMatrix;
                  points_moving = (true, true), parallelization_backend = default_backend(x),
                  eachindex_y = axes(y, 2))
     (; neighborhood_search, neighbor_lists) = search
 
-    # Update grid NHS
-    update!(neighborhood_search, x, y; eachindex_y, points_moving, parallelization_backend)
+    if eachindex_y != axes(y, 2)
+        error("this neighborhood search does not support inactive points")
+    end
+
+    # Update the internal neighborhood search
+    update!(neighborhood_search, x, y; points_moving, parallelization_backend)
 
     # Skip update if both point sets are static
     if any(points_moving)
         initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                                   parallelization_backend, eachindex_y)
+                                   parallelization_backend)
     end
+
+    return search
 end
 
 function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
-                                    parallelization_backend, eachindex_y)
+                                    parallelization_backend)
     # Initialize neighbor lists
     empty!(neighbor_lists)
     resize!(neighbor_lists, size(x, 2))
@@ -95,8 +117,19 @@ function initialize_neighbor_lists!(neighbor_lists, neighborhood_search, x, y,
     end
 
     # Fill neighbor lists
-    foreach_point_neighbor(x, y, neighborhood_search; parallelization_backend,
-                           points = eachindex_y) do point, neighbor, _, _
+    foreach_point_neighbor(x, y, neighborhood_search;
+                           parallelization_backend) do point, neighbor, _, _
+        push!(neighbor_lists[point], neighbor)
+    end
+end
+
+function initialize_neighbor_lists!(neighbor_lists::DynamicVectorOfVectors,
+                                    neighborhood_search, x, y, parallelization_backend)
+    resize!(neighbor_lists, size(x, 2))
+
+    # Fill neighbor lists
+    foreach_point_neighbor(x, y, neighborhood_search;
+                           parallelization_backend) do point, neighbor, _, _
         push!(neighbor_lists[point], neighbor)
     end
 end
@@ -132,8 +165,23 @@ end
 
 function copy_neighborhood_search(nhs::PrecomputedNeighborhoodSearch,
                                   search_radius, n_points; eachpoint = 1:n_points)
-    update_strategy_ = nhs.neighborhood_search.update_strategy
+    update_neighborhood_search = copy_neighborhood_search(nhs.neighborhood_search,
+                                                          search_radius, n_points;
+                                                          eachpoint)
+    max_neighbors = max_inner_length(nhs.neighbor_lists, 4 * ndims(nhs)^4)
     return PrecomputedNeighborhoodSearch{ndims(nhs)}(; search_radius, n_points,
                                                      periodic_box = nhs.periodic_box,
-                                                     update_strategy = update_strategy_)
+                                                     update_neighborhood_search,
+                                                     backend = typeof(nhs.neighbor_lists),
+                                                     max_neighbors)
+end
+
+# TODO move to `vector_of_vectors.jl`
+function max_inner_length(cells::DynamicVectorOfVectors, fallback)
+    return size(cells.backend, 1)
+end
+
+# Fallback when backend is a `Vector{Vector{T}}`. Only used for copying the cell list.
+function max_inner_length(::Any, fallback)
+    return fallback
 end

test/neighborhood_search.jl

@@ -59,6 +59,9 @@
                                                                            backend = Vector{Vector{Int32}})),
                 PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                                      periodic_box = periodic_boxes[i]),
+                PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                                     periodic_box = periodic_boxes[i],
+                                                     backend = Vector{Vector{Int32}}),
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                               periodic_box = periodic_boxes[i],
                                               cell_list = SpatialHashingCellList{NDIMS}(2 *
@@ -71,6 +74,7 @@
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`",
                 "`PrecomputedNeighborhoodSearch`",
+                "`PrecomputedNeighborhoodSearch` with `Vector{Vector}`",
                 "`GridNeighborhoodSearch` with `SpatialHashingCellList`"
             ]
 
@@ -86,6 +90,8 @@
                                                                            max_corner = periodic_boxes[i].max_corner,
                                                                            backend = Vector{Vector{Int32}})),
                 PrecomputedNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i]),
+                PrecomputedNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i],
+                                                     backend = Vector{Vector{Int32}}),
                 GridNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i],
                                               cell_list = SpatialHashingCellList{NDIMS}(2 *
                                                                                         n_points))
@@ -193,6 +199,8 @@
                                                                            search_radius,
                                                                            backend = Vector{Vector{Int}})),
                 PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points),
+                PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                                     backend = Vector{Vector{Int}}),
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                               cell_list = SpatialHashingCellList{NDIMS}(2 *
                                                                                         n_points))
@@ -207,6 +215,7 @@
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors` and `SemiParallelUpdate`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`",
                 "`PrecomputedNeighborhoodSearch`",
+                "`PrecomputedNeighborhoodSearch` with `Vector{Vector}`",
                 "`GridNeighborhoodSearch` with `SpatialHashingCellList`"
             ]
 
@@ -227,6 +236,7 @@
                                                                            max_corner,
                                                                            backend = Vector{Vector{Int32}})),
                 PrecomputedNeighborhoodSearch{NDIMS}(),
+                PrecomputedNeighborhoodSearch{NDIMS}(backend = Vector{Vector{Int32}}),
                 GridNeighborhoodSearch{NDIMS}(cell_list = SpatialHashingCellList{NDIMS}(2 *
                                                                                         n_points))
             ]