GPU compatibility

Author: efaulhaber

src/gpu.jl

@@ -11,7 +11,6 @@ Adapt.@adapt_structure FullGridCellList
 Adapt.@adapt_structure SpatialHashingCellList
 Adapt.@adapt_structure DynamicVectorOfVectors
 Adapt.@adapt_structure GridNeighborhoodSearch
-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.
@@ -24,3 +23,13 @@ end
 function Adapt.adapt_structure(to::typeof(Array), range::UnitRange)
     return range
 end
+
+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)
+
+    return PrecomputedNeighborhoodSearch{ndims(nhs)}(neighbor_lists, search_radius,
+                                                     periodic_box, neighborhood_search)
+end

src/nhs_precomputed.jl

@@ -43,25 +43,34 @@ struct PrecomputedNeighborhoodSearch{NDIMS, NL, ELTYPE, PB, NHS} <:
     periodic_box        :: PB
     neighborhood_search :: NHS
 
-    function PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
-                                                  periodic_box = nothing,
-                                                  update_strategy = nothing,
-                                                  update_neighborhood_search = GridNeighborhoodSearch{NDIMS}(;
-                                                                                                             search_radius,
-                                                                                                             n_points,
-                                                                                                             periodic_box,
-                                                                                                             update_strategy),
-                                                  backend = DynamicVectorOfVectors{Int32},
-                                                  max_neighbors = 4 * NDIMS^4) where {NDIMS}
-        neighbor_lists = construct_backend(nothing, backend, n_points, max_neighbors)
-
-        new{NDIMS, typeof(neighbor_lists),
-            typeof(search_radius), typeof(periodic_box),
-            typeof(update_neighborhood_search)}(neighbor_lists, search_radius,
-                                                periodic_box, update_neighborhood_search)
+    function PrecomputedNeighborhoodSearch{NDIMS}(neighbor_lists, search_radius,
+                                                  periodic_box,
+                                                  update_neighborhood_search) where {NDIMS}
+        return new{NDIMS, typeof(neighbor_lists),
+                   typeof(search_radius),
+                   typeof(periodic_box),
+                   typeof(update_neighborhood_search)}(neighbor_lists, search_radius,
+                                                       periodic_box,
+                                                       update_neighborhood_search)
     end
 end
 
+function PrecomputedNeighborhoodSearch{NDIMS}(; search_radius = 0.0, n_points = 0,
+                                              periodic_box = nothing,
+                                              update_strategy = nothing,
+                                              update_neighborhood_search = GridNeighborhoodSearch{NDIMS}(;
+                                                                                                         search_radius,
+                                                                                                         n_points,
+                                                                                                         periodic_box,
+                                                                                                         update_strategy),
+                                              backend = DynamicVectorOfVectors{Int32},
+                                              max_neighbors = 4 * NDIMS^4) where {NDIMS}
+    neighbor_lists = construct_backend(nothing, backend, n_points, max_neighbors)
+
+    PrecomputedNeighborhoodSearch{NDIMS}(neighbor_lists, search_radius,
+                                         periodic_box, update_neighborhood_search)
+end
+
 @inline Base.ndims(::PrecomputedNeighborhoodSearch{NDIMS}) where {NDIMS} = NDIMS
 
 @inline requires_update(::PrecomputedNeighborhoodSearch) = (true, true)