[SYCL] Key/Value sorting with fixed-size private array input

againull · againull · commit 665a5c867a05 · 2024-07-02T13:03:51.000-07:00
diff --git a/sycl/include/sycl/detail/group_sort_impl.hpp b/sycl/include/sycl/detail/group_sort_impl.hpp
@@ -17,6 +17,7 @@
 #include <sycl/group_barrier.hpp>
 #include <sycl/sycl_span.hpp>
 
+#include <iterator>
 #include <memory>
 
 namespace sycl {
@@ -52,6 +53,86 @@ static __SYCL_ALWAYS_INLINE T *align_scratch(sycl::span<std::byte> scratch,
 }
 #endif
 
+template <typename T1, typename T2> class key_value_iterator {
+public:
+  key_value_iterator(T1 *Keys, T2 *Values) : KeyValue{Keys, Values} {}
+
+  using difference_type = std::ptrdiff_t;
+  using value_type = std::tuple<T1, T2>;
+  using reference = std::tuple<T1 &, T2 &>;
+  using pointer = std::tuple<T1 *, T2 *>;
+  using iterator_category = std::random_access_iterator_tag;
+
+  reference operator*() const {
+    return std::tie(*(std::get<0>(KeyValue)), *(std::get<1>(KeyValue)));
+  }
+
+  reference operator[](difference_type i) const { return *(*this + i); }
+
+  difference_type operator-(const key_value_iterator &it) const {
+    return std::get<0>(KeyValue) - std::get<0>(it.KeyValue);
+  }
+
+  key_value_iterator &operator+=(difference_type i) {
+    KeyValue =
+        std::make_tuple(std::get<0>(KeyValue) + i, std::get<1>(KeyValue) + i);
+    return *this;
+  }
+  key_value_iterator &operator-=(difference_type i) { return *this += -i; }
+  key_value_iterator &operator++() { return *this += 1; }
+  key_value_iterator &operator--() { return *this -= 1; }
+  std::tuple<T1 *, T2 *> base() const { return KeyValue; }
+  key_value_iterator operator++(int) {
+    key_value_iterator it(*this);
+    ++(*this);
+    return it;
+  }
+  key_value_iterator operator--(int) {
+    key_value_iterator it(*this);
+    --(*this);
+    return it;
+  }
+
+  key_value_iterator operator-(difference_type i) const {
+    key_value_iterator it(*this);
+    return it -= i;
+  }
+  key_value_iterator operator+(difference_type i) const {
+    key_value_iterator it(*this);
+    return it += i;
+  }
+  friend key_value_iterator operator+(difference_type i,
+                                      const key_value_iterator &it) {
+    return it + i;
+  }
+
+  bool operator==(const key_value_iterator &it) const {
+    return *this - it == 0;
+  }
+
+  bool operator!=(const key_value_iterator &it) const { return !(*this == it); }
+  bool operator<(const key_value_iterator &it) const { return *this - it < 0; }
+  bool operator>(const key_value_iterator &it) const { return it < *this; }
+  bool operator<=(const key_value_iterator &it) const { return !(*this > it); }
+  bool operator>=(const key_value_iterator &it) const { return !(*this < it); }
+
+private:
+  std::tuple<T1 *, T2 *> KeyValue;
+};
+
+template <typename T> void swap(T &first, T &second) {
+  std::swap(first, second);
+}
+
+// Swap tuples of references.
+template <template <typename...> class Tuple, typename... T>
+void swap(Tuple<T &...> &&first, Tuple<T &...> &&second) {
+  auto lhs = first;
+  auto rhs = second;
+  // Do std::swap for each element of the tuple.
+  std::swap(lhs, rhs);
+}
+
 // ---- merge sort implementation
 
 // following two functions could be useless if std::[lower|upper]_bound worked
@@ -81,15 +162,6 @@ size_t upper_bound(Acc acc, const size_t first, const size_t last,
                              [comp](auto x, auto y) { return !comp(y, x); });
 }
 
-// swap for all data types including tuple-like types
-template <typename T> void swap_tuples(T &a, T &b) { std::swap(a, b); }
-
-template <template <typename...> class TupleLike, typename T1, typename T2>
-void swap_tuples(TupleLike<T1, T2> &&a, TupleLike<T1, T2> &&b) {
-  std::swap(std::get<0>(a), std::get<0>(b));
-  std::swap(std::get<1>(a), std::get<1>(b));
-}
-
 template <typename Iter> struct GetValueType {
   using type = typename std::iterator_traits<Iter>::value_type;
 };
