Add Backend for SpatialHashingCellList (#111)

* Add support for different backends in SpatialHashingCellList.

* Minor change, remove unnecessary call to prod().

* Add check_cell_bounds() for push_cell() for SHCL.
Move check_cell_bounds() and construct_backends() to separate cell_lists_util.jl to share functionality between different cell lists.

* Add constructor to SPCL to work with Adapt.jl

* Minor changes.

* Add SHCL to gpu.jl

* Resolve requested changes:
- Move functions from cell_lists_util.jl to cell_lists.jl
- Change dispatch in `supported_update_strategies()`
- Change doc string of SHCL (SpatialHashingCellList)

* Resolve requested changes:
- Add `@inbounds` in `push_cell_atomic!`
- Improved type dispatch for `supported_update_strategies`
- Clarified and cleaned up cell list initialization and emptying,
- General code cleanup.

* Resolve requested changes.

* Replace hashing function for the coordinates.
Make cell_list.coords primitive to use Atomix.

* Add `adapt_structure()` for SHCL.

* Use Atomix for updating cell's coords,

* Minor changes, runs on GPU.

* Add test for coordinate hash function.

* Minor changes.

* Resolve requested changes.
- Update the tests
- Improve the documentation
- Minor changes

---------

Co-authored-by: Erik Faulhaber <[email protected]>

Author: RubberLanding

src/cell_lists/cell_lists.jl

@@ -4,6 +4,47 @@ abstract type AbstractCellList end
 # able to be used in a threaded loop.
 @inline each_cell_index_threadable(cell_list::AbstractCellList) = each_cell_index(cell_list)
 
+@inline function check_cell_bounds(cell_list::AbstractCellList, cell::Integer)
+    (; cells) = cell_list
+
+    checkbounds(cells, cell)
+end
+
+function construct_backend(::Type{Vector{Vector{T}}},
+                           max_outer_length,
+                           max_inner_length) where {T}
+    return [T[] for _ in 1:max_outer_length]
+end
+
+function construct_backend(::Type{DynamicVectorOfVectors{T}},
+                           max_outer_length,
+                           max_inner_length) where {T}
+    cells = DynamicVectorOfVectors{T}(max_outer_length = max_outer_length,
+                                      max_inner_length = max_inner_length)
+    resize!(cells, max_outer_length)
+
+    return cells
+end
+
+# When `typeof(cell_list.cells)` is passed, we don't pass the type
+# `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(::Type{DynamicVectorOfVectors{T1, T2, T3, T4}}, max_outer_length,
+                           max_inner_length) where {T1, T2, T3, T4}
+    return construct_backend(DynamicVectorOfVectors{T1}, max_outer_length,
+                             max_inner_length)
+end
+
+function max_points_per_cell(cells::DynamicVectorOfVectors)
+    return size(cells.backend, 1)
+end
+
+# Fallback when backend is a `Vector{Vector{T}}`. Only used for copying the cell list.
+function max_points_per_cell(cells)
+    return 100
+end
+
 include("dictionary.jl")
 include("full_grid.jl")
 include("spatial_hashing.jl")

src/cell_lists/full_grid.jl

@@ -62,34 +62,13 @@ function FullGridCellList(; min_corner, max_corner,
         n_cells_per_dimension = ceil.(Int, (max_corner .- min_corner) ./ search_radius)
         linear_indices = LinearIndices(Tuple(n_cells_per_dimension))
 
-        cells = construct_backend(backend, n_cells_per_dimension, max_points_per_cell)
+        cells = construct_backend(backend, prod(n_cells_per_dimension),
+                                  max_points_per_cell)
     end
 
     return FullGridCellList(cells, linear_indices, min_corner, max_corner)
 end
 
-function construct_backend(::Type{Vector{Vector{T}}}, size, max_points_per_cell) where {T}
-    return [T[] for _ in 1:prod(size)]
-end
-
-function construct_backend(::Type{DynamicVectorOfVectors{T}}, size,
-                           max_points_per_cell) where {T}
-    cells = DynamicVectorOfVectors{T}(max_outer_length = prod(size),
-                                      max_inner_length = max_points_per_cell)
-    resize!(cells, prod(size))
-
-    return cells
-end
-
-# When `typeof(cell_list.cells)` is passed, we don't pass the type
-# `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(::Type{DynamicVectorOfVectors{T1, T2, T3, T4}}, size,
-                           max_points_per_cell) where {T1, T2, T3, T4}
-    return construct_backend(DynamicVectorOfVectors{T1}, size, max_points_per_cell)
-end
-
 @inline function cell_coords(coords, periodic_box::Nothing, cell_list::FullGridCellList,
                              cell_size)
     (; min_corner) = cell_list
