Get rid of last ranks argument in forwardQueriesAndCommunicateResults

aprokop · May 3, 2024 · 1b4832b · 1b4832b
1 parent b8c0a24
commit 1b4832b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 15 deletions.
diff --git a/src/details/ArborX_DetailsDistributedTreeNearest.hpp b/src/details/ArborX_DetailsDistributedTreeNearest.hpp
@@ -152,12 +152,11 @@ void DistributedTreeImpl::phaseII(ExecutionSpace const &space, Tree const &tree,
   // rather than nearest predicates.
   Kokkos::View<PairValueDistance<typename Values::value_type> *, MemorySpace>
       out(prefix + "::pairs_value_distance", 0);
-  Kokkos::View<int *, MemorySpace> ranks(prefix + "::ranks", 0);
   DistributedTree::forwardQueriesAndCommunicateResults(
       tree.getComm(), space, bottom_tree, predicates,
       CallbackWithDistance<BottomTree, Callback, typename Values::value_type,
                            true>(space, bottom_tree, callback),
-      nearest_ranks, offset, out, &ranks);
+      nearest_ranks, offset, out);
 
   // Unzip
   auto n = out.extent(0);

diff --git a/src/details/ArborX_DetailsDistributedTreeUtils.hpp b/src/details/ArborX_DetailsDistributedTreeUtils.hpp
@@ -318,15 +318,12 @@ 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 = Kokkos::View<int *, typename BottomTree::memory_space>>
+template <typename ExecutionSpace, typename BottomTree, typename Predicates,
+          typename Callback, typename RanksTo, typename Offset, typename Values>
 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_ptr = nullptr)
+    RanksTo const &ranks_to, Offset &offset, Values &values)
 {
   std::string prefix =
       "ArborX::DistributedTree::query::forwardQueriesAndCommunicateResults";
@@ -335,13 +332,11 @@ void forwardQueriesAndCommunicateResults(
   using Query = typename Predicates::value_type;
   using MemorySpace = typename BottomTree::memory_space;
 
-  Ranks ranks_aux(prefix + "::ranks", 0);
-  auto &ranks = (ranks_ptr != nullptr ? *ranks_ptr : ranks_aux);
-
   // Forward predicates
   Kokkos::View<int *, MemorySpace> ids(prefix + "::query_ids", 0);
   Kokkos::View<Query *, MemorySpace> fwd_predicates(prefix + "::fwd_predicates",
                                                     0);
+  Kokkos::View<int *, MemorySpace> ranks(prefix + "::ranks", 0);
   forwardQueries(comm, space, predicates, ranks_to, offset, fwd_predicates, ids,
                  ranks);
 
@@ -356,10 +351,7 @@ void forwardQueriesAndCommunicateResults(
   // Merge results
   int const n_predicates = predicates.size();
   countResults(space, n_predicates, ids, offset);
-  if (ranks_ptr != nullptr)
-    sortResultsByKey(space, ids, values, ranks);
-  else
-    sortResultsByKey(space, ids, values);
+  sortResultsByKey(space, ids, values);
 
   Kokkos::Profiling::popRegion();
 }