Compact Vector of Vectors

.gitignore

@@ -14,3 +14,8 @@ out/*
 .DS_Store
 
 LocalPreferences.toml
+
+benchmarks/plot.jl
+benchmarks/benchmark_initialize.jl
+gpu_env
+Project.toml

src/cell_lists/cell_lists.jl

@@ -16,6 +16,11 @@ function construct_backend(::Type{Vector{Vector{T}}},
     return [T[] for _ in 1:max_outer_length]
 end
 
+function construct_backend(::Type{CompactVectorOfVectors{T}},
+                           max_outer_length, _) where {T}
+    return CompactVectorOfVectors{T}(n_bins = max_outer_length)
+end
+
 function construct_backend(::Type{DynamicVectorOfVectors{T}},
                            max_outer_length,
                            max_inner_length) where {T}
@@ -30,12 +35,20 @@ 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(::Type{DynamicVectorOfVectors{T1, T2, T3, T4}}, max_outer_length,
+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 construct_backend(::Type{CompactVectorOfVectors{T1, T2, T3, T4}},
+                           max_outer_length,
+                           max_inner_length) where {T1, T2, T3, T4}
+    return construct_backend(CompactVectorOfVectors{T1}, max_outer_length,
+                             max_inner_length)
+end
+
 function max_points_per_cell(cells::DynamicVectorOfVectors)
     return size(cells.backend, 1)
 end

src/cell_lists/full_grid.jl

@@ -38,6 +38,10 @@ function supported_update_strategies(::FullGridCellList{<:DynamicVectorOfVectors
             SerialIncrementalUpdate, SerialUpdate)
 end
 
+function supported_update_strategies(::FullGridCellList{<:CompactVectorOfVectors})
+    return (ParallelUpdate, SerialUpdate)
+end
+
 function supported_update_strategies(::FullGridCellList)
     return (SemiParallelUpdate, SerialIncrementalUpdate, SerialUpdate)
 end
@@ -183,7 +187,6 @@ end
 
 function copy_cell_list(cell_list::FullGridCellList, search_radius, periodic_box)
     (; min_corner, max_corner) = cell_list
-
     return FullGridCellList(; min_corner, max_corner, search_radius,
                             backend = typeof(cell_list.cells),
                             max_points_per_cell = max_points_per_cell(cell_list.cells))

src/gpu.jl

@@ -8,7 +8,9 @@
 #
 # `Adapt.@adapt_structure` automatically generates the `adapt` function for our custom types.
 Adapt.@adapt_structure FullGridCellList
+Adapt.@adapt_structure SpatialHashingCellList
 Adapt.@adapt_structure DynamicVectorOfVectors
+Adapt.@adapt_structure CompactVectorOfVectors
 
 # `adapt(CuArray, ::SVector)::SVector`, but `adapt(Array, ::SVector)::Vector`.
 # We don't want to change the type of the `SVector` here.

src/nhs_grid.jl

@@ -239,6 +239,33 @@ function initialize_grid!(neighborhood_search::GridNeighborhoodSearch, y::Abstra
     return neighborhood_search
 end
 
+function initialize_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any, ParallelUpdate,
+                                                                      <:FullGridCellList{<:CompactVectorOfVectors}},
+                          y::AbstractMatrix;
+                          parallelization_backend = default_backend(y),
+                          eachindex_y = axes(y, 2))
+    (; cell_list) = neighborhood_search
+
+    if eachindex_y != axes(y, 2)
+        # Incremental update doesn't support inactive points
+        error("this neighborhood search/update strategy does not support inactive points")
+    end
+
+    if neighborhood_search.search_radius < eps()
+        # Cannot initialize with zero search radius.
+        # This is used in TrixiParticles when a neighborhood search is not used.
+        return neighborhood_search
+    end
+
+    resize!(cell_list.cells.values, size(y, 2))
+    cell_list.cells.values .= eachindex_y
+
+    point_to_cell = point_to_cell_wrapper(neighborhood_search, y)
+    update!(cell_list.cells, point_to_cell)
+
+    return neighborhood_search
+end
+
 function initialize_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any,
                                                                       ParallelUpdate},
                           y::AbstractMatrix; parallelization_backend = default_backend(y),
@@ -463,6 +490,19 @@ function update_grid!(neighborhood_search::Union{GridNeighborhoodSearch{<:Any,
     initialize_grid!(neighborhood_search, y; parallelization_backend, eachindex_y)
 end
 
+function update_grid!(neighborhood_search::GridNeighborhoodSearch{<:Any, ParallelUpdate,
+                                                                  <:FullGridCellList{<:CompactVectorOfVectors}},
+                      y::AbstractMatrix; parallelization_backend = default_backend(y),
+                      eachindex_y = axes(y, 2))
+    if eachindex_y != axes(y, 2)
+        # Incremental update doesn't support inactive points
+        error("this neighborhood search/update strategy does not support inactive points")
+    end
+
+    point_to_cell = point_to_cell_wrapper(neighborhood_search, y)
+    update!(neighborhood_search.cell_list.cells, point_to_cell)
+end
+
 function check_collision(neighbor_cell_, neighbor_coords, cell_list, nhs)
     # This is only relevant for the `SpatialHashingCellList`
     return false
@@ -517,7 +557,6 @@ end
 
         for neighbor_ in eachindex(neighbors)
             neighbor = @inbounds neighbors[neighbor_]
-
             # Making the following `@inbounds` yields a ~2% speedup on an NVIDIA H100.
             # But we don't know if `neighbor` (extracted from the cell list) is in bounds.
             neighbor_coords = extract_svector(neighbor_system_coords,
@@ -615,3 +654,12 @@ function copy_neighborhood_search(nhs::GridNeighborhoodSearch, search_radius, n_
                                               cell_list,
                                               update_strategy = nhs.update_strategy)
 end
+
+@inline function point_to_cell_wrapper(neighborhood_search, y)
+    @inline function point_to_cell(point)
+        point_coords = @inbounds extract_svector(y, Val(ndims(neighborhood_search)), point)
+        cell = cell_coords(point_coords, neighborhood_search)
+        return PointNeighbors.cell_index(neighborhood_search.cell_list, cell)
+    end
+    return point_to_cell
+end

src/vector_of_vectors.jl

@@ -161,3 +161,61 @@ end
 
     return vov
 end
+
+# As opposed to `DynamicVectorOfVectors`, this data structure does not support
+# modifying operations like `push!`.
+# It can be updated by assigning each contained integer value a new bin index.
+struct CompactVectorOfVectors{T, V, L, I}
+    values          :: V # Vector{Int} containing all values sorted by bin
+    n_bins          :: L # Ref{Int32}: Number of bins
+    first_bin_index :: I # Vector{Int} containing the first index in `values` of each bin
+
+    # This constructor is necessary for Adapt.jl to work with this struct.
+    # See the comments in gpu.jl for more details.
+    function CompactVectorOfVectors(values, n_bins, first_bin_index)
+        new{eltype(values), typeof(values), typeof(n_bins), typeof(first_bin_index)}(values,
+                                                                                     n_bins,
+                                                                                     first_bin_index)
+    end
+end
+
+function CompactVectorOfVectors{T}(; n_bins::Int) where {T}
+    first_bin_index = Vector{Int}(undef, n_bins+1)
+    first_bin_index[1] = 1 # required for the update
+    values = Vector{T}(undef, 0)
+    n_bins = Ref{Int32}(n_bins)
+
+    return CompactVectorOfVectors(values,
+                                  n_bins,
+                                  first_bin_index)
+end
+
+# Mhh should this may be changed to length(vov.values)?
+@inline Base.size(vov::CompactVectorOfVectors) = (vov.n_bins[],)
+
+@inline function Base.getindex(vov::CompactVectorOfVectors, i)
+    (; values, first_bin_index) = vov
+
+    start = first_bin_index[i]
+    stop = first_bin_index[i + 1] - 1
+    return view(values, start:stop)
+end
+
+@inline function update!(vov::CompactVectorOfVectors, f)
+    (; values, first_bin_index, n_bins) = vov
+
+    # TODO figure out how to do that fast and on the GPU
+    sort!(values, by = f)
+
+    # TODO figure out how to do that fast and on the GPU
+    n_particles_per_cell = zeros(n_bins[])
+    for val in values
+        n_particles_per_cell[f(val)] += 1
+    end
+
+    # Add 1 since first_bin_index starts at 1
+    n_particles_per_cell[1] += 1
+
+    # TODO avoid allocations
+    first_bin_index[2:end] .= cumsum(n_particles_per_cell)
+end

test/neighborhood_search.jl

@@ -57,6 +57,12 @@
                                                                            max_corner,
                                                                            search_radius,
                                                                            backend = Vector{Vector{Int32}})),
+                GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                              periodic_box = periodic_boxes[i],
+                                              cell_list = FullGridCellList(; min_corner,
+                                                                           max_corner,
+                                                                           search_radius,
+                                                                           backend = PointNeighbors.CompactVectorOfVectors{Int32})),
                 PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                                      periodic_box = periodic_boxes[i]),
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
@@ -70,6 +76,7 @@
                 "`GridNeighborhoodSearch`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`",
+                "`GridNeighborhoodSearch` with `FullGridCellList` with `CompactVectorOfVectors`",
                 "`PrecomputedNeighborhoodSearch`",
                 "`GridNeighborhoodSearch` with `SpatialHashingCellList`"
             ]
@@ -85,6 +92,10 @@
                                               cell_list = FullGridCellList(min_corner = periodic_boxes[i].min_corner,
                                                                            max_corner = periodic_boxes[i].max_corner,
                                                                            backend = Vector{Vector{Int32}})),
+                GridNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i],
+                                              cell_list = FullGridCellList(min_corner = periodic_boxes[i].min_corner,
+                                                                           max_corner = periodic_boxes[i].max_corner,
+                                                                           backend = PointNeighbors.CompactVectorOfVectors{Int32})),
                 PrecomputedNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i]),
                 GridNeighborhoodSearch{NDIMS}(periodic_box = periodic_boxes[i],
                                               cell_list = SpatialHashingCellList{NDIMS}(list_size = 2 *
@@ -192,6 +203,11 @@
                                                                            max_corner,
                                                                            search_radius,
                                                                            backend = Vector{Vector{Int}})),
+                GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                              cell_list = FullGridCellList(; min_corner,
+                                                                           max_corner,
+                                                                           search_radius,
+                                                                           backend = PointNeighbors.CompactVectorOfVectors{Int32})),
                 PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points),
                 GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
                                               cell_list = SpatialHashingCellList{NDIMS}(list_size = 2 *
@@ -210,6 +226,7 @@
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors` and `ParallelIncrementalUpdate`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors` and `SemiParallelUpdate`",
                 "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`",
+                "`GridNeighborhoodSearch` with `FullGridCellList` with `CompactVectorOfVectors`",
                 "`PrecomputedNeighborhoodSearch`",
                 "`GridNeighborhoodSearch` with `SpatialHashingCellList` with `DynamicVectorOfVectors`",
                 "`GridNeighborhoodSearch` with `SpatialHashingCellList` with `Vector{Vector}`"
@@ -231,6 +248,9 @@
                 GridNeighborhoodSearch{NDIMS}(cell_list = FullGridCellList(; min_corner,
                                                                            max_corner,
                                                                            backend = Vector{Vector{Int32}})),
+                GridNeighborhoodSearch{NDIMS}(cell_list = FullGridCellList(; min_corner,
+                                                                           max_corner,
+                                                                           backend = PointNeighbors.CompactVectorOfVectors{Int32})),
                 PrecomputedNeighborhoodSearch{NDIMS}(),
                 GridNeighborhoodSearch{NDIMS}(cell_list = SpatialHashingCellList{NDIMS}(list_size = 2 *
                                                                                                     n_points)),