@@ -210,15 +189,6 @@ function copy_cell_list(cell_list::FullGridCellList, search_radius, periodic_box
                             max_points_per_cell = max_points_per_cell(cell_list.cells))
 end
 
-function max_points_per_cell(cells::DynamicVectorOfVectors)
-    return size(cells.backend, 1)
-end
-
-# Fallback when backend is a `Vector{Vector{T}}`. Only used for copying the cell list.
-function max_points_per_cell(cells)
-    return 100
-end
-
 @inline function check_cell_bounds(cell_list::FullGridCellList{<:DynamicVectorOfVectors{<:Any,
                                                                                         <:Array}},
                                    cell::Tuple)
@@ -246,9 +216,3 @@ end
         error("particle coordinates are NaN or outside the domain bounds of the cell list")
     end
 end
-
-@inline function check_cell_bounds(cell_list::FullGridCellList, cell::Integer)
-    (; cells) = cell_list
-
-    checkbounds(cells, cell)
-end

src/cell_lists/spatial_hashing.jl

@@ -1,5 +1,7 @@
 """
-    SpatialHashingCellList{NDIMS}(; list_size)
+    SpatialHashingCellList{NDIMS}(; list_size, 
+                                  backend = DynamicVectorOfVectors{Int32},
+                                  max_points_per_cell = 100)
 
 A basic spatial hashing implementation. Similar to [`DictionaryCellList`](@ref), the domain is discretized into cells, 
 and the particles in each cell are stored in a hash map. The hash is computed using the spatial location of each cell, 
@@ -9,80 +11,132 @@ to balance memory consumption against the likelihood of hash collisions.
 
 # Arguments
 - `NDIMS::Int`: Number of spatial dimensions (e.g., `2` or `3`).
-- `list_size::Int`: Size of the hash map (e.g., `2 * n_points`) .
+- `list_size::Int`: Size of the hash map (e.g., `2 * n_points`).
+- `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.
+- `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.
 """
 
 struct SpatialHashingCellList{NDIMS, CL, CI, CF} <: AbstractCellList
-    points     :: CL
+    cells      :: CL
     coords     :: CI
     collisions :: CF
     list_size  :: Int
+
+    # This constructor is necessary for Adapt.jl to work with this struct
+    function SpatialHashingCellList(NDIMS, cells, coords, collisions, list_size)
+        return new{NDIMS, typeof(cells),
+                   typeof(coords), typeof(collisions)}(cells, coords,
+                                                       collisions, list_size)
+    end
 end
 
 @inline index_type(::SpatialHashingCellList) = Int32
 
 @inline Base.ndims(::SpatialHashingCellList{NDIMS}) where {NDIMS} = NDIMS
 
+function supported_update_strategies(::SpatialHashingCellList{<:Any,
+                                                              <:DynamicVectorOfVectors})
+    return (ParallelUpdate, SerialUpdate)
+end
+
 function supported_update_strategies(::SpatialHashingCellList)
     return (SerialUpdate,)
 end
 
-function SpatialHashingCellList{NDIMS}(list_size) where {NDIMS}
-    points = [Int[] for _ in 1:list_size]
+function SpatialHashingCellList{NDIMS}(; list_size,
+                                       backend = DynamicVectorOfVectors{Int32},
+                                       max_points_per_cell = 100) where {NDIMS}
+    cells = construct_backend(backend, list_size,
+                              max_points_per_cell)
     collisions = [false for _ in 1:list_size]
