Switch to using new helper for the nearest distributed

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

Hide file tree

Showing 2 changed files with 23 additions and 54 deletions.
diff --git a/src/details/ArborX_DetailsDistributedTreeNearest.hpp b/src/details/ArborX_DetailsDistributedTreeNearest.hpp
@@ -357,57 +357,25 @@ DistributedTreeImpl::queryDispatchImpl(NearestPredicateTag, Tree const &tree,
   {
     implementStrategy(space, tree, queries, distances, nearest_ranks, offset);
 
-    {
-      // NOTE_COMM_NEAREST: The communication pattern here for the nearest
-      // search is identical to that of the spatial search (see
-      // NOTE_COMM_SPATIAL). The code differences are:
-      // - no callbacks
-      // - explicit distances
-      // - results filtering
-
-      // Forward queries
-      using Query = typename Predicates::value_type;
-      Kokkos::View<int *, MemorySpace> ids(
-          "ArborX::DistributedTree::query::nearest::query_ids", 0);
-      Kokkos::View<Query *, MemorySpace> fwd_queries(
-          "ArborX::DistributedTree::query::nearest::fwd_queries", 0);
-      forwardQueries(comm, space, queries, nearest_ranks, offset, fwd_queries,
-                     ids, ranks);
-
-      // Perform queries that have been received
-      Kokkos::View<PairValueDistance<Value> *, MemorySpace> out(
-          "ArborX::DistributedTree::query::nearest::pairs_index_distance", 0);
-      bottom_tree.query(space, fwd_queries, callback_with_distance, out,
-                        offset);
-
-      // Unzip
-      auto const n = out.extent(0);
-      KokkosExt::reallocWithoutInitializing(space, values, n);
-      KokkosExt::reallocWithoutInitializing(space, distances, n);
-      Kokkos::parallel_for(
-          "ArborX::DistributedTree::query::nearest::split_"
-          "index_distance_pairs",
-          Kokkos::RangePolicy<ExecutionSpace>(space, 0, n),
-          KOKKOS_LAMBDA(int i) {
-            values(i) = out(i).value;
-            distances(i) = out(i).distance;
-          });
-
-      // Communicate results back
-      communicateResultsBack(comm, space, values, offset, ranks, ids,
-                             &distances);
-
-      // Merge results
-      Kokkos::Profiling::pushRegion(
-          "ArborX::DistributedTree::query::nearest::postprocess_results");
-
-      int const n_queries = queries.size();
-      countResults(space, n_queries, ids, offset);
-      sortResultsByKey(space, ids, values, ranks, distances);
-      filterResults(space, queries, distances, values, offset, ranks);
-
-      Kokkos::Profiling::popRegion();
-    }
+    Kokkos::View<PairValueDistance<Value> *, MemorySpace> out(
+        "ArborX::DistributedTree::query::nearest::pairs_index_distance", 0);
+    forwardQueriesAndCommunicateResults(comm, space, bottom_tree, queries,
+                                        callback_with_distance, nearest_ranks,
+                                        offset, out, ranks);
+
+    // Unzip
+    auto const n = out.extent(0);
+    KokkosExt::reallocWithoutInitializing(space, values, n);
+    KokkosExt::reallocWithoutInitializing(space, distances, n);
+    Kokkos::parallel_for(
+        "ArborX::DistributedTree::query::nearest::"
+        "split_index_distance_pairs",
+        Kokkos::RangePolicy<ExecutionSpace>(space, 0, n), KOKKOS_LAMBDA(int i) {
+          values(i) = out(i).value;
+          distances(i) = out(i).distance;
+        });
+
+    filterResults(space, queries, distances, values, offset, ranks);
   }
 }
 

diff --git a/src/details/ArborX_DetailsDistributedTreeUtils.hpp b/src/details/ArborX_DetailsDistributedTreeUtils.hpp
@@ -108,8 +108,9 @@ void sortResultsByKey(ExecutionSpace const &space, View keys,
   if (n == 0)
     return;
 
-  if constexpr (sizeof...(OtherViews) == 1 &&
-                std::tuple_element_t<0, std::tuple<OtherViews...>>::rank == 1)
+  using ViewType = std::tuple_element_t<0, std::tuple<OtherViews...>>;
+  if constexpr (sizeof...(OtherViews) == 1 && ViewType::rank == 1 &&
+                std::is_arithmetic_v<typename ViewType::value_type>)
   {
     // If there's only one 1D view to process, we can avoid computing the
     // permutation.
@@ -363,7 +364,7 @@ void forwardQueriesAndCommunicateResults(
   // Merge results
   int const n_predicates = predicates.size();
   countResults(space, n_predicates, ids, offset);
-  sortResultsByKey(space, ids, values);
+  sortResultsByKey(space, ids, values, ranks);
 
   Kokkos::Profiling::popRegion();
 }