@@ -205,18 +277,20 @@ void bubble_sort(Iter first, const size_t begin, const size_t end,
   if (begin < end) {
     for (size_t i = begin; i < end; ++i) {
       // Handle intermediate items
-      for (size_t idx = i + 1; idx < end; ++idx) {
-        if (comp(first[idx], first[i])) {
-          detail::swap_tuples(first[i], first[idx]);
+      for (size_t idx = begin; idx < begin + (end - 1 - i); ++idx) {
+        auto refs_1 = first[idx + 1];
+        auto refs_2 = first[idx];
+        if (comp(refs_1, refs_2)) {
+          detail::swap(first[idx], first[idx + 1]);
         }
       }
     }
   }
 }
 
-template <typename Group, typename Iter, typename T, typename Compare>
+template <typename Group, typename Iter, typename ScratchIter, typename Compare>
 void merge_sort(Group group, Iter first, const size_t n, Compare comp,
-                T *scratch) {
+                ScratchIter scratch) {
   const size_t idx = group.get_local_linear_id();
   const size_t local = group.get_local_range().size();
   const size_t chunk = (n - 1) / local + 1;
diff --git a/sycl/include/sycl/ext/oneapi/experimental/group_helpers_sorters.hpp b/sycl/include/sycl/ext/oneapi/experimental/group_helpers_sorters.hpp
@@ -113,7 +113,7 @@ template <typename Compare = std::less<>> class default_sorter {
                   [[maybe_unused]] Ptr last) {
 #ifdef __SYCL_DEVICE_ONLY__
     using T = typename sycl::detail::GetValueType<Ptr>::type;
-    size_t n = last - first;
+    size_t n = std::distance(first, last);
     T *scratch_begin = sycl::detail::align_scratch<T>(scratch, g, n);
     sycl::detail::merge_sort(g, first, n, comp, scratch_begin);
 #else
@@ -206,8 +206,8 @@ class radix_sorter {
     sycl::detail::privateDynamicSort</*is_key_value=*/false,
                                      OrderT == sorting_order::ascending,
                                      /*empty*/ 1, BitsPerPass>(
-        g, first, /*empty*/ first, last - first, scratch.data(), first_bit,
-        last_bit);
+        g, first, /*empty*/ first, std::distance(first, last), scratch.data(),
+        first_bit, last_bit);
 #else
     throw sycl::exception(
         std::error_code(PI_ERROR_INVALID_DEVICE, sycl::sycl_category()),
@@ -266,7 +266,7 @@ template <typename CompareT = std::less<>> class joint_sorter {
                   [[maybe_unused]] Ptr last) {
 #ifdef __SYCL_DEVICE_ONLY__
     using T = typename sycl::detail::GetValueType<Ptr>::type;
-    size_t n = last - first;
+    size_t n = std::distance(first, last);
     T *scratch_begin = sycl::detail::align_scratch<T>(scratch, g, n);
     sycl::detail::merge_sort(g, first, n, comp, scratch_begin);
 #else
@@ -374,6 +374,68 @@ class group_key_value_sorter {
         g, KeyValue(key, value));
   }
 
+  template <typename Group, typename Properties>
+  void operator()(Group g, sycl::span<KeyTy, ElementsPerWorkItem> keys,
+                  sycl::span<ValueTy, ElementsPerWorkItem> values,
+                  Properties property) {
+#ifdef __SYCL_DEVICE_ONLY__
+    auto range_size = g.get_local_linear_range();
+    std::size_t local_id = g.get_local_linear_id();
+    auto number_of_elements = range_size * ElementsPerWorkItem;
+    KeyTy *keys_scratch_begin =
+        sycl::detail::align_scratch<KeyTy>(scratch, g, number_of_elements);
+    size_t KeysScratchSpace =
+        number_of_elements * sizeof(KeyTy) + alignof(KeyTy);
+    size_t ValuesScratchSpace =
+        number_of_elements * sizeof(ValueTy) + alignof(ValueTy);
+    ValueTy *values_scratch_begin = sycl::detail::align_scratch<ValueTy>(
+        scratch.subspan(/* offset */ KeysScratchSpace,
+                        /* extent */ ValuesScratchSpace),
+        g, number_of_elements);
+
+    KeyTy *keys_scratch_temp_begin = sycl::detail::align_scratch<KeyTy>(
+        scratch.subspan(KeysScratchSpace + ValuesScratchSpace,
+                        KeysScratchSpace),
+        g, number_of_elements);
+    ValueTy *values_scratch_temp_begin = sycl::detail::align_scratch<ValueTy>(
+        scratch.subspan(2 * KeysScratchSpace + ValuesScratchSpace,
+                        ValuesScratchSpace),
+        g, number_of_elements);
+
+    std::size_t shift{};
+    for (std::uint32_t i = 0; i < ElementsPerWorkItem; ++i) {
+      if constexpr (detail::isInputBlocked(property)) {
+        shift = local_id * ElementsPerWorkItem + i;
+      } else {
+        shift = i * range_size + local_id;
+      }
+      keys_scratch_begin[shift] = keys[i];
+      values_scratch_begin[shift] = values[i];
+    }
+
+    auto scratch_begin = sycl::detail::key_value_iterator(keys_scratch_begin,
+                                                          values_scratch_begin);
+    auto scratch_temp_begin = sycl::detail::key_value_iterator(
+        keys_scratch_temp_begin, values_scratch_temp_begin);
+    sycl::detail::merge_sort(
+        g, scratch_begin, number_of_elements,
+        [this](auto x, auto y) { return comp(std::get<0>(x), std::get<0>(y)); },
+        scratch_temp_begin);
+
+    // from temp
+    for (std::uint32_t i = 0; i < ElementsPerWorkItem; ++i) {
+      if constexpr (detail::isOutputBlocked(property)) {
+        shift = local_id * ElementsPerWorkItem + i;
+      } else {
+        shift = i * range_size + local_id;
+      }
+
+      keys[i] = std::get<0>(scratch_begin[shift]);
+      values[i] = std::get<1>(scratch_begin[shift]);
+    }
+#endif
+  }
+
   static std::size_t memory_required(sycl::memory_scope scope,
                                      std::size_t range_size) {
     return group_sorter<std::tuple<KeyTy, ValueTy>, CompareT,
@@ -422,8 +484,8 @@ class joint_sorter {
     sycl::detail::privateDynamicSort</*is_key_value=*/false,
                                      OrderT == sorting_order::ascending,
                                      /*empty*/ 1, BitsPerPass>(
-        g, first, /*empty*/ first, last - first, scratch.data(), first_bit,
-        last_bit);
+        g, first, /*empty*/ first, std::distance(first, last), scratch.data(),
+        first_bit, last_bit);
 #else
     throw sycl::exception(
         std::error_code(PI_ERROR_INVALID_DEVICE, sycl::sycl_category()),
@@ -550,6 +612,18 @@ class group_key_value_sorter {
     return {key, val};
   }
 
+  template <typename Group, typename Properties>
+  void operator()(Group g, sycl::span<KeyTy, ElementsPerWorkItem> keys,
+                  sycl::span<ValueTy, ElementsPerWorkItem> vals,
+                  Properties properties) {
+#ifdef __SYCL_DEVICE_ONLY__
+    sycl::detail::privateStaticSort<
+        /*is_key_value=*/true, detail::isOutputBlocked(properties),
+        Order == sorting_order::ascending, ElementsPerWorkItem, bits>(
+        g, keys.data(), vals.data(), scratch.data(), first_bit, last_bit);
+#endif
+  }
+
   static constexpr std::size_t memory_required(sycl::memory_scope,
                                                std::size_t range_size) {
     return (std::max)(range_size * ElementsPerWorkItem *
diff --git a/sycl/include/sycl/ext/oneapi/experimental/group_sort.hpp b/sycl/include/sycl/ext/oneapi/experimental/group_sort.hpp
@@ -81,6 +81,21 @@ struct is_key_value_sorter<
                        std::tuple<Key, Value>> &&
         sycl::is_group_v<Group>>> : std::true_type {};
 
+template <typename Sorter, typename Group, typename Key, typename Value,
+          typename Properties, size_t ElementsPerWorkItem, typename = void>
+struct is_array_key_value_sorter : std::false_type {};
+
+template <typename Sorter, typename Group, typename Key, typename Value,
+          typename Properties, size_t ElementsPerWorkItem>
+struct is_array_key_value_sorter<
+    Sorter, Group, Key, Value, Properties, ElementsPerWorkItem,
+    std::enable_if_t<
+        std::is_same_v<std::invoke_result_t<
+                           Sorter, Group, sycl::span<Key, ElementsPerWorkItem>,
+                           sycl::span<Value, ElementsPerWorkItem>, Properties>,
+                       void> &&
+        sycl::is_group_v<Group>>> : std::true_type {};
+
 } // namespace detail
 
 // ---- sort_over_group
@@ -239,6 +254,63 @@ sort_key_value_over_group(
           exec.get_memory()));
 }
 
+template <std::size_t ElementsPerWorkItem, typename Group, typename T,
+          typename U, typename ArraySorter,
+          typename Properties = ext::oneapi::experimental::empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_property_list_v<
+        std::decay_t<Properties>> &&
+        detail::is_array_key_value_sorter<ArraySorter, Group, T, U, Properties,
+                                          ElementsPerWorkItem>::value,
+    void>
+sort_key_value_over_group(Group group, sycl::span<T, ElementsPerWorkItem> keys,
+                          sycl::span<U, ElementsPerWorkItem> values,
+                          ArraySorter array_sorter,
+                          Properties properties = {}) {
+  array_sorter(group, keys, values, properties);
+}
+
+template <typename Group, typename T, typename U, std::size_t Extent,
+          std::size_t ElementsPerWorkItem, typename Compare,
+          typename Properties = ext::oneapi::experimental::empty_properties_t>
+std::enable_if_t<
+    sycl::ext::oneapi::experimental::is_property_list_v<
+        std::decay_t<Properties>> &&
+        !sycl::ext::oneapi::experimental::is_property_list_v<
+            std::decay_t<Compare>> &&
+        !detail::is_array_key_value_sorter<Compare, Group, T, U, Properties,
+                                           ElementsPerWorkItem>::value,
+    void>
+sort_key_value_over_group(
+    experimental::group_with_scratchpad<Group, Extent> exec,
+    sycl::span<T, ElementsPerWorkItem> keys,
+    sycl::span<U, ElementsPerWorkItem> values, Compare comp,
+    Properties properties = {}) {
+  return experimental::sort_key_value_over_group(
+      exec.get_group(), keys, values,
+      typename experimental::default_sorters::group_key_value_sorter<
+          T, U, Compare, ElementsPerWorkItem>(exec.get_memory(), comp),
+      properties);
+}
+
+template <typename Group, typename T, typename U, std::size_t Extent,
+          std::size_t ElementsPerWorkItem,
+          typename Properties = ext::oneapi::experimental::empty_properties_t>
+std::enable_if_t<sycl::is_group_v<std::decay_t<Group>> &&
+                     sycl::ext::oneapi::experimental::is_property_list_v<
+                         std::decay_t<Properties>>,
+                 void>
+sort_key_value_over_group(
+    experimental::group_with_scratchpad<Group, Extent> exec,
+    sycl::span<T, ElementsPerWorkItem> keys,
+    sycl::span<U, ElementsPerWorkItem> values, Properties properties = {}) {
+  return experimental::sort_key_value_over_group(
+      exec.get_group(), keys, values,
+      typename experimental::default_sorters::group_key_value_sorter<
+          T, U, std::less<>, ElementsPerWorkItem>(exec.get_memory()),
+      properties);
+}
+
 } // namespace ext::oneapi::experimental
 } // namespace _V1
 } // namespace sycl
diff --git a/sycl/test-e2e/GroupAlgorithm/SYCL2020/group_sort/array_input_sort.cpp b/sycl/test-e2e/GroupAlgorithm/SYCL2020/group_sort/array_input_sort.cpp
@@ -209,9 +209,9 @@ template <class T> void RunOverType(sycl::queue &Q, size_t DataSize) {
   for (T &Elem : ArrayDataRandom)
     Elem = T(distribution(generator));
 
-  auto blocked = oneapi_exp::properties{oneapi_exp::input_data_placement<
+  auto blocked = oneapi_exp::properties{oneapi_exp::output_data_placement<
       oneapi_exp::group_algorithm_data_placement::blocked>};
-  auto striped = oneapi_exp::properties{oneapi_exp::input_data_placement<
+  auto striped = oneapi_exp::properties{oneapi_exp::output_data_placement<
       oneapi_exp::group_algorithm_data_placement::striped>};
   RunOnData<UseGroupT::WorkGroup, 1, PerWI>(Q, ArrayDataRandom, std::less<T>{},
                                             blocked);
diff --git a/sycl/test-e2e/GroupAlgorithm/SYCL2020/group_sort/key_value_array_input_sort.cpp b/sycl/test-e2e/GroupAlgorithm/SYCL2020/group_sort/key_value_array_input_sort.cpp
