Address review comments

algorithms/src/sorting/Kokkos_SortByKeyPublicAPI.hpp

@@ -29,7 +29,7 @@ namespace Kokkos::Experimental {
 
 template <class ExecutionSpace, class KeysDataType, class... KeysProperties,
           class ValuesDataType, class... ValuesProperties>
-void sort_by_key([[maybe_unused]] const ExecutionSpace& exec,
+void sort_by_key(const ExecutionSpace& exec,
                  Kokkos::View<KeysDataType, KeysProperties...>& keys,
                  Kokkos::View<ValuesDataType, ValuesProperties...>& values) {
   // constraints
@@ -48,8 +48,8 @@ void sort_by_key([[maybe_unused]] const ExecutionSpace& exec,
       "Kokkos::sort: execution space instance is not able to access "
       "the memory space of the values View argument!");
 
-  // FIXME: what's the right way to check this condition?
-  assert(values.extent(0) >= keys.extent(0));
+  if (values.size() != keys.size())
+    Kokkos::abort("values and keys extents must be the same");
 
   if (keys.extent(0) <= 1) {
     return;
@@ -66,7 +66,7 @@ void sort_by_key([[maybe_unused]] const ExecutionSpace& exec,
 template <class ExecutionSpace, class ComparatorType, class KeysDataType,
           class... KeysProperties, class ValuesDataType,
           class... ValuesProperties>
-void sort_by_key([[maybe_unused]] const ExecutionSpace& exec,
+void sort_by_key(const ExecutionSpace& exec,
                  Kokkos::View<KeysDataType, KeysProperties...>& keys,
                  Kokkos::View<ValuesDataType, ValuesProperties...>& values,
                  const ComparatorType& comparator) {
@@ -86,8 +86,8 @@ void sort_by_key([[maybe_unused]] const ExecutionSpace& exec,
       "Kokkos::sort: execution space instance is not able to access "
       "the memory space of the values View argument!");
 
-  // FIXME: what's the right way to check this condition?
-  assert(values.extent(0) >= keys.extent(0));
+  if (values.size() != keys.size())
+    Kokkos::abort("values and keys extents must be the same");
 
   if (keys.extent(0) <= 1) {
     return;

algorithms/src/sorting/impl/Kokkos_SortByKeyImpl.hpp

@@ -17,8 +17,6 @@
 #ifndef KOKKOS_SORT_BY_KEY_FREE_FUNCS_IMPL_HPP_
 #define KOKKOS_SORT_BY_KEY_FREE_FUNCS_IMPL_HPP_
 
-#include "../Kokkos_BinOpsPublicAPI.hpp"
-#include "../Kokkos_BinSortPublicAPI.hpp"
 #include <Kokkos_Core.hpp>
 
 #if defined(KOKKOS_ENABLE_CUDA)
@@ -160,7 +158,10 @@ void applyPermutation(ExecutionSpace const& space,
                       ViewType const& view) {
   static_assert(std::is_integral<typename PermutationView::value_type>::value);
 
-  auto view_copy = Kokkos::create_mirror(space, view);
+  auto view_copy = Kokkos::create_mirror(
+      Kokkos::view_alloc(space, typename ExecutionSpace::memory_space{},
+                         Kokkos::WithoutInitializing),
+      view);
   Kokkos::deep_copy(space, view_copy, view);
   Kokkos::parallel_for(
       "Kokkos::sort_by_key_via_sort::permute_" + view.label(),

algorithms/unit_tests/TestSortByKey.hpp

@@ -79,24 +79,12 @@ TEST(TEST_CATEGORY, SortByKeyEmptyView) {
 
   // does not matter if we use int or something else
   Kokkos::View<int *, ExecutionSpace> keys("keys", 0);
-  Kokkos::View<float *, ExecutionSpace> values("values", 2);
+  Kokkos::View<float *, ExecutionSpace> values("values", 0);
 
   ASSERT_NO_THROW(
       Kokkos::Experimental::sort_by_key(ExecutionSpace(), keys, values));
 }
 
-TEST(TEST_CATEGORY, SortByKeyKeysLargerThanValues) {
-  using ExecutionSpace = TEST_EXECSPACE;
-
-  // does not matter if we use int or something else
-  Kokkos::View<int *, ExecutionSpace> keys("keys", 3);
-  Kokkos::View<float *, ExecutionSpace> values("values", 1);
-
-  // FIXME depends on how we check (keys.extent(0) <= values.extent(0))
-  // ASSERT_THROW(Kokkos::Experimental::sort_by_key(ExecutionSpace(), keys,
-  // values));
-}
-
 TEST(TEST_CATEGORY, SortByKey) {
   using ExecutionSpace = TEST_EXECSPACE;
   using MemorySpace    = typename ExecutionSpace::memory_space;
@@ -173,5 +161,59 @@ TEST(TEST_CATEGORY, SortByKeyWithComparator) {
   }
 }
 
+TEST(TEST_CATEGORY, SortByKeyWithStrides) {
+  using ExecutionSpace = TEST_EXECSPACE;
+
+  ExecutionSpace space{};
+
+  auto const n = 10;
+  Kokkos::View<int ***, ExecutionSpace> keys("keys", n, n, n);
+  Kokkos::View<int ***, ExecutionSpace> values("values", n, n, n);
+
+  Kokkos::parallel_for(
+      "create_data",
+      Kokkos::MDRangePolicy<Kokkos::Rank<3>, ExecutionSpace>(space, {0, 0, 0},
+                                                             {n, n, n}),
+      KOKKOS_LAMBDA(int i, int j, int k) {
+        keys(i, j, k)   = n - i;
+        values(i, j, k) = i;
+      });
+
+  auto keys_sub   = Kokkos::subview(keys, Kokkos::ALL(), 0, 0);
+  auto values_sub = Kokkos::subview(values, Kokkos::ALL(), 0, 0);
+
+  auto keys_orig = Kokkos::create_mirror(space, keys_sub);
+  Kokkos::deep_copy(space, keys_orig, keys_sub);
+
+  Kokkos::Experimental::sort_by_key(space, keys_sub, values_sub);
+
+  unsigned int sort_fails = 0;
+  Kokkos::parallel_reduce(
+      Kokkos::RangePolicy<ExecutionSpace>(space, 0, n),
+      SortImpl::is_sorted_by_key_struct<ExecutionSpace, decltype(keys_sub),
+                                        decltype(values_sub)>(
+          keys_sub, keys_orig, values_sub),
+      sort_fails);
+
+  ASSERT_EQ(sort_fails, 0u);
+}
+
+#if !defined(NDEBUG) || defined(KOKKOS_ENABLE_DEBUG)
+TEST(TEST_CATEGORY, SortByKeyKeysLargerThanValues) {
+  using ExecutionSpace = TEST_EXECSPACE;
+
+  // does not matter if we use int or something else
+  Kokkos::View<int *, ExecutionSpace> keys("keys", 3);
+  Kokkos::View<float *, ExecutionSpace> values("values", 1);
+
+  ASSERT_DEATH(
+      Kokkos::Experimental::sort_by_key(ExecutionSpace(), keys, values),
+      "values and keys extents must be the same");
+  ASSERT_DEATH(Kokkos::Experimental::sort_by_key(ExecutionSpace(), keys, values,
+                                                 SortImpl::Greater{}),
+               "values and keys extents must be the same");
+}
+#endif
+
 }  // namespace Test
 #endif