Add function to strip not GPU-compatible data structures

Author: efaulhaber

benchmarks/plot.jl

@@ -68,8 +68,10 @@ function plot_benchmarks(benchmark, n_points_per_dimension, iterations;
         for i in eachindex(neighborhood_searches)
             neighborhood_search = neighborhood_searches[i]
             initialize!(neighborhood_search, coordinates, coordinates)
+            # Remove unnecessary data structures that are only used for initialization
+            neighborhood_search_ = PointNeighbors.freeze_neighborhood_search(neighborhood_search)
 
-            time = benchmark(neighborhood_search, coordinates; parallelization_backend)
+            time = benchmark(neighborhood_search_, coordinates; parallelization_backend)
             times[iter, i] = time
             time_string = BenchmarkTools.prettytime(time * 1e9)
             println("$(neighborhood_searches_names[i])")

src/gpu.jl

@@ -28,7 +28,7 @@ function Adapt.adapt_structure(to, nhs::PrecomputedNeighborhoodSearch)
     neighbor_lists = Adapt.adapt_structure(to, nhs.neighbor_lists)
     search_radius = Adapt.adapt_structure(to, nhs.search_radius)
     periodic_box = Adapt.adapt_structure(to, nhs.periodic_box)
-    neighborhood_search = nothing# Adapt.adapt_structure(to, nhs.neighborhood_search)
+    neighborhood_search = Adapt.adapt_structure(to, nhs.neighborhood_search)
 
     return PrecomputedNeighborhoodSearch{ndims(nhs)}(neighbor_lists, search_radius,
                                                      periodic_box, neighborhood_search)

src/neighborhood_search.jl

@@ -108,6 +108,14 @@ GridNeighborhoodSearch{2, Float64, ...}(...)
     return nothing
 end
 
+@inline function freeze_neighborhood_search(search::AbstractNeighborhoodSearch)
+    # Indicate that the neighborhood search is static and will not be updated anymore.
+    # Some implementations might use this to strip unnecessary data structures for updating.
+    # Notably, this is used for the `PrecomputedNeighborhoodSearch` to strip a potentially
+    # not GPU-compatible inner neighborhood search that is used only for initialization.
+    return search
+end
+
 """
     PeriodicBox(; min_corner, max_corner)
 

src/nhs_precomputed.jl

@@ -190,6 +190,16 @@ function copy_neighborhood_search(nhs::PrecomputedNeighborhoodSearch,
                                                      max_neighbors)
 end
 
+@inline function freeze_neighborhood_search(search::PrecomputedNeighborhoodSearch)
+    # Indicate that the neighborhood search is static and will not be updated anymore.
+    # For the `PrecomputedNeighborhoodSearch`, strip the inner neighborhood search,
+    # which is used only for initialization and updating.
+    return PrecomputedNeighborhoodSearch{ndims(search)}(search.neighbor_lists,
+                                                        search.search_radius,
+                                                        search.periodic_box,
+                                                        nothing)
+end
+
 # TODO move to `vector_of_vectors.jl`
 function max_inner_length(cells::DynamicVectorOfVectors, fallback)
     return size(cells.backend, 1)