[SYCL] Add support for key/value sorting APIs

againull · againull · commit 6aa79957f141 · 2024-06-10T10:33:33.000-07:00
diff --git a/sycl/include/sycl/detail/group_sort_impl.hpp b/sycl/include/sycl/detail/group_sort_impl.hpp
@@ -578,7 +578,7 @@ template <size_t items_per_work_item, uint32_t radix_bits, bool is_comp_asc,
           typename ValsT, typename GroupT>
 void performRadixIterStaticSize(GroupT group, const uint32_t radix_iter,
                                 const uint32_t last_iter, KeysT *keys,
-                                ValsT vals, const ScratchMemory &memory) {
+                                ValsT *vals, const ScratchMemory &memory) {
   const uint32_t radix_states = getStatesInBits(radix_bits);
   const size_t wgsize = group.get_local_linear_range();
   const size_t idx = group.get_local_linear_id();
diff --git a/sycl/include/sycl/ext/oneapi/experimental/group_helpers_sorters.hpp b/sycl/include/sycl/ext/oneapi/experimental/group_helpers_sorters.hpp
@@ -343,6 +343,41 @@ class group_sorter {
   }
 };
 
+template <typename T, typename U, typename CompareT = std::less<>,
+          std::size_t ElementsPerWorkItem = 1>
+class group_key_value_sorter {
+  CompareT comp;
+  sycl::span<std::byte> scratch;
+
+public:
+  template <std::size_t Extent>
+  group_key_value_sorter(sycl::span<std::byte, Extent> scratch_,
+                         CompareT comp_ = {})
+      : comp(comp_), scratch(scratch_) {}
+
+  template <typename Group>
+  std::tuple<T, U> operator()(Group g, T key, U value) {
+    static_assert(ElementsPerWorkItem == 1,
+                  "ElementsPerWorkItem must be equal 1");
+
+    using KeyValue = std::tuple<T, U>;
+    auto this_comp = this->comp;
+    auto comp_key_value = [this_comp](const KeyValue &lhs,
+                                      const KeyValue &rhs) {
+      return this_comp(std::get<0>(lhs), std::get<0>(rhs));
+    };
+    return group_sorter<KeyValue, decltype(comp_key_value),
+                        ElementsPerWorkItem>(scratch, comp_key_value)(
+        g, KeyValue(key, value));
+  }
+
+  static constexpr std::size_t memory_required(sycl::memory_scope scope,
+                                               std::size_t range_size) {
+    return group_sorter<std::tuple<T, U>, CompareT,
+                        ElementsPerWorkItem>::memory_required(scope,
+                                                              range_size);
+  }
+};
 } // namespace default_sorters
 
 // Radix sorters provided by the second version of the extension specification.
@@ -455,6 +490,56 @@ class group_sorter {
   }
 };
 
+template <typename T, typename U,
+          sorting_order Order = sorting_order::ascending,
+          size_t ElementsPerWorkItem = 1, unsigned int BitsPerPass = 4>
+class group_key_value_sorter {
+  sycl::span<std::byte> scratch;
+  uint32_t first_bit;
+  uint32_t last_bit;
+
+  static constexpr uint32_t bits = BitsPerPass;
+  using bitset_t = std::bitset<sizeof(T) * CHAR_BIT>;
+
+public:
+  template <std::size_t Extent>
+  group_key_value_sorter(sycl::span<std::byte, Extent> scratch_,
+                         const bitset_t mask = bitset_t{}.set())
+      : scratch(scratch_) {
+    static_assert(
+        (std::is_arithmetic<T>::value || std::is_same<T, sycl::half>::value),
+        "radix sort is not usable");
+    for (first_bit = 0; first_bit < mask.size() && !mask[first_bit];
+         ++first_bit)
+      ;
+    for (last_bit = first_bit; last_bit < mask.size() && mask[last_bit];
+         ++last_bit)
+      ;
+  }
+
+  template <typename Group>
+  std::tuple<T, U> operator()([[maybe_unused]] Group g, T key, U val) {
+    static_assert(ElementsPerWorkItem == 1, "ElementsPerWorkItem must be 1");
+    T key_result[]{key};
+    U val_result[]{val};
+#ifdef __SYCL_DEVICE_ONLY__
+    sycl::detail::privateStaticSort<
+        /*is_key_value=*/true,
+        /*is_blocked=*/true, Order == sorting_order::ascending, 1, bits>(
+        g, key_result, val_result, scratch.data(), first_bit, last_bit);
+#endif
+    key = key_result[0];
+    val = val_result[0];
+    return {key, val};
+  }
+
+  static constexpr std::size_t memory_required(sycl::memory_scope scope,
+                                               std::size_t range_size) {
+    return (std::max)(range_size * ElementsPerWorkItem *
+                          (sizeof(T) + sizeof(U)),
+                      range_size * (1 << bits) * sizeof(uint32_t));
+  }
+};
 } // namespace radix_sorters
 
 } // namespace ext::oneapi::experimental
