[SYCL][HIP] Fix AMD joint_matrix use::a row-major load and multi-sub-group lane indexing

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<T, Space, IsDecorated> 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<S, double>) {
     const auto thread_x = idx % N;
@@ -211,33 +214,52 @@ template <
 void load_multiplicand_hip(joint_matrix_hip<S, Use, M, N, Layout> &res,
                            multi_ptr<T, Space, IsDecorated> 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<S, double>) {
-    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<T, Space, IsDecorated> 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<T, double>) {
     const auto thread_x = idx % N;

sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_apply.hpp

@@ -60,7 +60,7 @@ void hip_matrix_apply() {
                 joint_matrix<sub_group, OutType, use::accumulator, M, N> sub_c;
                 joint_matrix<sub_group, InType, use::b, K, N, layout::row_major>
                     sub_b;
-                joint_matrix<sub_group, InType, use::a, M, K, layout::col_major>
+                joint_matrix<sub_group, InType, use::a, M, K, layout::row_major>
                     sub_a;
 
                 joint_matrix_load(

sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_copy.hpp

@@ -69,7 +69,7 @@ void hip_matrix_copy() {
                     sub_c_copy;
                 joint_matrix<sub_group, InType, use::b, K, N, layout::row_major>
                     sub_b, sub_b_copy;
-                joint_matrix<sub_group, InType, use::a, M, K, layout::col_major>
+                joint_matrix<sub_group, InType, use::a, M, K, layout::row_major>
                     sub_a, sub_a_copy;
 
                 joint_matrix_load(

sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_fill.hpp

@@ -61,7 +61,7 @@ void hip_matrix_fill() {
                     sub_c{};
                 joint_matrix<sub_group, InType, use::b, K, N, layout::row_major>
                     sub_b{};
-                joint_matrix<sub_group, InType, use::a, M, K, layout::col_major>
+                joint_matrix<sub_group, InType, use::a, M, K, layout::row_major>
                     sub_a{};
 
                 joint_matrix_fill(sg, sub_a, 1);

sycl/test-e2e/Matrix/Inputs/joint_matrix_hip_mfma.hpp

@@ -68,7 +68,7 @@ void hip_matrix_mfma() {
                 joint_matrix<sub_group, OutType, use::accumulator, M, N> sub_c;
                 joint_matrix<sub_group, InType, use::b, K, N, layout::row_major>
                     sub_b;
-                joint_matrix<sub_group, InType, use::a, M, K, layout::col_major>
+                joint_matrix<sub_group, InType, use::a, M, K, layout::row_major>
                     sub_a;
 
                 joint_matrix_load(

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 <cassert>
+#include <cmath>
+#include <vector>
+
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/oneapi/matrix/matrix.hpp>
+
+using namespace sycl;
+using namespace sycl::ext::oneapi::experimental::matrix;
+using sycl::ext::oneapi::bfloat16;
+
+constexpr int M = 16, N = 16, K = 16;
+
+template <layout AL, layout BL> bool run_combo() {
+  std::vector<float> Alog(M * K), Blog(K * N), Ref(M * N, 0.f);
+  for (int i = 0; i < M * K; ++i)
+    Alog[i] = static_cast<float>((i * 3 + 1) % 7) - 3.f;
+  for (int i = 0; i < K * N; ++i)
+    Blog[i] = static_cast<float>((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<bfloat16> 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<float> C(M * N, 0.f);
+
+  {
+    queue q;
+    buffer<bfloat16> bA(Amem.data(), range{M * K});
+    buffer<bfloat16> bB(Bmem.data(), range{K * N});
+    buffer<float> 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<sub_group, float, use::accumulator, M, N> c;
+             joint_matrix<sub_group, bfloat16, use::a, M, K, AL> a;
+             joint_matrix<sub_group, bfloat16, use::b, K, N, BL> b;
+             joint_matrix_fill(sg, c, 0.f);
+             joint_matrix_load(
+                 sg, a, accA.template get_multi_ptr<access::decorated::yes>(),
+                 lda);
+             joint_matrix_load(
+                 sg, b, accB.template get_multi_ptr<access::decorated::yes>(),
+                 ldb);
+             joint_matrix_mad(sg, c, a, b, c);
+             joint_matrix_store(
+                 sg, c, accC.template get_multi_ptr<access::decorated::yes>(), 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<layout::row_major, layout::row_major>();
+  bool rc = run_combo<layout::row_major, layout::col_major>();
+  bool cr = run_combo<layout::col_major, layout::row_major>();
+  bool cc = run_combo<layout::col_major, layout::col_major>();
+  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;
+}

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 <cassert>
+#include <cmath>
+#include <vector>
+
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/oneapi/matrix/matrix.hpp>
+
+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<float> Alog(S * S), Blog(S * S), Ref(S * S, 0.f);
+  for (int i = 0; i < S * S; ++i)
+    Alog[i] = static_cast<float>((i * 3 + 1) % 7) - 3.f;
+  for (int i = 0; i < S * S; ++i)
+    Blog[i] = static_cast<float>((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<bfloat16> 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<float> C(S * S, 0.f);
+
+  {
+    queue q;
+    buffer<bfloat16> bA(Amem.data(), range{(size_t)S * S});
+    buffer<bfloat16> bB(Bmem.data(), range{(size_t)S * S});
+    buffer<float> 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<access::decorated::yes>();
+             auto pb = accB.template get_multi_ptr<access::decorated::yes>();
+             auto pc = accC.template get_multi_ptr<access::decorated::yes>();
+             joint_matrix<sub_group, float, use::accumulator, T, T> c;
+             joint_matrix<sub_group, bfloat16, use::a, T, T, layout::row_major> a;
+             joint_matrix<sub_group, bfloat16, use::b, T, T, layout::col_major> 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;
+}