[SYCL][Joint Matrix][E2E] Fix several tests (#16548)

with this fix joint_matrix_colA_rowB_colC.cpp is passing on CPU

Author: YuriPlyakhin

sycl/test-e2e/Matrix/SG32/joint_matrix_bfloat16_colmajorA_colmajorB.cpp

@@ -16,16 +16,7 @@
 
 // XFAIL: gpu
 // XFAIL-TRACKER: GSD-5768
-#include "../common.hpp"
-#include <iostream>
-#include <sycl/detail/core.hpp>
-#include <sycl/ext/oneapi/matrix/matrix.hpp>
-
-using namespace sycl;
-using namespace sycl::ext::oneapi::experimental::matrix;
-using bfloat16 = sycl::ext::oneapi::bfloat16;
 
+#include "../common.hpp"
 #define SG_SZ 32
-constexpr size_t TN = 16;
-
 #include "../joint_matrix_bfloat16_colmajorA_colmajorB_impl.hpp"

sycl/test-e2e/Matrix/SG32/joint_matrix_colA_rowB_colC.cpp

@@ -11,12 +11,9 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
-// XFAIL: run-mode
+// XFAIL: gpu && run-mode
 // XFAIL-TRACKER: GSD-5768
 
 #include "../common.hpp"
-
 #define SG_SZ 32
-constexpr size_t TN = 16;
-
 #include "../joint_matrix_colA_rowB_colC_impl.hpp"

sycl/test-e2e/Matrix/common.hpp

@@ -174,7 +174,7 @@ bool matrix_compare(unsigned int rows, unsigned int cols, T1 *src, T2 *ref) {
                                 std::is_same_v<T2, double>))) {
         float diff = std::fabs(src[i * cols + j] - (T1)ref[i * cols + j]);
         if (diff > FLOAT_EPSILON || std::isnan(src[i * cols + j])) {
-          std::cout << "Incorrect result in matrix. "
+          std::cerr << "Incorrect result in matrix. "
                     << "i: " << i << ", j: " << j
                     << ", Ref: " << (T1)ref[i * cols + j]
                     << ", Val: " << src[i * cols + j] << ", Diff: " << diff
@@ -183,14 +183,14 @@ bool matrix_compare(unsigned int rows, unsigned int cols, T1 *src, T2 *ref) {
         }
       } else if constexpr (exact || std::is_integral_v<T1>) {
         if (src[i * cols + j] != ref[i * cols + j]) {
-          std::cout << "Incorrect result in matrix."
+          std::cerr << "Incorrect result in matrix."
                     << "i: " << i << ", j: " << j
                     << ", Ref: " << ref[i * cols + j]
                     << ", Val: " << src[i * cols + j] << "\n";
           return false;
         }
       } else {
-        std::cout << "Unsupported type in matrix_compare\n";
+        std::cerr << "Unsupported type in matrix_compare\n";
         return false;
       }
     }

sycl/test-e2e/Matrix/joint_matrix_bf16_fill_k_cache_impl.hpp

