Rename sortResults -> sortResultsByKey and allow modifying keys

masterleinad · Apr 10, 2024 · 902d1c8 · 902d1c8
1 parent 3ef91ae
commit 902d1c8
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 19 deletions.
diff --git a/src/details/ArborX_DetailsDistributedTreeNearest.hpp b/src/details/ArborX_DetailsDistributedTreeNearest.hpp
@@ -403,7 +403,7 @@ DistributedTreeImpl::queryDispatchImpl(NearestPredicateTag, Tree const &tree,
 
       int const n_queries = queries.size();
       countResults(space, n_queries, ids, offset);
-      sortResults(space, ids, values, ranks, distances);
+      sortResultsByKey(space, ids, values, ranks, distances);
       filterResults(space, queries, distances, values, offset, ranks);
 
       Kokkos::Profiling::popRegion();

diff --git a/src/details/ArborX_DetailsDistributedTreeSpatial.hpp b/src/details/ArborX_DetailsDistributedTreeSpatial.hpp
@@ -88,7 +88,7 @@ DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
     // Merge results
     int const n_predicates = predicates.size();
     countResults(space, n_predicates, ids, offset);
-    sortResults(space, ids, values);
+    sortResultsByKey(space, ids, values);
 
     Kokkos::Profiling::popRegion();
   }

diff --git a/src/details/ArborX_DetailsDistributedTreeUtils.hpp b/src/details/ArborX_DetailsDistributedTreeUtils.hpp
@@ -98,8 +98,8 @@ sendAcrossNetwork(ExecutionSpace const &space, Distributor const &distributor,
 }
 
 template <typename ExecutionSpace, typename View, typename... OtherViews>
-void sortResults(ExecutionSpace const &space, View keys,
-                 OtherViews... other_views)
+void sortResultsByKey(ExecutionSpace const &space, View keys,
+                      OtherViews... other_views)
 {
   auto const n = keys.extent(0);
   // If they were no queries, min_val and max_val values won't change after
@@ -108,18 +108,16 @@ void sortResults(ExecutionSpace const &space, View keys,
   if (n == 0)
     return;
 
-  // We don't want to sort the keys, so we create a copy.
-  auto keys_clone = KokkosExt::clone(space, keys);
   if constexpr (sizeof...(OtherViews) == 1 &&
                 std::tuple_element_t<0, std::tuple<OtherViews...>>::rank == 1)
   {
     // If there's only one 1D view to process, we can avoid computing the
     // permutation.
-    KokkosExt::sortByKey(space, keys_clone, other_views...);
+    KokkosExt::sortByKey(space, keys, other_views...);
   }
   else
   {
-    auto const permutation = ArborX::Details::sortObjects(space, keys_clone);
+    auto const permutation = ArborX::Details::sortObjects(space, keys);
 
     // Call applyPermutation for every entry in the parameter pack.
     // We need to use the comma operator here since the function returns void.

diff --git a/test/tstDetailsDistributedTreeImpl.cpp b/test/tstDetailsDistributedTreeImpl.cpp
@@ -68,12 +68,8 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(sort_results, DeviceType, ARBORX_DEVICE_TYPES)
   Kokkos::deep_copy(ranks, ranks_host);
 
   using ExecutionSpace = typename DeviceType::execution_space;
-  ArborX::Details::DistributedTree::sortResults(ExecutionSpace{}, ids, results,
-                                                ranks);
-
-  // COMMENT: ids are untouched
-  Kokkos::deep_copy(ids_host, ids);
-  BOOST_TEST(ids_host == ids_, tt::per_element());
+  ArborX::Details::DistributedTree::sortResultsByKey(ExecutionSpace{}, ids,
+                                                     results, ranks);
 
   Kokkos::deep_copy(results_host, results);
   Kokkos::deep_copy(ranks_host, ranks);
@@ -85,7 +81,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(sort_results, DeviceType, ARBORX_DEVICE_TYPES)
     }
 
   Kokkos::View<int *, DeviceType> not_sized_properly("", m);
-  BOOST_CHECK_THROW(ArborX::Details::DistributedTree::sortResults(
+  BOOST_CHECK_THROW(ArborX::Details::DistributedTree::sortResultsByKey(
                         ExecutionSpace{}, ids, not_sized_properly),
                     ArborX::SearchException);
 }
@@ -109,8 +105,8 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(sort_results_2d, DeviceType, ARBORX_DEVICE_TYPES)
   Kokkos::deep_copy(results_2d, results_2d_host);
 
   using ExecutionSpace = typename DeviceType::execution_space;
-  ArborX::Details::DistributedTree::sortResults(ExecutionSpace{}, ids,
-                                                results_2d);
+  ArborX::Details::DistributedTree::sortResultsByKey(ExecutionSpace{}, ids,
+                                                     results_2d);
 
   Kokkos::deep_copy(results_2d_host, results_2d);
   for (int i = 0; i < 5; ++i)
@@ -146,8 +142,8 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(sort_results_3d, DeviceType, ARBORX_DEVICE_TYPES)
   Kokkos::deep_copy(results_3d, results_3d_host);
 
   using ExecutionSpace = typename DeviceType::execution_space;
-  ArborX::Details::DistributedTree::sortResults(ExecutionSpace{}, ids,
-                                                results_3d);
+  ArborX::Details::DistributedTree::sortResultsByKey(ExecutionSpace{}, ids,
+                                                     results_3d);
 
   Kokkos::deep_copy(results_3d_host, results_3d);
   for (int i = 0; i < 5; ++i)