[SYCL][Matrix] Add support for more types to joint_matrix_mad (intel#4486)

With this patch, more cases for C=A*B+C can be realized:
1. A is uint8, B is int8
2. A is int8, B is uint8
3. A is uint8, B is uint8

Author: yubingex007-a11y

sycl/include/CL/__spirv/spirv_ops.hpp

@@ -50,6 +50,42 @@ __spirv_JointMatrixMadINTEL(
     __spv::__spirv_JointMatrixINTEL<T2, M, N, LC, S> *C,
     __spv::Scope::Flag Sc = __spv::Scope::Flag::Subgroup);
 
+template <typename T1, typename T2, typename T3, std::size_t M, std::size_t K,
+          std::size_t N, __spv::MatrixLayout LA = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LB = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LC = __spv::MatrixLayout::RowMajor,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern SYCL_EXTERNAL __spv::__spirv_JointMatrixINTEL<T3, M, N, LC, S> *
+__spirv_JointMatrixUUMadINTEL(
+    __spv::__spirv_JointMatrixINTEL<T1, M, K, LA, S> *A,
+    __spv::__spirv_JointMatrixINTEL<T2, K, N, LB, S> *B,
+    __spv::__spirv_JointMatrixINTEL<T3, M, N, LC, S> *C,
+    __spv::Scope::Flag Sc = __spv::Scope::Flag::Subgroup);
+
+template <typename T1, typename T2, typename T3, std::size_t M, std::size_t K,
+          std::size_t N, __spv::MatrixLayout LA = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LB = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LC = __spv::MatrixLayout::RowMajor,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern SYCL_EXTERNAL __spv::__spirv_JointMatrixINTEL<T3, M, N, LC, S> *
+__spirv_JointMatrixUSMadINTEL(
+    __spv::__spirv_JointMatrixINTEL<T1, M, K, LA, S> *A,
+    __spv::__spirv_JointMatrixINTEL<T2, K, N, LB, S> *B,
+    __spv::__spirv_JointMatrixINTEL<T3, M, N, LC, S> *C,
+    __spv::Scope::Flag Sc = __spv::Scope::Flag::Subgroup);
+
+template <typename T1, typename T2, typename T3, std::size_t M, std::size_t K,
+          std::size_t N, __spv::MatrixLayout LA = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LB = __spv::MatrixLayout::RowMajor,
+          __spv::MatrixLayout LC = __spv::MatrixLayout::RowMajor,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern SYCL_EXTERNAL __spv::__spirv_JointMatrixINTEL<T3, M, N, LC, S> *
+__spirv_JointMatrixSUMadINTEL(
+    __spv::__spirv_JointMatrixINTEL<T1, M, K, LA, S> *A,
+    __spv::__spirv_JointMatrixINTEL<T2, K, N, LB, S> *B,
+    __spv::__spirv_JointMatrixINTEL<T3, M, N, LC, S> *C,
+    __spv::Scope::Flag Sc = __spv::Scope::Flag::Subgroup);
+
 #ifndef __SPIRV_BUILTIN_DECLARATIONS__
 #error                                                                         \
     "SPIR-V built-ins are not available. Please set -fdeclare-spirv-builtins flag."

sycl/include/sycl/ext/oneapi/matrix/matrix-jit.hpp

@@ -160,16 +160,27 @@ joint_matrix_store(Group sg,
 #endif // __SYCL_DEVICE_ONLY__
 }
 
-template <typename Group, typename T1, typename T2, size_t M, size_t K,
-          size_t N, matrix_layout LayoutA, matrix_layout LayoutB,
+template <typename Group, typename T1, typename T2, typename T3, size_t M,
+          size_t K, size_t N, matrix_layout LayoutA, matrix_layout LayoutB,
           matrix_layout LayoutC>
-inline __SYCL_ALWAYS_INLINE joint_matrix<T2, M, N, LayoutC, Group>
+inline __SYCL_ALWAYS_INLINE joint_matrix<T3, M, N, LayoutC, Group>
 joint_matrix_mad(Group sg, joint_matrix<T1, M, K, LayoutA, Group> &mA,
-                 joint_matrix<T1, K, N, LayoutB, Group> &mB,
-                 joint_matrix<T2, M, N, LayoutC, Group> &mC) {
+                 joint_matrix<T2, K, N, LayoutB, Group> &mB,
+                 joint_matrix<T3, M, N, LayoutC, Group> &mC) {
 #ifdef __SYCL_DEVICE_ONLY__
-  joint_matrix<T2, M, N, LayoutC, Group> res(sg);
-  res.spvm = __spirv_JointMatrixMadINTEL(mA.spvm, mB.spvm, mC.spvm);
+  joint_matrix<T3, M, N, LayoutC, Group> res(sg);
+  if constexpr (std::is_same<T1, uint16_t>::value &&
+                std::is_same<T2, uint16_t>::value &&
+                std::is_same<T3, float>::value)
+    res.spvm = __spirv_JointMatrixMadINTEL(mA.spvm, mB.spvm, mC.spvm);
+  else if constexpr (std::is_unsigned<T1>::value && std::is_unsigned<T2>::value)
+    res.spvm = __spirv_JointMatrixUUMadINTEL(mA.spvm, mB.spvm, mC.spvm);
+  else if constexpr (std::is_signed<T1>::value && std::is_unsigned<T2>::value)
+    res.spvm = __spirv_JointMatrixSUMadINTEL(mA.spvm, mB.spvm, mC.spvm);
+  else if constexpr (std::is_unsigned<T1>::value && std::is_signed<T2>::value)
+    res.spvm = __spirv_JointMatrixUSMadINTEL(mA.spvm, mB.spvm, mC.spvm);
+  else
+    res.spvm = __spirv_JointMatrixMadINTEL(mA.spvm, mB.spvm, mC.spvm);
   return res;
 #else
   (void)sg;

sycl/test/matrix/matrix-bf16-test-SG-16.cpp

@@ -1,4 +1,4 @@
-// RUN: %clangxx -march=sapphirerapids -fsycl -O2 %s -o %t.out
+// RUN: %clangxx -fsycl -O2 %s -o %t.out
 #include <CL/sycl.hpp>
 #if (SYCL_EXT_ONEAPI_MATRIX == 2)
 #include <iostream>

sycl/test/matrix/matrix-bf16-test.cpp

@@ -1,4 +1,4 @@
-// RUN: %clangxx -march=sapphirerapids -fsycl -O2 %s -o %t.out
+// RUN: %clangxx -fsycl -O2 %s -o %t.out
 #include <CL/sycl.hpp>
 #if (SYCL_EXT_ONEAPI_MATRIX == 2)
 #include <iostream>

sycl/test/matrix/matrix-int8-test-SG-16.cpp

@@ -1,4 +1,4 @@
-// RUN: %clangxx -march=sapphirerapids -fsycl -O2 %s -o %t.out
+// RUN: %clangxx -fsycl -O2 %s -o %t.out
 #include <CL/sycl.hpp>
 #if (SYCL_EXT_ONEAPI_MATRIX == 2)
 #include <iostream>

sycl/test/matrix/matrix-int8-test.cpp

@@ -1,4 +1,4 @@
-// RUN: %clangxx -march=sapphirerapids -fsycl -O2 %s -o %t.out
+// RUN: %clangxx -fsycl -O2 %s -o %t.out
 #include <CL/sycl.hpp>
 #if (SYCL_EXT_ONEAPI_MATRIX == 2)
 #include <iostream>