Create a helper forwardQueriesAndCommunicateResults

aprokop · Apr 29, 2024 · 5e66108 · 5e66108
1 parent 6803cce
commit 5e66108
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 40 deletions.
diff --git a/src/details/ArborX_DetailsDistributedTreeSpatial.hpp b/src/details/ArborX_DetailsDistributedTreeSpatial.hpp
@@ -45,53 +45,18 @@ DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
     return;
   }
 
-  using namespace DistributedTree;
   using MemorySpace = typename Tree::memory_space;
 
-  auto const &top_tree = tree._top_tree;
-  auto const &bottom_tree = tree._bottom_tree;
-  auto comm = tree.getComm();
-
   Kokkos::View<int *, MemorySpace> intersected_ranks(
       "ArborX::DistributedTree::query::spatial::intersected_ranks", 0);
-  top_tree.query(space, predicates, LegacyDefaultCallback{}, intersected_ranks,
-                 offset);
+  tree._top_tree.query(space, predicates, LegacyDefaultCallback{},
+                       intersected_ranks, offset);
 
   Kokkos::View<int *, MemorySpace> ranks(
       "ArborX::DistributedTree::query::spatial::ranks", 0);
-  {
-    // NOTE_COMM_SPATIAL: The communication pattern here for the spatial search
-    // is identical to that of the nearest search (see NOTE_COMM_NEAREST). The
-    // code differences are:
-    // - usage of callbacks
-    // - no explicit distances
-    // - no results filtering
-
-    // Forward predicates
-    using Query = typename Predicates::value_type;
-    Kokkos::View<int *, MemorySpace> ids(
-        "ArborX::DistributedTree::query::spatial::query_ids", 0);
-    Kokkos::View<Query *, MemorySpace> fwd_predicates(
-        "ArborX::DistributedTree::query::spatial::fwd_predicates", 0);
-    forwardQueries(comm, space, predicates, intersected_ranks, offset,
-                   fwd_predicates, ids, ranks);
-
-    // Perform predicates that have been received
-    bottom_tree.query(space, fwd_predicates, callback, values, offset);
-
-    // Communicate results back
-    communicateResultsBack(comm, space, values, offset, ranks, ids);
-
-    Kokkos::Profiling::pushRegion(
-        "ArborX::DistributedTree::spatial::postprocess_results");
-
-    // Merge results
-    int const n_predicates = predicates.size();
-    countResults(space, n_predicates, ids, offset);
-    sortResultsByKey(space, ids, values);
-
-    Kokkos::Profiling::popRegion();
-  }
+  DistributedTree::forwardQueriesAndCommunicateResults(
+      tree.getComm(), space, tree._bottom_tree, predicates, callback,
+      intersected_ranks, offset, values, ranks);
 }
 
 template <typename Tree, typename ExecutionSpace, typename Predicates,

diff --git a/src/details/ArborX_DetailsDistributedTreeUtils.hpp b/src/details/ArborX_DetailsDistributedTreeUtils.hpp
@@ -330,6 +330,44 @@ void communicateResultsBack(MPI_Comm comm, ExecutionSpace const &space,
   }
 }
 
+template <typename ExecutionSpace, typename BottomTree, typename Predicates,
+          typename Callback, typename RanksTo, typename Offset, typename Values,
+          typename Ranks>
+void forwardQueriesAndCommunicateResults(
+    MPI_Comm comm, ExecutionSpace const &space, BottomTree const &bottom_tree,
+    Predicates const &predicates, Callback const &callback,
+    RanksTo const &ranks_to, Offset &offset, Values &values, Ranks &ranks)
+{
+  std::string prefix =
+      "ArborX::DistributedTree::query::forwardQueriesAndCommunicateResults";
+  Kokkos::Profiling::ScopedRegion guard(prefix);
+
+  using Query = typename Predicates::value_type;
+  using MemorySpace = typename BottomTree::memory_space;
+
+  // Forward predicates
+  Kokkos::View<int *, MemorySpace> ids(prefix + "::query_ids", 0);
+  Kokkos::View<Query *, MemorySpace> fwd_predicates(prefix + "::fwd_predicates",
+                                                    0);
+  forwardQueries(comm, space, predicates, ranks_to, offset, fwd_predicates, ids,
+                 ranks);
+
+  // Perform predicates that have been received
+  bottom_tree.query(space, fwd_predicates, callback, values, offset);
+
+  // Communicate results back
+  communicateResultsBack(comm, space, values, offset, ranks, ids);
+
+  Kokkos::Profiling::pushRegion(prefix + "postprocess_results");
+
+  // Merge results
+  int const n_predicates = predicates.size();
+  countResults(space, n_predicates, ids, offset);
+  sortResultsByKey(space, ids, values);
+
+  Kokkos::Profiling::popRegion();
+}
+
 template <typename ExecutionSpace, typename MemorySpace, typename Predicates,
           typename Values, typename Offset, typename Ranks>
 void filterResults(ExecutionSpace const &space, Predicates const &queries,