joint_matrix performance issues and wrong results for AMD MI210

[fabilj825]

Describe the bug

I implemented the attached kernel for matrix multiplication. It works on a Nvidia A30 and returns correct results if I compare the results to a naive CPU implementation. The performance is for NVIDIA is up to 5 Tflops depending on the matrix size.

If I compile the kernel for the AMD MI210, the performance drops significantly to ≈ 0.5 Gflops. Furthermore, the results of the kernel are wrong if I compare them to the naive CPU implementation.

A second issue is that if I increase the BLOCK_SIZE to 32, the AMD kernel crashes throwing a core dump error whereas on Nvidia the kernel still executes fine and even faster as expected when increasing the block size.

My question is if someone experienced the same problems when using the joint_matrix extension on AMD, especially the wrong results or if it is a user error?

To reproduce

template <unsigned int matrix_Size, typename TOperand, typename TResult,
          size_t tileM, size_t tileN, size_t tileK, class kernel_name>
void compute_joint_fast( size_t sg_Size, queue &q, int testIterations) {

double t;
constexpr int TILE_SIZE = 16;
constexpr int BLOCK_SIZE = 16;
const int padded_N = ((matrix_Size % BLOCK_SIZE) == 0) ? matrix_Size : matrix_Size + (BLOCK_SIZE - (matrix_Size % BLOCK_SIZE));


TOperand *A = malloc_shared<TOperand>(padded_N * padded_N, q);
TOperand *B = malloc_shared<TOperand>(padded_N * padded_N, q);
TResult *C = malloc_shared<TResult>(padded_N * padded_N, q);
  // Initialize; fill matrices
fill_matrix_random(padded_N, matrix_Size, A);
fill_matrix_random(padded_N, matrix_Size, B);


#ifdef CHECKRESULTS
  TResult *refC = malloc_shared<TResult>(padded_N * padded_N, q);
  matrix_multiply_CPU(A, B, refC, padded_N);
#endif
range<2> global{padded_N / TILE_SIZE, (padded_N / TILE_SIZE) * sg_Size};
range<2> local{BLOCK_SIZE / TILE_SIZE, BLOCK_SIZE / TILE_SIZE * sg_Size};


  std::chrono::steady_clock::time_point start =
      std::chrono::steady_clock::now();
  for (unsigned int i = 0; i < testIterations; i++) {
    event mk = q.submit(
        [&](handler &h)

        {
          h.parallel_for<kernel_name>(nd_range<2>{global, local}, [=](nd_item<2>
                                                                          it)[[sycl::reqd_sub_group_size(64)]] {
            auto pA =
                address_space_cast<sycl::access::address_space::global_space,
                                   sycl::access::decorated::no>(A);

            auto pB =
                address_space_cast<sycl::access::address_space::global_space,
                                   sycl::access::decorated::no>(B);

            auto pC =
                address_space_cast<sycl::access::address_space::global_space,
                                   sycl::access::decorated::no>(C);

            auto m2 = it.get_group(0);
            auto n2 = it.get_group(1);
            auto m1 = it.get_local_id(0);
            auto n1 = it.get_local_id(1);
            auto sg = it.get_sub_group();

            int warpM = (n2 * it.get_local_range(1) + n1) / sg_Size;
            int warpN = (m2 * it.get_local_range(0) + m1);

            joint_matrix<sub_group, TResult, use::accumulator, tileM, tileN> tC;
            joint_matrix<sub_group, TOperand, use::a, tileM, tileK,
                         layout::row_major>
                tA;
            joint_matrix<sub_group, TOperand, use::b, tileK, tileN,
                         layout::row_major>
                tB;
            joint_matrix_fill(sg, tC, 0);

            for (unsigned int i = 0; i < padded_N; i += tileK) {
              int aCol = i;
              int aRow = warpM * tileM;
              int bCol = i;
              int bRow = warpN * tileN;

              joint_matrix_load(sg, tA, pA + aCol + aRow * padded_N,
                                padded_N);
              joint_matrix_load(sg, tB, pB + bRow + bCol * padded_N,
                                padded_N);
              joint_matrix_mad(sg, tC, tA, tB, tC);
            }
            int cCol = warpN * tileN;
            int cRow = warpM * tileM;
            joint_matrix_store(sg, tC, pC + cCol + cRow * padded_N,
                               padded_N, layout::row_major);
          });
        });
  }
  q.wait();

Compile:

icpx -O3 -fsycl -fsycl-targets=amd_gpu_gfx90a translateCuda.cpp 

Environment

• OS: Linux
• Target device and vendor: AMD MI210
• DPC++ version: Intel(R) oneAPI DPC++/C++ Compiler 2025.0.0 (2025.0.0.20241008)

Additional context

No response

[dkhaldi]

@mmoadeli, would you be able to take a look at this?
I will also try the code on Intel GPUs to check the correctness.

[dkhaldi]

@fabilj825, please provide a full compilable and runnable SYCL code. Thanks!

[fabilj825]

Of course, thanks for your efforts! It compiles with following line to also check the results.

icpx -O3 -fsycl -fsycl-targets=nvidia_gpu_sm_80 test.cpp -DCHECKRESULTS

helper.hpp

#include <random>
#include <iomanip> // For std::setw and formatting
#include <sycl/sycl.hpp>
#include <array>
#include <type_traits>
using namespace sycl;
using namespace sycl::ext::oneapi::experimental::matrix;
using bfloat16 = sycl::ext::oneapi::bfloat16;
constexpr float tol_float = 10e-2;


float bf16_to_fp32(bfloat16 val) {
  uint32_t fp32 = *((uint32_t *)&val) << 16;  
  return *reinterpret_cast<float*>(&fp32);  
}

template <typename T>
void fill_matrix_random(unsigned int padded_size, unsigned int matrix_size, T *src)
{
  std::random_device dev;
  std::uniform_real_distribution<float> fdistr(-1, 1);
  std::uniform_int_distribution idistr((int)-1, (int)1);

  for (unsigned int i = 0; i < padded_size; i++)
  {
    for (unsigned int j = 0; j < padded_size; j++)
    {
      if ((i < matrix_size && j < matrix_size))
      {
        if constexpr (std::is_integral_v<T>)
        {
          src[i * padded_size + j] = T(idistr(dev));
        }
        else
        {
          src[i * padded_size + j] = T(fdistr(dev));
        }
      }
      else
      {
        src[i * padded_size + j] = T(0);
      }
    }
  }
}



template <typename TOperand, typename TResult>
void matrix_multiply_CPU(TOperand *A, TOperand *B, TResult *C, int matrix_size ){
  for (unsigned int m = 0; m < matrix_size; m++) {
    for (unsigned int n = 0; n < matrix_size; n++) {
      TResult value = *(C + (m * matrix_size + n));
      for (unsigned int k = 0; k < matrix_size; k++) {

        TOperand *value_a = (TOperand *)(A + (m * matrix_size + k));
        TOperand *value_b = (TOperand *)(B + (k * matrix_size + n));


          if constexpr (std::is_same_v<TOperand, bfloat16> )
            value += bf16_to_fp32(value_a[0]) * bf16_to_fp32(value_b[0]);
          else if constexpr (std::is_same_v<TOperand, sycl::half>)
            value += static_cast<float>(value_a[0]) * static_cast<float>(value_b[0]);
          else if constexpr ((std::is_integral_v<TOperand> && std::is_integral_v<TResult>) || (std::is_same_v<TOperand, sycl::half>) ||
                             (std::is_same_v<TOperand, double> &&   std::is_same_v<TResult, double>))
            value += value_a[0] * value_b[0];

      }
      *(C + (m * matrix_size + n)) = value;
    }
  }
}

template <typename TResult>
bool matrix_compare (unsigned int matrix_size, TResult *src, TResult *ref) {
  for (int i = 0; i < matrix_size; i++) {
    for (int j = 0; j < matrix_size; j++) {
      if (std::is_same_v<TResult, float>  || std::is_same_v<TResult, double>) {
        float diff = std::fabs(src[i * matrix_size + j] - ref[i * matrix_size + j]);
        if (diff > tol_float) {
          std::cout << "Incorrect result in matrix. " << "i: " << i
                    << ", j: " << j << "\n";
          return false;
        }
      } else if constexpr (std::is_same_v<TResult, int32_t>) {
        if (src[i * matrix_size + j] != ref[i * matrix_size + j]) {
          std::cout << "Incorrect result in matrix." << "i: " << i
                    << ", j: " << j  << "\n";
          return false;
        }
      } 
    }
  }
  return true;
}

void print_performance_metrics (unsigned int matrix_size, size_t tileM, size_t tileN, size_t tileK, double duration, double testIterations) {
  double msecPerMatrixMul = duration / static_cast<double>(testIterations);
  double gflops = (2.f * matrix_size * matrix_size * matrix_size * 1.0e-9f) /
                  (msecPerMatrixMul / 1000.f);

  std::cerr << "DONE for GEMM size " << matrix_size << "x" << matrix_size << "x"
            << matrix_size << " Matrix Combination is " << tileM << "x" << tileN
            << "x" << tileK << std::endl;
  std::cerr << "Average test time is " << msecPerMatrixMul << " ms"
            << std::endl;

  std::cerr << "GOPS is " << gflops << " Gop/s" << std::endl;

  std::cerr << "_________________________________" << std::endl;
  std::cerr << "_________________________________" << std::endl;
  std::cerr << "                                 " << std::endl;
  std::cerr << "                                 " << std::endl;
}

test.cpp

#include "helper.hpp"

#ifndef MATRIX_SIZE
#define MATRIX_SIZE 2048
#endif


template <unsigned int matrix_Size, typename TOperand, typename TResult,
          size_t tileM, size_t tileN, size_t tileK, class kernel_name>
void compute_joint_fast( size_t sg_Size, queue &q, int testIterations) {
  double t;

#ifdef FP64_GEMM
constexpr int WARP_SIZE = 8;
constexpr int BLOCK_SIZE = 32;
#else
constexpr int WARP_SIZE = 16;
constexpr int BLOCK_SIZE = 16;
#endif

  const int padded_N =
      ((matrix_Size % BLOCK_SIZE) == 0)
          ? matrix_Size
          : matrix_Size + (BLOCK_SIZE - (matrix_Size % BLOCK_SIZE));

  TOperand *A = malloc_shared<TOperand>(padded_N * padded_N, q);
  TOperand *B = malloc_shared<TOperand>(padded_N * padded_N, q);
  TResult *C = malloc_shared<TResult>(padded_N * padded_N, q);
  // Initialize; fill matrices
  fill_matrix_random(padded_N, matrix_Size, A);
  fill_matrix_random(padded_N, matrix_Size, B);


#ifdef CHECKRESULTS
  TResult *refC = malloc_shared<TResult>(padded_N * padded_N, q);
  matrix_multiply_CPU(A, B, refC, padded_N);
#endif
  range<2> global{padded_N / WARP_SIZE, (padded_N / WARP_SIZE) * sg_Size};
  range<2> local{BLOCK_SIZE / WARP_SIZE, BLOCK_SIZE / WARP_SIZE * sg_Size};



  std::chrono::steady_clock::time_point start =
      std::chrono::steady_clock::now();
  for (unsigned int i = 0; i < testIterations; i++) {
    event mk = q.submit(
        [&](handler &h)

        {
          h.parallel_for<kernel_name>(nd_range<2>{global, local}, [=](nd_item<2>
                                                                          it)[[sycl::reqd_sub_group_size(32)]] {
            auto pA =
                address_space_cast<sycl::access::address_space::global_space,
                                   sycl::access::decorated::no>(A);

            auto pB =
                address_space_cast<sycl::access::address_space::global_space,
                                   sycl::access::decorated::no>(B);

            auto pC =
                address_space_cast<sycl::access::address_space::global_space,
                                   sycl::access::decorated::no>(C);

            auto m2 = it.get_group(0);
            auto n2 = it.get_group(1);
            auto m1 = it.get_local_id(0);
            auto n1 = it.get_local_id(1);
            auto sg = it.get_sub_group();

            int warpM = (n2 * it.get_local_range(1) + n1) / sg_Size;
            int warpN = (m2 * it.get_local_range(0) + m1);

            joint_matrix<sub_group, TResult, use::accumulator, tileM, tileN> tC;
            joint_matrix<sub_group, TOperand, use::a, tileM, tileK,
                         layout::row_major>
                tA;
            joint_matrix<sub_group, TOperand, use::b, tileK, tileN,
                         layout::row_major>
                tB;
            joint_matrix_fill(sg, tC, 0);

            for (unsigned int i = 0; i < padded_N; i += tileK) {
              int aCol = i;
              int aRow = warpM * tileM;
              int bCol = i;
              int bRow = warpN * tileN;

              joint_matrix_load(sg, tA, pA + aCol + aRow * padded_N,
                                padded_N);
              joint_matrix_load(sg, tB, pB + bRow + bCol * padded_N,
                                padded_N);
              joint_matrix_mad(sg, tC, tA, tB, tC);
            }
            int cCol = warpN * tileN;
            int cRow = warpM * tileM;
            joint_matrix_store(sg, tC, pC + cCol + cRow * padded_N,
                               padded_N, layout::row_major);
          });
        });
  }
  q.wait();

  std::chrono::duration<double, std::milli> duration =
      std::chrono::steady_clock::now() - start;

 print_performance_metrics(padded_N, tileM, tileN, tileK, duration.count(),
                            testIterations);
#ifdef CHECKRESULTS
  matrix_compare(padded_N, C, refC);
#endif

  free(A, q);
  free(B, q);
  free(C, q);
#ifdef CHECKRESULTS
  free(refC, q);
#endif
}



int main() {
  queue q;
  std::cout << "Device: " << q.get_device().get_info<info::device::name>()
            << "\n";
            auto sg_sizes = q.get_device().get_info<info::device::sub_group_sizes>();

  const size_t sg_Size =   sg_sizes.empty() ? 32
                       : *std::max_element(sg_sizes.begin(), sg_sizes.end());
  std::cout << "Kernel uses subgroup size: " << sg_Size << std::endl;

  size_t local_mem_size =
      q.get_device().get_info<sycl::info::device::local_mem_size>();
  std::cout << "GPU local memory size is: " << local_mem_size << std::endl;
  std::cout << "                                 " << std::endl;

  constexpr unsigned int testIterations = 10;

#ifdef FP64_GEMM
  std::cout << "Matrix data types A,B -> double C-> double" << std::endl;
  compute_joint_fast<MATRIX_SIZE, double, double, 8, 8, 4, class double_8x8x4>(
      sg_Size, q, testIterations);

#else
  std::cout << "Matrix data types A,B -> bfloat16 C-> float" << std::endl;
  compute_joint_fast<MATRIX_SIZE, bfloat16, float, 16, 16, 16,
                     class bfloat_16x16x16>(sg_Size, q, testIterations);

  std::cout << "Matrix data types A,B -> signed char C-> signed int"
            << std::endl;
  compute_joint_fast<MATRIX_SIZE, signed char, signed int, 16, 16, 16,
                     class schar_16x16x16>(sg_Size, q, testIterations);

  std::cout << "Matrix data types A,B -> half C-> float" << std::endl;
  compute_joint_fast<MATRIX_SIZE, half, float, 16, 16, 16, class half_16x16x16>(
      sg_Size, q, testIterations);

#endif

  return 0;
}

[dkhaldi]

Do all the GEMMs: double, fp16, bf16, int8 fail when block size=32 or just a subset of them?

[fabilj825]

I did experience it for bf16 as it is the first one called. Sometimes it still executed without crashing but still with bad performance and wrong results.

[dkhaldi]

OK, I run the bf16 code on Intel GPU and it passes for me (block size=32). I also looked at the logic. It seems correct to me.
I created an internal issue (OFNAAO-420) and assign it to the original contributor.
Regarding the performance, from one device to the other, tuning techniques are different. I would recommend looking for a HIP-based GEMM code that gives you good performance, then use the tuning used there.

[fabilj825]

@dkhaldi thanks for testing the code on Intel and checking on this!