[Matrix] Enable joint_matrix_fill for joint_matrix feature (#4994)

Author: yubingex007-a11y

sycl/include/CL/__spirv/spirv_ops.hpp

@@ -86,6 +86,12 @@ __spirv_JointMatrixSUMadINTEL(
     __spv::__spirv_JointMatrixINTEL<T3, M, N, LC, S> *C,
     __spv::Scope::Flag Sc = __spv::Scope::Flag::Subgroup);
 
+template <typename T, std::size_t R, std::size_t C,
+          __spv::MatrixLayout L = __spv::MatrixLayout::RowMajor,
+          __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
+extern SYCL_EXTERNAL __spv::__spirv_JointMatrixINTEL<T, R, C, L, S> *
+__spirv_CompositeConstruct(const T v);
+
 template <typename T, std::size_t R, std::size_t C, __spv::MatrixLayout U,
           __spv::Scope::Flag S = __spv::Scope::Flag::Subgroup>
 extern SYCL_EXTERNAL size_t __spirv_JointMatrixWorkItemLengthINTEL(

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

@@ -202,6 +202,23 @@ joint_matrix_mad(Group sg, joint_matrix<T1, M, K, LayoutA, Group> &mA,
 #endif // __SYCL_DEVICE_ONLY__
 }
 
+template <typename Group, typename T, size_t NumRows, size_t NumCols,
+          matrix_layout Layout>
+inline __SYCL_ALWAYS_INLINE void
+joint_matrix_fill(Group sg,
+                  joint_matrix<T, NumRows, NumCols, Layout, Group> &res,
+                  const T v) {
+  // We kept the unused "sg" in joint_matrix_fill to match the other DPC++
+  // functions
+  (void)sg;
+#ifdef __SYCL_DEVICE_ONLY__
+  res.spvm = __spirv_CompositeConstruct<T, NumRows, NumCols>(v);
+#else
+  (void)res;
+  (void)v;
+#endif // __SYCL_DEVICE_ONLY__
+}
+
 template <typename T, size_t NumRows, size_t NumCols,
           matrix_layout Layout = matrix_layout::row_major,
           typename Group = sycl::sub_group>

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

@@ -1,4 +1,5 @@
 // RUN: %clangxx -fsycl -O2 %s -o %t.out
+// XFAIL: *
 #include <CL/sycl.hpp>
 #if (SYCL_EXT_ONEAPI_MATRIX == 2)
 #include <iostream>
@@ -68,10 +69,7 @@ void matrix_multiply(big_matrix<T1, NUM_ROWS_C, NUM_COLS_C> &C, big_matrix<T2, N
 
            // AMX: 8 register tiles : 1k byte size, SMmaxxSKmax =16x64
            // strideX = X's cols, so strideC = N, strideA = K, strideB = N*4
-           joint_matrix_load(sg, sub_c,
-                             accC.get_pointer() + (sg_startx * TM) * N +
-                                 sg_starty / SG_SZ * TN,
-                             N, matrix_layout::row_major);
+           joint_matrix_fill(sg, sub_c, 0);
            for (int k = 0; k < K / TK; k += 1) {
              joint_matrix_load(
                  sg, sub_a, accA.get_pointer() + (sg_startx * TM) * K + k * TK,
@@ -129,8 +127,8 @@ int main() {
   }
   for (int i = 0; i < MATRIX_M; i++) {
     for (int j = 0; j < MATRIX_N; j++) {
-      C[i][j] = 1;
-      D[i][j] = 1;
+      C[i][j] = 0;
+      D[i][j] = 0;
     }
   }