Fix benchmarks

Author: efaulhaber

benchmarks/n_body.jl

@@ -18,7 +18,9 @@ function benchmark_n_body(neighborhood_search, coordinates_;
     # allows us to change the type of the array to run the benchmark on the GPU.
     # Passing `parallel = true` or `parallel = false` will not change anything here.
     coordinates = PointNeighbors.Adapt.adapt(parallelization_backend, coordinates_)
-    nhs = PointNeighbors.Adapt.adapt(parallelization_backend, neighborhood_search)
+    # Remove unnecessary data structures that are only used for initialization
+    neighborhood_search_ = PointNeighbors.freeze_neighborhood_search(neighborhood_search)
+    nhs = PointNeighbors.Adapt.adapt(parallelization_backend, neighborhood_search_)
 
     # This preserves the data type of `coordinates`, which makes it work for GPU types
     mass = 1e10 * (rand!(similar(coordinates, size(coordinates, 2))) .+ 1)

benchmarks/plot.jl

@@ -68,10 +68,8 @@ 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])")

benchmarks/smoothed_particle_hydrodynamics.jl

@@ -54,7 +54,9 @@ function benchmark_wcsph(neighborhood_search, coordinates;
                                                density_diffusion = density_diffusion)
 
     system = PointNeighbors.Adapt.adapt(parallelization_backend, fluid_system)
-    nhs = PointNeighbors.Adapt.adapt(parallelization_backend, neighborhood_search)
+    # Remove unnecessary data structures that are only used for initialization
+    neighborhood_search_ = PointNeighbors.freeze_neighborhood_search(neighborhood_search)
+    nhs = PointNeighbors.Adapt.adapt(parallelization_backend, neighborhood_search_)
     semi = DummySemidiscretization(nhs, parallelization_backend)
 
     v = PointNeighbors.Adapt.adapt(parallelization_backend,
@@ -105,7 +107,9 @@ function benchmark_wcsph_fp32(neighborhood_search, coordinates_;
                                                density_diffusion = density_diffusion)
 
     system = PointNeighbors.Adapt.adapt(parallelization_backend, fluid_system)
-    nhs = PointNeighbors.Adapt.adapt(parallelization_backend, neighborhood_search)
+    # Remove unnecessary data structures that are only used for initialization
+    neighborhood_search_ = PointNeighbors.freeze_neighborhood_search(neighborhood_search)
+    nhs = PointNeighbors.Adapt.adapt(parallelization_backend, neighborhood_search_)
     semi = DummySemidiscretization(nhs, parallelization_backend)
 
     v = PointNeighbors.Adapt.adapt(parallelization_backend,