-    coords = [ntuple(_ -> typemin(Int), NDIMS) for _ in 1:list_size]
-    return SpatialHashingCellList{NDIMS, typeof(points), typeof(coords),
-                                  typeof(collisions)}(points, coords, collisions, list_size)
+    coords = [typemin(UInt128) for _ in 1:list_size]
+
+    return SpatialHashingCellList(NDIMS, cells, coords, collisions, list_size)
 end
 
 function Base.empty!(cell_list::SpatialHashingCellList)
-    (; list_size) = cell_list
+    (; cells) = cell_list
     NDIMS = ndims(cell_list)
 
-    Base.empty!.(cell_list.points)
-    cell_list.coords .= [ntuple(_ -> typemin(Int), NDIMS) for _ in 1:list_size]
+    # `Base.empty!.(cells)`, but for all backends
+    @threaded default_backend(cells) for i in eachindex(cells)
+        emptyat!(cells, i)
+    end
+
+    fill!(cell_list.coords, typemin(UInt128))
     cell_list.collisions .= false
     return cell_list
 end
 
 # For each entry in the hash table, store the coordinates of the cell of the first point being inserted at this entry.
 # If a point with a different cell coordinate is being added, we have found a collision.
+# We flatten the coordinate tuples to an `UInt128` number to make it work with atomics.
 function push_cell!(cell_list::SpatialHashingCellList, cell, point)
-    (; points, coords, collisions, list_size) = cell_list
+    (; cells, coords, collisions, list_size) = cell_list
     NDIMS = ndims(cell_list)
     hash_key = spatial_hash(cell, list_size)
-    push!(points[hash_key], point)
 
-    cell_coord = coords[hash_key]
-    if cell_coord == ntuple(_ -> typemin(Int), NDIMS)
+    @boundscheck check_cell_bounds(cell_list, hash_key)
+    @inbounds pushat!(cells, hash_key, point)
+
+    cell_coord_hash = coordinates_flattened(cell)
+    previous_cell_coord = coords[hash_key]
+    if previous_cell_coord == typemin(UInt128)
         # If this cell is not used yet, set cell coordinates
-        coords[hash_key] = cell
-    elseif cell_coord != cell
+        coords[hash_key] = cell_coord_hash
+    elseif previous_cell_coord != cell_coord_hash
         # If it is already used by a different cell, mark as collision
         collisions[hash_key] = true
     end
 end
 
+function push_cell_atomic!(cell_list::SpatialHashingCellList, cell, point)
+    (; cells, coords, collisions, list_size) = cell_list
+    NDIMS = ndims(cell_list)
+    hash_key = spatial_hash(cell, list_size)
+
+    cell_coord_hash = coordinates_flattened(cell)
+
+    @boundscheck check_cell_bounds(cell_list, hash_key)
+    @inbounds pushat_atomic!(cells, hash_key, point)
+
+    cell_coord = @inbounds coords[hash_key]
+    if cell_coord == ntuple(_ -> typemin(UInt128), Val(NDIMS))
+        # If this cell is not used yet, set cell coordinates
+        @inbounds Atomix.@atomic coords[hash_key] = cell_coord_hash
+    elseif cell_coord != cell_coord_hash
+        # If it is already used by a different cell, mark as collision
+        @inbounds Atomix.@atomic collisions[hash_key] = true
+    end
+end
+
 function deleteat_cell!(cell_list::SpatialHashingCellList, cell, i)
     deleteat!(cell_list[cell], i)
 end
 
-@inline each_cell_index(cell_list::SpatialHashingCellList) = eachindex(cell_list.points)
+@inline each_cell_index(cell_list::SpatialHashingCellList) = eachindex(cell_list.cells)
 
 function copy_cell_list(cell_list::SpatialHashingCellList, search_radius,
                         periodic_box)
     (; list_size) = cell_list
     NDIMS = ndims(cell_list)
 
