Make it work on GPUs

Author: efaulhaber

src/gpu.jl

@@ -9,6 +9,7 @@
 # `Adapt.@adapt_structure` automatically generates the `adapt` function for our custom types.
 Adapt.@adapt_structure FullGridCellList
 Adapt.@adapt_structure DynamicVectorOfVectors
+# 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.
@@ -31,3 +32,11 @@ 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, nhs::PrecomputedNeighborhoodSearch)
+    neighbor_lists = Adapt.adapt_structure(to, nhs.neighbor_lists)
+    periodic_box = Adapt.adapt_structure(to, nhs.periodic_box)
+    neighborhood_search = TrivialNeighborhoodSearch{ndims(nhs)}(search_radius = search_radius(nhs))
+
+    return PrecomputedNeighborhoodSearch(neighborhood_search, neighbor_lists, periodic_box)
+end

src/nhs_precomputed.jl

@@ -37,21 +37,26 @@ struct PrecomputedNeighborhoodSearch{NDIMS, NHS, NL, PB} <: AbstractNeighborhood
     neighbor_lists      :: NL
     periodic_box        :: PB
 
-    function PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
-                                                  periodic_box = nothing,
-                                                  update_strategy = nothing,
-                                                  backend = VectorOfVectors{Int}) where {NDIMS}
-        nhs = GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
-                                            periodic_box, update_strategy)
-
-        neighbor_lists = backend()
-
-        new{NDIMS, typeof(nhs),
-            typeof(neighbor_lists),
-            typeof(periodic_box)}(nhs, neighbor_lists, periodic_box)
+    # This constructor is necessary for Adapt.jl to work with this struct.
+    # See the comments in gpu.jl for more details.
+    function PrecomputedNeighborhoodSearch(nhs, neighbor_lists, periodic_box)
+        new{ndims(nhs), typeof(nhs),
+            typeof(neighbor_lists), typeof(periodic_box)}(nhs, neighbor_lists, periodic_box)
     end
 end
 
+function PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
+                                                periodic_box = nothing,
+                                                update_strategy = nothing,
+                                                backend = VectorOfVectors{Int}) where {NDIMS}
+    nhs = GridNeighborhoodSearch{NDIMS}(; search_radius, n_points,
+                                        periodic_box, update_strategy)
+
+    neighbor_lists = backend()
+
+    return PrecomputedNeighborhoodSearch(nhs, neighbor_lists, periodic_box)
+end
+
 @inline Base.ndims(::PrecomputedNeighborhoodSearch{NDIMS}) where {NDIMS} = NDIMS
 
 @inline requires_update(::PrecomputedNeighborhoodSearch) = (true, true)