@@ -119,7 +119,7 @@ double joint_matmul(TOperand *A, TOperand *B, TResult *C, queue &q, int i
           // along the workgroup prefetch for B matrix. For A matrix, sgId is
           // enough.
           size_t pm1B = sgId / 16;   // prefetch m1 (sgId/16)
-          size_t pn1B = sgId & 0x15; // prefetch n1 (sgId%16)
+          size_t pn1B = sgId & 0xF;  // prefetch n1 (sgId%16)
 #else                                // VNNI
           size_t pm1B = sgId / 8;   // prefetch m1 (sgId/8)
           size_t pn1B = sgId & 0x7; // prefetch n1 (sgId%8)

sycl/test-e2e/Matrix/joint_matrix_bfloat16_colmajorA_colmajorB.cpp

@@ -17,7 +17,4 @@
 // XFAIL-TRACKER: GSD-5768
 
 #include "common.hpp"
-
-constexpr size_t TN = 16;
-
 #include "joint_matrix_bfloat16_colmajorA_colmajorB_impl.hpp"

sycl/test-e2e/Matrix/joint_matrix_bfloat16_colmajorA_colmajorB_impl.hpp

@@ -6,8 +6,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-#define TM 8
-#define TK 16
+constexpr size_t TM = 8;
+constexpr size_t TN = 16;
+constexpr size_t TK = 16;
 
 template <typename T1, typename T2, size_t M, size_t N, size_t K>
 void matrix_multiply(big_matrix<T1, M, N> &C, big_matrix<T2, M, K> &A,
@@ -43,7 +44,6 @@ void matrix_multiply(big_matrix<T1, M, N> &C, big_matrix<T2, M, K> &A,
            sub_group sg = spmd_item.get_sub_group();
            joint_matrix<sub_group, bfloat16, use::a, TM, TK, layout::col_major>
                sub_a;
-           // For B, we assume B has been already VNNIed.
            joint_matrix<sub_group, bfloat16, use::b, TK, TN, layout::col_major>
                sub_b;
            joint_matrix<sub_group, float, use::accumulator, TM, TN> sub_c;

sycl/test-e2e/Matrix/joint_matrix_colA_rowB_colC.cpp

@@ -10,11 +10,8 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
-// XFAIL: run-mode
+// XFAIL: gpu && run-mode
 // XFAIL-TRACKER: GSD-5768
 
 #include "common.hpp"
-
-constexpr size_t TN = 16;
-
 #include "joint_matrix_colA_rowB_colC_impl.hpp"

sycl/test-e2e/Matrix/joint_matrix_colA_rowB_colC_impl.hpp

@@ -11,6 +11,7 @@
 #include <sycl/usm.hpp>
 
 constexpr size_t TM = 8;
+constexpr size_t TN = 16;
 constexpr size_t TK = 16;
 
 template <typename T1, typename T2, size_t NUM_ROWS_A, size_t NUM_COLS_A,
@@ -60,14 +61,15 @@ void matrix_multiply(T1 *C, T2 *A, T2 *B, queue q) {
            joint_matrix<sub_group, float, use::accumulator, TM, TN> sub_c;
            joint_matrix_fill(sg, sub_c, 1);
            for (int k = 0; k < K; k += TK) {
-             joint_matrix_load(sg, sub_a, pA + (sg_startx * TM) * K + k, K);
+             joint_matrix_load(sg, sub_a, pA + k * M + sg_startx * TM, M);
              joint_matrix_load(sg, sub_b, pB + k * N + sg_starty / sg_size * TN,
                                N);
              joint_matrix_mad(sg, sub_c, sub_a, sub_b, sub_c);
            }
-           joint_matrix_store(
-               sg, sub_c, pC + (sg_startx * TM) * N + sg_starty / sg_size * TN,
-               N, layout::col_major);
+           joint_matrix_store(sg, sub_c,
+                              pC + (sg_startx * TM) +
+                                  (sg_starty / sg_size * TN) * M,
+                              M, layout::col_major);
          }); // parallel for
    }).wait();
 }
@@ -76,23 +78,28 @@ int main() {
   static constexpr size_t MATRIX_M = 1024;
   static constexpr size_t MATRIX_N = 1024;
   static constexpr size_t MATRIX_K = 1024;
+
   queue q;
-  bfloat16 *A = malloc_shared<bfloat16>(MATRIX_M * MATRIX_K, q);
+  bfloat16 *A = malloc_shared<bfloat16>(MATRIX_K * MATRIX_M, q);
   bfloat16 *B = malloc_shared<bfloat16>(MATRIX_K * MATRIX_N, q);
-  float *C = malloc_shared<float>(MATRIX_M * MATRIX_N, q);
-  float *D = malloc_shared<float>(MATRIX_M * MATRIX_N, q);
+  float *C = malloc_shared<float>(MATRIX_N * MATRIX_M, q);
+  float *D = malloc_shared<float>(MATRIX_N * MATRIX_M, q);
 
-  matrix_rand(MATRIX_M, MATRIX_K, A, (bfloat16)5);
+  matrix_rand(MATRIX_K, MATRIX_M, A, (bfloat16)5);
   matrix_rand(MATRIX_K, MATRIX_N, B, (bfloat16)5);
-  matrix_fill(MATRIX_M, MATRIX_N, C, (float)1.0);
-  matrix_fill(MATRIX_M, MATRIX_N, D, (float)1.0);
+  matrix_fill(MATRIX_N, MATRIX_M, D, (float)1.0);
 
   matrix_multiply<float, bfloat16, MATRIX_M, MATRIX_K, MATRIX_K, MATRIX_N,
                   MATRIX_M, MATRIX_N>(C, A, B, q);
   matrix_multiply_ref(A, B, D, MATRIX_M, MATRIX_N, MATRIX_K,
-                      true /*transposed c*/);
+                      /*transposed c*/ true, /*colmajor a*/ true);
+
+  bool res = matrix_compare(MATRIX_N, MATRIX_M, C, D);
 
-  bool res = matrix_compare(MATRIX_M, MATRIX_N, C, D);
+  sycl::free(A, q);
+  sycl::free(B, q);
+  sycl::free(C, q);
+  sycl::free(D, q);
 
   std::cout << (res ? "passed" : "failed") << std::endl;
   return !res;