Handling of tree destruction and unused views

arborx · Dec 27, 2023 · fcaefb0 · fcaefb0
1 parent 2f814b5
commit fcaefb0
Showing 1 changed file with 32 additions and 16 deletions.
diff --git a/src/interpolation/ArborX_InterpMovingLeastSquares.hpp b/src/interpolation/ArborX_InterpMovingLeastSquares.hpp
@@ -147,22 +147,8 @@ class MovingLeastSquares
     // There must be enough source points
     KOKKOS_ASSERT(0 < _num_neighbors && _num_neighbors <= _source_size);
 
-    // Organize the source points as a tree
-    BoundingVolumeHierarchy<MemorySpace, ArborX::PairValueIndex<SourcePoint>>
-        source_tree(space, ArborX::Experimental::attach_indices(source_points));
-
-    // Create the predicates
-    Details::MLSTargetPointsPredicateWrapper<TargetPoints> predicates{
-        {target_points}, _num_neighbors};
-
-    // Query the source
-    _indices = Kokkos::View<int *, MemorySpace>(
-        "ArborX::MovingLeastSquares::indices", 0);
-    Kokkos::View<int *, MemorySpace> offsets(
-        "ArborX::MovingLeastSquares::offsets", 0);
-    source_tree.query(space, predicates,
-                      ArborX::Details::LegacyDefaultCallback{}, _indices,
-                      offsets);
+    // Search for neighbors
+    searchNeighbors(space, source_points, target_points);
 
     // Fill in the value indices object so values can be transferred from a 1D
     // source data to a properly distributed 2D array for each target.
@@ -258,6 +244,36 @@ class MovingLeastSquares
     return source_view;
   }
 
+  template <typename ExecutionSpace, typename SourcePoints,
+            typename TargetPoints>
+  void searchNeighbors(ExecutionSpace const &space,
+                       SourcePoints const &source_points,
+                       TargetPoints const &target_points)
+  {
+    auto guard = Kokkos::Profiling::ScopedRegion(
+        "ArborX::MovingLeastSquares::searchNeighbors");
+
+    // Organize the source points as a tree
+    using SourcePoint =
+        typename ArborX::Details::AccessValues<SourcePoints,
+                                               PrimitivesTag>::value_type;
+    BoundingVolumeHierarchy<MemorySpace, ArborX::PairValueIndex<SourcePoint>>
+        source_tree(space, ArborX::Experimental::attach_indices(source_points));
+
+    // Create the predicates
+    Details::MLSTargetPointsPredicateWrapper<TargetPoints> predicates{
+        {target_points}, _num_neighbors};
+
+    // Query the source
+    _indices = Kokkos::View<int *, MemorySpace>(
+        "ArborX::MovingLeastSquares::indices", 0);
+    Kokkos::View<int *, MemorySpace> offsets(
+        "ArborX::MovingLeastSquares::offsets", 0);
+    source_tree.query(space, predicates,
+                      ArborX::Details::LegacyDefaultCallback{}, _indices,
+                      offsets);
+  }
+
   Kokkos::View<FloatingCalculationType **, MemorySpace> _coeffs;
   Kokkos::View<int *, MemorySpace> _indices;
   int _num_targets;