diff --git a/sycl/include/sycl/ext/oneapi/experimental/group_sort.hpp b/sycl/include/sycl/ext/oneapi/experimental/group_sort.hpp
@@ -68,6 +68,29 @@ struct is_sorter_impl<Sorter, Group, Ptr,
 template <typename Sorter, typename Group, typename ValOrPtr>
 struct is_sorter : decltype(is_sorter_impl<Sorter, Group, ValOrPtr>::test(0)) {
 };
+
+template <typename Sorter, typename Group, typename Key, typename Value>
+struct is_key_value_sorter_impl {
+  template <typename G>
+  using is_expected_return_type =
+      typename std::is_same<std::tuple<Key, Value>,
+                            decltype(std::declval<Sorter>()(
+                                std::declval<G>(), std::declval<Key>(),
+                                std::declval<Value>()))>;
+
+  template <typename G = Group>
+  static decltype(std::integral_constant<bool,
+                                         is_expected_return_type<G>::value &&
+                                             sycl::is_group_v<G>>{})
+  test(int);
+
+  template <typename = Group> static std::false_type test(...);
+};
+
+template <typename Sorter, typename Group, typename Key, typename Value>
+struct is_key_value_sorter
+    : decltype(is_key_value_sorter_impl<Sorter, Group, Key, Value>::test(0)){};
+
 } // namespace detail
 
 // ---- sort_over_group
@@ -131,6 +154,43 @@ joint_sort(experimental::group_with_scratchpad<Group, Extent> exec, Iter first,
              default_sorters::joint_sorter<>(exec.get_memory()));
 }
 
+template <typename Group, typename T, typename U, typename Sorter>
+std::enable_if_t<detail::is_key_value_sorter<Sorter, Group, T, U>::value,
+                 std::tuple<T, U>>
+sort_key_value_over_group([[maybe_unused]] Group g, [[maybe_unused]] T key,
+                          [[maybe_unused]] U value,
+                          [[maybe_unused]] Sorter sorter) {
+#ifdef __SYCL_DEVICE_ONLY__
+  return sorter(g, key, value);
+#else
+  throw sycl::exception(
+      std::error_code(PI_ERROR_INVALID_DEVICE, sycl::sycl_category()),
+      "Group algorithms are not supported on host device.");
+#endif
+}
+
+template <typename Group, typename T, typename U, typename Compare,
+          std::size_t Extent>
+std::enable_if_t<!detail::is_key_value_sorter<Compare, Group, T, U>::value,
+                 std::tuple<T, U>>
+sort_key_value_over_group(
+    experimental::group_with_scratchpad<Group, Extent> exec, T key, U value,
+    Compare comp) {
+  return sort_key_value_over_group(
+      exec.get_group(), key, value,
+      default_sorters::group_key_value_sorter<T, U, Compare>(exec.get_memory(),
+                                                             comp));
+}
+
+template <typename T, typename U, typename Group, std::size_t Extent>
+std::enable_if_t<sycl::is_group_v<std::decay_t<Group>>, std::tuple<T, U>>
+sort_key_value_over_group(
+    experimental::group_with_scratchpad<Group, Extent> exec, T key, U value) {
+  return sort_key_value_over_group(
+      exec.get_group(), key, value,
+      default_sorters::group_key_value_sorter<T, U>(exec.get_memory()));
+}
+
 } // namespace ext::oneapi::experimental
 } // namespace _V1
 } // namespace sycl
diff --git a/sycl/test-e2e/GroupAlgorithm/SYCL2020/sort.cpp b/sycl/test-e2e/GroupAlgorithm/SYCL2020/sort.cpp