-    return SpatialHashingCellList{NDIMS}(list_size)
+    return SpatialHashingCellList{NDIMS}(list_size = list_size,
+                                         backend = typeof(cell_list.cells),
+                                         max_points_per_cell = max_points_per_cell(cell_list.cells))
 end
 
 @inline function Base.getindex(cell_list::SpatialHashingCellList, cell::Tuple)
-    return cell_list.points[spatial_hash(cell, length(cell_list.points))]
+    return cell_list.cells[spatial_hash(cell, length(cell_list.cells))]
 end
 
 @inline function Base.getindex(cell_list::SpatialHashingCellList, i::Integer)
-    return cell_list.points[i]
+    return cell_list.cells[i]
 end
 
 @inline function is_correct_cell(cell_list::SpatialHashingCellList{<:Any, Nothing},
@@ -106,3 +160,23 @@ function spatial_hash(cell::NTuple{3, Real}, list_size)
 
     return mod(xor(i * 73856093, j * 19349663, k * 83492791), list_size) + 1
 end
+
+@inline function check_cell_bounds(cell_list::SpatialHashingCellList, cell::Tuple)
+    check_cell_bounds(cell_list, spatial_hash(cell, cell_list.list_size))
+end
+
+# Compute a compact 128-bit representation by reinterpreting each coordinate as a UInt32
+# and bit-shifting them into a UInt128 slot (appending the `UInt32` bitstrings).
+function coordinates_flattened(cell_coordinate)
+    # Size check
+    @assert length(cell_coordinate) <= 3
+
+    result = UInt128(0)
+    for (i, coord) in enumerate(cell_coordinate)
+        ucoord = reinterpret(UInt32, Int32(coord))
+        # Shift the `i`-th coordinate by (i - 1) x 32 bits, so the used bits don't overlap
+        result = (UInt128(ucoord) << ((i-1) * 32)) | result
+    end
+
+    return result
+end

src/gpu.jl

@@ -31,3 +31,12 @@ function Adapt.adapt_structure(to, nhs::GridNeighborhoodSearch)
     return GridNeighborhoodSearch(cell_list, search_radius, periodic_box, n_cells,
                                   cell_size, update_buffer, nhs.update_strategy)
 end
+
+function Adapt.adapt_structure(to, cell_list::SpatialHashingCellList{NDIMS}) where {NDIMS}
+    (; list_size) = cell_list
+    cells = Adapt.adapt_structure(to, cell_list.cells)
+    coords = Adapt.adapt_structure(to, cell_list.coords)
+    collisions = Adapt.adapt_structure(to, cell_list.collisions)
+
+    return SpatialHashingCellList(NDIMS, cells, coords, collisions, list_size)
+end

src/nhs_grid.jl

@@ -390,6 +390,8 @@ end
 end
 
 # Fully parallel incremental update with atomic push.
+# TODO `cell_list.cells.lengths` and `cell_list.cells.backend` are hardcoded
+# for `FullGridCellList`, which is currently the only implementation
 function update_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                   ParallelIncrementalUpdate},
                       y::AbstractMatrix; parallelization_backend = default_backend(y),
@@ -470,7 +472,6 @@ end
 # with the `SpatialHashingCellList` if this cell has a collision.
 function check_collision(neighbor_cell_::CartesianIndex, neighbor_coords,
                          cell_list::SpatialHashingCellList, nhs)
-    (; list_size, collisions, coords) = cell_list
     neighbor_cell = periodic_cell_index(Tuple(neighbor_cell_), nhs)
 
     return neighbor_cell != cell_coords(neighbor_coords, nhs)
@@ -494,7 +495,8 @@ function check_cell_collision(neighbor_cell_::CartesianIndex,
     # `collisions[hash] == false` means points from only one cells are in this list.
     # We could still have a collision though, if this one cell is not `neighbor_cell`,
     # which is possible when `neighbor_cell` is empty.
-    return collisions[hash] || coords[hash] != neighbor_cell
+    return collisions[hash] ||
+           coords[hash] != PointNeighbors.coordinates_flattened(neighbor_cell)
 end
 
 # Specialized version of the function in `neighborhood_search.jl`, which is faster