Add distributed spatial query with pure callbacks

Co-authored-by: Daniel Arndt <[email protected]>
masterleinad · Jun 25, 2024 · 78f9d58 · 78f9d58
1 parent 71df935
commit 78f9d58
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 2 deletions.
diff --git a/src/ArborX_DistributedTree.hpp b/src/ArborX_DistributedTree.hpp
@@ -200,6 +200,13 @@ class DistributedTree<MemorySpace, Details::LegacyDefaultTemplateValue,
                      std::forward<OffsetView>(offset));
   }
 
+  template <typename ExecutionSpace, typename UserPredicates, typename Callback>
+  void query(ExecutionSpace const &space, UserPredicates const &user_predicates,
+             Callback &&callback) const
+  {
+    base_type::query(space, user_predicates, std::forward<Callback>(callback));
+  }
+
   template <typename ExecutionSpace, typename UserPredicates, typename Callback,
             typename Indices, typename Offset>
   void query(ExecutionSpace const &space, UserPredicates const &user_predicates,

diff --git a/src/details/ArborX_DetailsDistributedTreeImpl.hpp b/src/details/ArborX_DetailsDistributedTreeImpl.hpp
@@ -40,6 +40,13 @@ struct DistributedTreeImpl
                 ExecutionSpace const &space, Predicates const &queries,
                 Callback const &callback, OutputView &out, OffsetView &offset);
 
+  template <typename DistributedTree, typename ExecutionSpace,
+            typename Predicates, typename Callback>
+  static void queryDispatch(SpatialPredicateTag, DistributedTree const &tree,
+                            ExecutionSpace const &space,
+                            Predicates const &predicates,
+                            Callback const &callback);
+
   // nearest neighbors queries
   template <typename DistributedTree, typename ExecutionSpace,
             typename Predicates, typename Callback, typename Indices,

diff --git a/src/details/ArborX_DetailsDistributedTreeSpatial.hpp b/src/details/ArborX_DetailsDistributedTreeSpatial.hpp
@@ -47,16 +47,52 @@ DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
 
   using MemorySpace = typename Tree::memory_space;
 
+  auto const &top_tree = tree._top_tree;
+
   Kokkos::View<int *, MemorySpace> intersected_ranks(
       "ArborX::DistributedTree::query::spatial::intersected_ranks", 0);
-  tree._top_tree.query(space, predicates, LegacyDefaultCallback{},
-                       intersected_ranks, offset);
+  top_tree.query(space, predicates, intersected_ranks, offset);
 
   DistributedTree::forwardQueriesAndCommunicateResults(
       tree.getComm(), space, tree._bottom_tree, predicates, callback,
       intersected_ranks, offset, values);
 }
 
+template <typename Tree, typename ExecutionSpace, typename Predicates,
+          typename Callback>
+void DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
+                                        ExecutionSpace const &space,
+                                        Predicates const &predicates,
+                                        Callback const &callback)
+{
+  std::string prefix = "ArborX::DistributedTree::query::spatial(pure)";
+
+  Kokkos::Profiling::ScopedRegion guard(prefix);
+
+  if (tree.empty())
+    return;
+
+  using MemorySpace = typename Tree::memory_space;
+  using namespace DistributedTree;
+
+  auto const &top_tree = tree._top_tree;
+  auto const &bottom_tree = tree._bottom_tree;
+  auto comm = tree.getComm();
+
+  Kokkos::View<int *, MemorySpace> intersected_ranks(
+      prefix + "::intersected_ranks", 0);
+  Kokkos::View<int *, MemorySpace> offset(prefix + "::offset", 0);
+  top_tree.query(space, predicates, intersected_ranks, offset);
+
+  using Query = typename Predicates::value_type;
+  Kokkos::View<Query *, MemorySpace> fwd_predicates(prefix + "::fwd_predicates",
+                                                    0);
+  forwardQueries(comm, space, predicates, intersected_ranks, offset,
+                 fwd_predicates);
+
+  bottom_tree.query(space, fwd_predicates, callback);
+}
+
 template <typename Tree, typename ExecutionSpace, typename Predicates,
           typename Values, typename Offset>
 std::enable_if_t<Kokkos::is_view_v<Values> && Kokkos::is_view_v<Offset>>