[SYCL] Make joint_reduce work with sub_group (#8786)

Note: the unqualified name lookup of joint_reduce in the overload of
joint_reduce without an init param was not finding the overload of
joint_reduce with an init param (because that declaration was located
after it), so it searched for joint_reduce via ADL. With sycl::group,
ADL can find both overloads of joint_reduce, but with sycl::sub_group =
sycl::ext::oneapi::sub_group, ADL finds no joint_reduce in
sycl::ext::oneapi.

Fixes #8348

---------

Signed-off-by: Cai, Justin <justin.cai@intel.com>

Author: jzc

sycl/include/sycl/group_algorithm.hpp

@@ -315,29 +315,6 @@ reduce_over_group(Group g, V x, T init, BinaryOperation binary_op) {
 }
 
 // ---- joint_reduce
-template <typename Group, typename Ptr, class BinaryOperation>
-detail::enable_if_t<
-    (is_group_v<std::decay_t<Group>> && detail::is_pointer<Ptr>::value &&
-     detail::is_arithmetic_or_complex<
-         typename detail::remove_pointer<Ptr>::type>::value &&
-     detail::is_plus_or_multiplies_if_complex<
-         typename detail::remove_pointer<Ptr>::type, BinaryOperation>::value),
-    typename detail::remove_pointer<Ptr>::type>
-joint_reduce(Group g, Ptr first, Ptr last, BinaryOperation binary_op) {
-#ifdef __SYCL_DEVICE_ONLY__
-  using T = typename detail::remove_pointer<Ptr>::type;
-  T init = detail::identity_for_ga_op<T, BinaryOperation>();
-  return joint_reduce(g, first, last, init, binary_op);
-#else
-  (void)g;
-  (void)first;
-  (void)last;
-  (void)binary_op;
-  throw runtime_error("Group algorithms are not supported on host.",
-                      PI_ERROR_INVALID_DEVICE);
-#endif
-}
-
 template <typename Group, typename Ptr, typename T, class BinaryOperation>
 detail::enable_if_t<
     (is_group_v<std::decay_t<Group>> && detail::is_pointer<Ptr>::value &&
@@ -373,6 +350,29 @@ joint_reduce(Group g, Ptr first, Ptr last, T init, BinaryOperation binary_op) {
 #endif
 }
 
+template <typename Group, typename Ptr, class BinaryOperation>
+detail::enable_if_t<
+    (is_group_v<std::decay_t<Group>> && detail::is_pointer<Ptr>::value &&
+     detail::is_arithmetic_or_complex<
+         typename detail::remove_pointer<Ptr>::type>::value &&
+     detail::is_plus_or_multiplies_if_complex<
+         typename detail::remove_pointer<Ptr>::type, BinaryOperation>::value),
+    typename detail::remove_pointer<Ptr>::type>
+joint_reduce(Group g, Ptr first, Ptr last, BinaryOperation binary_op) {
+#ifdef __SYCL_DEVICE_ONLY__
+  using T = typename detail::remove_pointer<Ptr>::type;
+  T init = detail::identity_for_ga_op<T, BinaryOperation>();
+  return joint_reduce(g, first, last, init, binary_op);
+#else
+  (void)g;
+  (void)first;
+  (void)last;
+  (void)binary_op;
+  throw runtime_error("Group algorithms are not supported on host.",
+                      PI_ERROR_INVALID_DEVICE);
+#endif
+}
+
 // ---- any_of_group
 template <typename Group>
 detail::enable_if_t<is_group_v<std::decay_t<Group>>, bool>

sycl/test-e2e/GroupAlgorithm/reduce_sycl2020.cpp

@@ -34,13 +34,18 @@ void test(queue q, InputContainer input, OutputContainer output,
       cgh.parallel_for<SpecializationKernelName>(
           nd_range<1>(G, G), [=](nd_item<1> it) {
             group<1> g = it.get_group();
+            auto sg = it.get_sub_group();
             int lid = it.get_local_id(0);
             out[0] = reduce_over_group(g, in[lid], binary_op);
             out[1] = reduce_over_group(g, in[lid], init, binary_op);
             out[2] = joint_reduce(g, in.get_pointer(), in.get_pointer() + N,
                                   binary_op);
             out[3] = joint_reduce(g, in.get_pointer(), in.get_pointer() + N,
                                   init, binary_op);
+            out[4] = joint_reduce(sg, in.get_pointer(), in.get_pointer() + N,
+                                  binary_op);
+            out[5] = joint_reduce(sg, in.get_pointer(), in.get_pointer() + N,
+                                  init, binary_op);
           });
     });
   }
@@ -54,6 +59,10 @@ void test(queue q, InputContainer input, OutputContainer output,
          std::accumulate(input.begin(), input.end(), identity, binary_op));
   assert(output[3] ==
          std::accumulate(input.begin(), input.end(), init, binary_op));
+  assert(output[4] ==
+         std::accumulate(input.begin(), input.end(), identity, binary_op));
+  assert(output[5] ==
+         std::accumulate(input.begin(), input.end(), init, binary_op));
 }
 
 int main() {
@@ -65,7 +74,7 @@ int main() {
 
   constexpr int N = 128;
   std::array<int, N> input;
-  std::array<int, 4> output;
+  std::array<int, 6> output;
   std::iota(input.begin(), input.end(), 0);
   std::fill(output.begin(), output.end(), 0);
 
@@ -93,7 +102,7 @@ int main() {
   // sycl::plus binary operation.
 #ifdef SYCL_EXT_ONEAPI_COMPLEX_ALGORITHMS
   std::array<std::complex<float>, N> input_cf;
-  std::array<std::complex<float>, 4> output_cf;
+  std::array<std::complex<float>, 6> output_cf;
   std::iota(input_cf.begin(), input_cf.end(), 0);
   std::fill(output_cf.begin(), output_cf.end(), 0);
   test<class KernelNamePlusComplexF>(q, input_cf, output_cf,