From 6eafc610ea5b49fa4aa81e3003e3b1fb7d969706 Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Sun, 19 Jul 2026 20:11:18 -0700 Subject: [PATCH] [SYCL][HIP] Fix AMD joint_matrix use::a row-major load and multi-sub-group lane indexing Two correctness fixes in the AMD (HIP) joint_matrix backend load path: 1. The use::a row-major load loaded the A multiplicand transposed. The access pattern was chosen from the layout enum alone, but "row-major" implies opposite in-memory stride patterns for A and B (the free dimension is strided for A, contiguous for B). Select the pattern from the combined (use, layout) condition so both layouts are correct for A. The double and non-double paths are unified (Size == 1 reduces to the former double case), which also fixes use::b col-major for double. 2. The per-lane fragment index was computed from the work-group linear id (get_group_linear_id() * sub_group_size + lane) instead of the sub-group local id. With a single sub-group per work-group this is accidentally correct, but kernels launching multiple sub-groups per work-group computed out-of-range element offsets, giving wrong results or memory faults. The hip Inputs helpers previously declared use::a as col_major while storing A row-major, relying on the old (buggy) behavior; they now use row_major. Adds two regression e2e tests: all four A/B layout combinations, and a tiled GEMM with multiple sub-groups per work-group. Co-authored-by: Cursor --- .../sycl/ext/oneapi/matrix/matrix-hip.hpp | 83 +++++++++----- .../Matrix/Inputs/joint_matrix_hip_apply.hpp | 2 +- .../Matrix/Inputs/joint_matrix_hip_copy.hpp | 2 +- .../Matrix/Inputs/joint_matrix_hip_fill.hpp | 2 +- .../Matrix/Inputs/joint_matrix_hip_mfma.hpp | 2 +- .../Matrix/joint_matrix_hip_all_layouts.cpp | 106 ++++++++++++++++++ .../joint_matrix_hip_multi_subgroup.cpp | 99 ++++++++++++++++ 7 files changed, 263 insertions(+), 33 deletions(-) create mode 100644 sycl/test-e2e/Matrix/joint_matrix_hip_all_layouts.cpp create mode 100644 sycl/test-e2e/Matrix/joint_matrix_hip_multi_subgroup.cpp diff --git a/sycl/include/sycl/ext/oneapi/matrix/matrix-hip.hpp b/sycl/include/sycl/ext/oneapi/matrix/matrix-hip.hpp index 6f55ae3f1d6cf..de70895bf3eb0 100644 --- a/sycl/include/sycl/ext/oneapi/matrix/matrix-hip.hpp +++ b/sycl/include/sycl/ext/oneapi/matrix/matrix-hip.hpp @@ -117,8 +117,11 @@ void load_accumulator_layoutT( S, sycl::ext::oneapi::experimental::matrix::use::accumulator, M, N, sycl::ext::oneapi::experimental::matrix::layout::dynamic> &res, multi_ptr src, size_t stride, Group &sg) { - const auto idx = sg.get_group_linear_id() * sg.get_local_range()[0] + - sg.get_local_linear_id(); + // The fragment lane index is the work-item's position within its own + // sub-group (0 .. sub_group_size-1). It must NOT include the sub-group's + // index within the work-group, otherwise kernels that launch more than one + // sub-group per work-group compute out-of-range element offsets. + const auto idx = sg.get_local_linear_id(); if constexpr (std::is_same_v) { const auto thread_x = idx % N; @@ -211,33 +214,52 @@ template < void load_multiplicand_hip(joint_matrix_hip &res, multi_ptr src, size_t stride, Group &sg) { - const auto idx = sg.get_group_linear_id() * sg.get_local_range()[0] + - sg.get_local_linear_id(); - - if constexpr (std::is_same_v) { - if constexpr (Layout == - sycl::ext::oneapi::experimental::matrix::layout::row_major) { - res.wi_marray[0] = src[idx]; - } else { - res.wi_marray[0] = src[(idx % M) * stride + idx / M]; + // The fragment lane index is the work-item's position within its own + // sub-group (0 .. sub_group_size-1). It must NOT include the sub-group's + // index within the work-group, otherwise kernels that launch more than one + // sub-group per work-group compute out-of-range element offsets. + const auto idx = sg.get_local_linear_id(); + + // The AMD MFMA fragment layout is described in terms of the matrix's "free" + // dimension (rows M for the A multiplicand, columns N for the B multiplicand) + // and the shared contraction dimension K. Whether the free dimension is the + // strided one in memory depends on BOTH the use and the layout: + // - A row_major: element (m,k) at m*stride + k -> free dim (m) is strided. + // - A col_major: element (m,k) at k*stride + m -> free dim (m) is contiguous. + // - B row_major: element (k,n) at k*stride + n -> free dim (n) is contiguous. + // - B col_major: element (k,n) at n*stride + k -> free dim (n) is strided. + // So the strided-free-dim access pattern applies to (A, row_major) and + // (B, col_major); the contiguous-free-dim pattern applies to (A, col_major) + // and (B, row_major). Selecting the pattern from this combined condition + // (rather than from the layout alone) makes both layouts behave correctly for + // the A multiplicand, which previously loaded A transposed for row_major. + constexpr bool StridedFreeDim = + (Use == sycl::ext::oneapi::experimental::matrix::use::a) + ? (Layout == + sycl::ext::oneapi::experimental::matrix::layout::row_major) + : (Layout == + sycl::ext::oneapi::experimental::matrix::layout::col_major); + + // Free dimension: M for the A multiplicand, N for the B multiplicand. The + // wavefront is split into (WAVEFRONT_SIZE / FreeDim) groups along the + // contraction dimension K, and each work-item holds one contiguous K-block of + // `Size` elements (Size == 1 for the double shapes). + constexpr int FreeDim = + (Use == sycl::ext::oneapi::experimental::matrix::use::a) ? M : N; + constexpr int Size = (M * N) / WAVEFRONT_SIZE; + + const auto thread_x = idx % FreeDim; + const auto thread_y = idx / FreeDim; + + if constexpr (StridedFreeDim) { + for (int i = 0; i < Size; ++i) { + const int c_idx = thread_x * stride + i + thread_y * Size; + res.wi_marray[i] = src[c_idx]; } } else { - constexpr int Dim = (M == 16) ? 16 : 32; - - const auto thread_x = idx % Dim; - const auto thread_y = idx / Dim; - - if constexpr (Layout == - sycl::ext::oneapi::experimental::matrix::layout::col_major) { - for (int i = 0; i < 4; ++i) { - const int c_idx = thread_x * stride + i + thread_y * 4; - res.wi_marray[i] = src[c_idx]; - } - } else { - for (int i = 0; i < 4; ++i) { - const int r_idx = thread_x + i * stride + thread_y * stride * 4; - res.wi_marray[i] = src[r_idx]; - } + for (int i = 0; i < Size; ++i) { + const int r_idx = thread_x + i * stride + thread_y * stride * Size; + res.wi_marray[i] = src[r_idx]; } } } @@ -251,8 +273,11 @@ void store_layoutT( T, sycl::ext::oneapi::experimental::matrix::use::accumulator, M, N, sycl::ext::oneapi::experimental::matrix::layout::dynamic> &src, multi_ptr dst, size_t stride, Group &sg) { - const auto idx = sg.get_group_linear_id() * sg.get_local_range()[0] + - sg.get_local_linear_id(); + // The fragment lane index is the work-item's position within its own + // sub-group (0 .. sub_group_size-1). It must NOT include the sub-group's + // index within the work-group, otherwise kernels that launch more than one + // sub-group per work-group compute out-of-range element offsets. + const auto idx = sg.get_local_linear_id(); if constexpr (std::is_same_v) { const auto thread_x = idx % N; diff --git a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_apply.hpp b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_apply.hpp index 01f187f2368a7..307964fc0c00e 100644 --- a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_apply.hpp +++ b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_apply.hpp @@ -60,7 +60,7 @@ void hip_matrix_apply() { joint_matrix sub_c; joint_matrix sub_b; - joint_matrix + joint_matrix sub_a; joint_matrix_load( diff --git a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_copy.hpp b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_copy.hpp index e56fecf8df5b3..a560d3400dd3b 100644 --- a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_copy.hpp +++ b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_copy.hpp @@ -69,7 +69,7 @@ void hip_matrix_copy() { sub_c_copy; joint_matrix sub_b, sub_b_copy; - joint_matrix + joint_matrix sub_a, sub_a_copy; joint_matrix_load( diff --git a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_fill.hpp b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_fill.hpp index f997836670ebf..bca0909819cc2 100644 --- a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_fill.hpp +++ b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_fill.hpp @@ -61,7 +61,7 @@ void hip_matrix_fill() { sub_c{}; joint_matrix sub_b{}; - joint_matrix + joint_matrix sub_a{}; joint_matrix_fill(sg, sub_a, 1); diff --git a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_mfma.hpp b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_mfma.hpp index d066350b638f0..9d46c271037e0 100644 --- a/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_mfma.hpp +++ b/sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_mfma.hpp @@ -68,7 +68,7 @@ void hip_matrix_mfma() { joint_matrix sub_c; joint_matrix sub_b; - joint_matrix + joint_matrix sub_a; joint_matrix_load( diff --git a/sycl/test-e2e/Matrix/joint_matrix_hip_all_layouts.cpp b/sycl/test-e2e/Matrix/joint_matrix_hip_all_layouts.cpp new file mode 100644 index 0000000000000..512d7de3f70df --- /dev/null +++ b/sycl/test-e2e/Matrix/joint_matrix_hip_all_layouts.cpp @@ -0,0 +1,106 @@ +//===---joint_matrix_hip_all_layouts.cpp - DPC++ joint_matrix-------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +// RUN: %clangxx -fsycl -fsycl-targets=amdgcn-amd-amdhsa %amd_arch_options %s -o %t.out +// RUN: %{run} %t.out + +// REQUIRES: target-amd + +// Regression test: a 16x16x16 bf16 -> float tile computed for every +// combination of A/B memory layouts (row/col major), each stored in memory +// according to its declared layout, and checked against a standard row-major +// C = A*B reference. Guards against the AMD `use::a` load-path layout bug where +// the row-major A multiplicand was loaded transposed. + +#include +#include +#include + +#include +#include + +using namespace sycl; +using namespace sycl::ext::oneapi::experimental::matrix; +using sycl::ext::oneapi::bfloat16; + +constexpr int M = 16, N = 16, K = 16; + +template bool run_combo() { + std::vector Alog(M * K), Blog(K * N), Ref(M * N, 0.f); + for (int i = 0; i < M * K; ++i) + Alog[i] = static_cast((i * 3 + 1) % 7) - 3.f; + for (int i = 0; i < K * N; ++i) + Blog[i] = static_cast((i * 2 + 1) % 5) - 2.f; + for (int m = 0; m < M; ++m) + for (int n = 0; n < N; ++n) { + float e = 0.f; + for (int k = 0; k < K; ++k) + e += Alog[m * K + k] * Blog[k * N + n]; + Ref[m * N + n] = e; + } + + const int lda = (AL == layout::row_major) ? K : M; + const int ldb = (BL == layout::row_major) ? N : K; + std::vector Amem(M * K), Bmem(K * N); + for (int m = 0; m < M; ++m) + for (int k = 0; k < K; ++k) + Amem[(AL == layout::row_major) ? m * lda + k : k * lda + m] = + bfloat16(Alog[m * K + k]); + for (int k = 0; k < K; ++k) + for (int n = 0; n < N; ++n) + Bmem[(BL == layout::row_major) ? k * ldb + n : n * ldb + k] = + bfloat16(Blog[k * N + n]); + std::vector C(M * N, 0.f); + + { + queue q; + buffer bA(Amem.data(), range{M * K}); + buffer bB(Bmem.data(), range{K * N}); + buffer bC(C.data(), range{M * N}); + q.submit([&](handler &h) { + accessor accA{bA, h, read_only}; + accessor accB{bB, h, read_only}; + accessor accC{bC, h, write_only}; + h.parallel_for( + nd_range<2>{{4, 16}, {4, 16}}, [=](nd_item<2> it) { + auto sg = it.get_sub_group(); + joint_matrix c; + joint_matrix a; + joint_matrix b; + joint_matrix_fill(sg, c, 0.f); + joint_matrix_load( + sg, a, accA.template get_multi_ptr(), + lda); + joint_matrix_load( + sg, b, accB.template get_multi_ptr(), + ldb); + joint_matrix_mad(sg, c, a, b, c); + joint_matrix_store( + sg, c, accC.template get_multi_ptr(), N, + layout::row_major); + }); + }).wait(); + } + + for (int i = 0; i < M * N; ++i) + if (std::fabs(C[i] - Ref[i]) > 2.f) + return false; + return true; +} + +int main() { + bool rr = run_combo(); + bool rc = run_combo(); + bool cr = run_combo(); + bool cc = run_combo(); + assert(rr && "A=row B=row"); + assert(rc && "A=row B=col"); + assert(cr && "A=col B=row"); + assert(cc && "A=col B=col"); + return 0; +} diff --git a/sycl/test-e2e/Matrix/joint_matrix_hip_multi_subgroup.cpp b/sycl/test-e2e/Matrix/joint_matrix_hip_multi_subgroup.cpp new file mode 100644 index 0000000000000..7d76dee85fa51 --- /dev/null +++ b/sycl/test-e2e/Matrix/joint_matrix_hip_multi_subgroup.cpp @@ -0,0 +1,99 @@ +//===---joint_matrix_hip_multi_subgroup.cpp - DPC++ joint_matrix----------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +// RUN: %clangxx -fsycl -fsycl-targets=amdgcn-amd-amdhsa %amd_arch_options %s -o %t.out +// RUN: %{run} %t.out + +// REQUIRES: target-amd + +// Regression test: a tiled GEMM whose work-group contains MORE THAN ONE +// sub-group (a 64x64 output tile computed by a 4x4 grid of wavefronts). Guards +// against the AMD fragment lane-index bug where the per-lane offset was derived +// from the work-group linear id instead of the sub-group local id, which made +// every sub-group beyond the first read/write out of range. + +#include +#include +#include + +#include +#include + +using namespace sycl; +using namespace sycl::ext::oneapi::experimental::matrix; +using sycl::ext::oneapi::bfloat16; + +constexpr int T = 16; // MFMA tile M=N=K +constexpr int TILE = 64; // work-group output tile (TILE x TILE) +constexpr int S = 128; // full square matrix dimension (multiple of TILE) + +int main() { + std::vector Alog(S * S), Blog(S * S), Ref(S * S, 0.f); + for (int i = 0; i < S * S; ++i) + Alog[i] = static_cast((i * 3 + 1) % 7) - 3.f; + for (int i = 0; i < S * S; ++i) + Blog[i] = static_cast((i * 2 + 1) % 5) - 2.f; + for (int m = 0; m < S; ++m) + for (int n = 0; n < S; ++n) { + float e = 0.f; + for (int k = 0; k < S; ++k) + e += Alog[m * S + k] * Blog[k * S + n]; + Ref[m * S + n] = e; + } + + std::vector Amem(S * S), Bmem(S * S); // A row-major, B col-major + for (int m = 0; m < S; ++m) + for (int k = 0; k < S; ++k) + Amem[m * S + k] = bfloat16(Alog[m * S + k]); + for (int k = 0; k < S; ++k) + for (int n = 0; n < S; ++n) + Bmem[n * S + k] = bfloat16(Blog[k * S + n]); + std::vector C(S * S, 0.f); + + { + queue q; + buffer bA(Amem.data(), range{(size_t)S * S}); + buffer bB(Bmem.data(), range{(size_t)S * S}); + buffer bC(C.data(), range{(size_t)S * S}); + // Work-group: {1, TILE/T, (TILE/T)*WAVE} where WAVE=64 -> 16 sub-groups. + const int nwarp = TILE / T; // 4 + range<2> lws{(size_t)nwarp, (size_t)nwarp * 64}; + range<2> gws{(size_t)(S / TILE) * nwarp, (size_t)(S / TILE) * nwarp * 64}; + q.submit([&](handler &h) { + accessor accA{bA, h, read_only}; + accessor accB{bB, h, read_only}; + accessor accC{bC, h, write_only}; + h.parallel_for( + nd_range<2>{gws, lws}, [=](nd_item<2> it) { + auto sg = it.get_sub_group(); + const int warpM = it.get_local_id(0); + const int warpN = it.get_local_id(1) / 64; + const int row = (it.get_group(0) * (TILE / T) + warpM) * T; + const int col = (it.get_group(1) * (TILE / T) + warpN) * T; + auto pa = accA.template get_multi_ptr(); + auto pb = accB.template get_multi_ptr(); + auto pc = accC.template get_multi_ptr(); + joint_matrix c; + joint_matrix a; + joint_matrix b; + joint_matrix_fill(sg, c, 0.f); + for (int kk = 0; kk < S; kk += T) { + joint_matrix_load(sg, a, pa + row * S + kk, S); + joint_matrix_load(sg, b, pb + col * S + kk, S); + joint_matrix_mad(sg, c, a, b, c); + } + joint_matrix_store(sg, c, pc + row * S + col, S, + layout::row_major); + }); + }).wait(); + } + + for (int i = 0; i < S * S; ++i) + assert(std::fabs(C[i] - Ref[i]) < 4.f && "multi-subgroup GEMM mismatch"); + return 0; +}