[SYCL] Follow up fixes for group_sort extension

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>{});