[SYCL] Follow up fixes for group_sort extension (#14591)

* Change `memory_required` queries to return the exact size of local
memory which needs to be allocated for `group_sort` algorithm, currently
it is a bit more than required.
* Fix a bug in `align_key_value_scratch` where I didn't take into
account that `std::align` changes value of `KeysScratchSpace`. To find
scratch of the values I need to use original value of this variable.
* Fix mistake in the test where keys/values were mixed up. This is more
of a cosmetic change.
RunSortKeyValueOverGroup functions accepts two vectors - the first is
considered keys, and the second is considered values.
Verification is done in the same function (with the same assumption that
the first vector is keys and the second is values).
So it doesn't actually matter in which order aforementioned arguments
are provided, test still serves it purpose.

Author: againull

sycl/include/sycl/detail/group_sort_impl.hpp

@@ -72,7 +72,7 @@ align_key_value_scratch(sycl::span<std::byte> scratch, Group g,
     scratch_ptr =
         std::align(alignof(KeyTy), KeysSize, scratch_ptr, KeysScratchSpace);
     keys_scratch_begin = ::new (scratch_ptr) KeyTy[number_of_elements];
-    scratch_ptr = scratch.data() + KeysScratchSpace;
+    scratch_ptr = scratch.data() + KeysSize + alignof(KeyTy);
     scratch_ptr = std::align(alignof(ValueTy), ValuesSize, scratch_ptr,
                              ValuesScratchSpace);
     values_scratch_begin = ::new (scratch_ptr) ValueTy[number_of_elements];

sycl/include/sycl/ext/oneapi/experimental/group_helpers_sorters.hpp

@@ -341,10 +341,13 @@ class group_sorter {
 #endif
   }
 
-  static std::size_t memory_required(sycl::memory_scope scope,
+  static std::size_t memory_required([[maybe_unused]] sycl::memory_scope scope,
                                      size_t range_size) {
-    return 2 * joint_sorter<>::template memory_required<T>(
-                   scope, range_size * ElementsPerWorkItem);
+    // We need a space (in bytes) for the buffer of output values and the
+    // temporary buffer. Where number of elements in each buffer is range_size
+    // (group size) multiplied by elements per work item. Also we have to align
+    // these two buffers, so need an additional space of size alignof(T).
+    return 2 * range_size * ElementsPerWorkItem * sizeof(T) + alignof(T);
   }
 };
 
@@ -449,11 +452,18 @@ class group_key_value_sorter {
 #endif
   }
 
-  static std::size_t memory_required(sycl::memory_scope scope,
+  static std::size_t memory_required([[maybe_unused]] sycl::memory_scope scope,
                                      std::size_t range_size) {
-    return group_sorter<std::tuple<KeyTy, ValueTy>, CompareT,
-                        ElementsPerWorkItem>::memory_required(scope,
-                                                              range_size);
+    // We need a space (in bytes) for the following buffers:
+    // 1. Output buffer for keys and temporary buffer for keys.
+    // 2. Output buffer for values and temporary buffer for values.
+    // Where number of elements in each buffer is range_size (group size)
+    // multiplied by elements per work item. We have to align buffers of keys
+    // and buffers of values, so need an additional space equal to maximum
+    // between alignment requirements of types KeyTy and ValueTy.
+    return 2 * range_size * ElementsPerWorkItem *
+               (sizeof(KeyTy) + sizeof(ValueTy)) +
+           (std::max)(alignof(KeyTy), alignof(ValueTy));
   }
 };
 } // namespace default_sorters

sycl/test-e2e/GroupAlgorithm/SYCL2020/group_sort/key_value_sort.cpp

@@ -269,9 +269,9 @@ void RunOverType(sycl::queue &Q, size_t DataSize) {
   auto RunOnDataAndComp = [&](const std::vector<KeyTy> &Keys,
                               const std::vector<ValueTy> &Data,
                               const auto &Comparator) {
-    RunSortKeyValueOverGroup<UseGroupT::WorkGroup, 1>(Q, Data, Keys,
+    RunSortKeyValueOverGroup<UseGroupT::WorkGroup, 1>(Q, Keys, Data,
                                                       Comparator);
-    RunSortKeyValueOverGroup<UseGroupT::WorkGroup, 2>(Q, Data, Keys,
+    RunSortKeyValueOverGroup<UseGroupT::WorkGroup, 2>(Q, Keys, Data,
                                                       Comparator);
 
     if (Q.get_backend() == sycl::backend::ext_oneapi_cuda ||
@@ -280,8 +280,8 @@ void RunOverType(sycl::queue &Q, size_t DataSize) {
       return;
     }
 
-    RunSortKeyValueOverGroup<UseGroupT::SubGroup, 1>(Q, Data, Keys, Comparator);
-    RunSortKeyValueOverGroup<UseGroupT::SubGroup, 2>(Q, Data, Keys, Comparator);
+    RunSortKeyValueOverGroup<UseGroupT::SubGroup, 1>(Q, Keys, Data, Comparator);
+    RunSortKeyValueOverGroup<UseGroupT::SubGroup, 2>(Q, Keys, Data, Comparator);
   };
 
   RunOnDataAndComp(KeysRandom, DataRandom, std::less<KeyTy>{});