[SYCL] Update group_broadcast to support vec types (#8886)

Update the `OpGroupBroadcast` definition to allow vector value and
result types. Add special case for `vec` types to enable
`GenericBroadcast` since certain backends don't support vector types in
`OpGroupBroadcast`.

Signed-off-by: Michael Aziz <[email protected]>

Author: 0x12CC

clang/lib/Sema/SPIRVBuiltins.td

@@ -926,9 +926,9 @@ foreach name = ["GroupAll", "GroupAny"] in {
 
 foreach name = ["GroupBroadcast"] in {
   foreach IDType = TLAllInts.List in {
-    def : SPVBuiltin<name, [AGenType1, UInt, AGenType1, IDType], Attr.Convergent>;
-    def : SPVBuiltin<name, [AGenType1, UInt, AGenType1, VectorType<IDType, 2>], Attr.Convergent>;
-    def : SPVBuiltin<name, [AGenType1, UInt, AGenType1, VectorType<IDType, 3>], Attr.Convergent>;
+    def : SPVBuiltin<name, [AGenTypeN, UInt, AGenTypeN, IDType], Attr.Convergent>;
+    def : SPVBuiltin<name, [AGenTypeN, UInt, AGenTypeN, VectorType<IDType, 2>], Attr.Convergent>;
+    def : SPVBuiltin<name, [AGenTypeN, UInt, AGenTypeN, VectorType<IDType, 3>], Attr.Convergent>;
     def : SPVBuiltin<name, [Bool, UInt, Bool, IDType], Attr.Convergent>;
     def : SPVBuiltin<name, [Bool, UInt, Bool, VectorType<IDType, 2>], Attr.Convergent>;
     def : SPVBuiltin<name, [Bool, UInt, Bool, VectorType<IDType, 3>], Attr.Convergent>;

sycl/include/sycl/detail/spirv.hpp

@@ -103,10 +103,12 @@ template <typename Group> bool GroupAny(bool pred) {
 }
 
 // Native broadcasts map directly to a SPIR-V GroupBroadcast intrinsic
-// FIXME: Do not special-case for half once all backends support all data types.
+// FIXME: Do not special-case for half or vec once all backends support all data
+// types.
 template <typename T>
-using is_native_broadcast = bool_constant<detail::is_arithmetic<T>::value &&
-                                          !std::is_same<T, half>::value>;
+using is_native_broadcast =
+    bool_constant<detail::is_arithmetic<T>::value &&
+                  !std::is_same<T, half>::value && !detail::is_vec<T>::value>;
 
 template <typename T, typename IdT = size_t>
 using EnableIfNativeBroadcast = detail::enable_if_t<

sycl/test-e2e/GroupAlgorithm/SYCL2020/group_broadcast.cpp

@@ -12,6 +12,18 @@
 #include <sycl/sycl.hpp>
 using namespace sycl;
 
+template <typename T> bool equal(const T &a, const T &b) { return a == b; }
+
+template <typename T, int N>
+bool equal(const vec<T, N> &a, const vec<T, N> &b) {
+  for (int i = 0; i < N; i++) {
+    if (a[i] != b[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
 template <typename kernel_name, typename InputContainer,
           typename OutputContainer>
 void test(queue q, InputContainer input, OutputContainer output) {
@@ -37,9 +49,9 @@ void test(queue q, InputContainer input, OutputContainer output) {
       });
     });
   }
-  assert(output[0] == input[0]);
-  assert(output[1] == input[1 * G + 2]);
-  assert(output[2] == input[2 * G + 1]);
+  assert(equal(output[0], input[0]));
+  assert(equal(output[1], input[1 * G + 2]));
+  assert(equal(output[2], input[2 * G + 1]));
 }
 
 int main() {
@@ -71,6 +83,17 @@ int main() {
     test<class KernelName_NrqELzFQToOSPsRNMi>(q, input, output);
   }
 
+  // Test vector types
+  {
+    std::array<vec<int, 4>, N> input;
+    std::array<vec<int, 4>, 3> output;
+    for (int i = 0; i < N; ++i) {
+      input[i] = vec<int, 4>{i, i, i, i};
+    }
+    std::fill(output.begin(), output.end(), vec<int, 4>{0, 0, 0, 0});
+    test<class KernelName_VectorGroupBroadcast>(q, input, output);
+  }
+
   // Test user-defined type
   // - Use complex as a proxy for this
   // - Test float and double to test 64-bit and 128-bit types