Skip to content

[SYCLomatic][PTX][MMA] Support migration of PTX mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 #2941

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tomflinda merged 1 commit into from
Jul 23, 2025
Merged

[SYCLomatic][PTX][MMA] Support migration of PTX mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 #2941

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions clang/lib/DPCT/RulesAsm/AsmMigration.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1554,10 +1554,11 @@ class SYCLGen : public SYCLGenBase {
N = "8";
K = "16";

// Only f16/s8 types are supported for A and B matrices of m16n8k16
if (AType->getKind() == InlineAsmBuiltinType::f16) {
InMatrixType[0] = "uint32_t"; // A type is .f16x2
InMatrixType[1] = "uint32_t"; // B type is .f16x2
// Only f16/s8/bf16 types are supported for A and B matrices of m16n8k16
if (AType->getKind() == InlineAsmBuiltinType::f16 ||
AType->getKind() == InlineAsmBuiltinType::bf16) {
InMatrixType[0] = "uint32_t"; // A type is .f16/.bf16x2
InMatrixType[1] = "uint32_t"; // B type is .f16/.bf16x2

// If A matrix type is f16, then C&D matrix types can only be f32
if (CType->getKind() == InlineAsmBuiltinType::f32) {
Expand Down
2 changes: 1 addition & 1 deletion clang/runtime/dpct-rt/include/dpct/math.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2394,7 +2394,7 @@ template <typename T> struct MMAType {
/// - m8n8k4 (f32.f16.f16.f32)
/// - m8n8k16 (s32.s8.s8.s32)
/// - m16n8k8 (f32.f16.f16.f32 & f32.bf16.bf16.f32)
/// - m16n8k16 (f32.f16.f16.f32 & s32.s8.s8.s32)
/// - m16n8k16 (f32.f16.f16.f32 & f32.bf16.bf16.f32 & s32.s8.s8.s32)
/// - m16n8k32 (s32.s8.s8.s32)
/// Here, m, n & k define the shapes of A, B & C matrices respectively
/// (A = [m x k], B = [k x n], C = [m x n]).
Expand Down
19 changes: 18 additions & 1 deletion clang/test/dpct/asm/mma.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,8 @@ As per PTX ASM 8.1, below is the status of supported configurations
m8n8k4 .f16 .f16 .f32
m8n8k16 .s8 .s8 .s32
m16n8k8 .f16/.bf16 .f16/.bf16 .f32
m16n8k16 .f16 .f16 .f32
m16n8k16 .f16 .f16 .f32
.bf16 .bf16 .f32
.s8 .s8 .s32
m16n8k32 .s8 .s8 .s32

Expand Down Expand Up @@ -116,6 +117,22 @@ __global__ void mma_kernel_m16n8k16(int *a, int *b, int *c, float *fc, int *d) {
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
"r"(b[0]), "r"(b[1]));

// CHECK: {
// CHECK-NEXT: volatile void *d_mat_frag_ct1[4] = { &fc[0], &fc[1], &fc[2], &fc[3] };
// CHECK-NEXT: sycl::vec<uint32_t, 4> a_mat_frag_ct1(a[0], a[1], a[2], a[3]);
// CHECK-NEXT: sycl::vec<uint32_t, 2> b_mat_frag_ct1(b[0], b[1]);
// CHECK-NEXT: sycl::vec<float, 4> c_mat_frag_ct1(fc[0], fc[1], fc[2], fc[3]);
// CHECK-NEXT: dpct::experimental::matrix::mma<16, 8, 16, sycl::ext::oneapi::bfloat16, float>(reinterpret_cast<volatile void **>(d_mat_frag_ct1), &a_mat_frag_ct1, &b_mat_frag_ct1, &c_mat_frag_ct1);
// CHECK-NEXT: }
asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
" { %0, %1, %2, %3 }, "
" { %4, %5, %6, %7 }, "
" { %8, %9 }, "
" { %0, %1, %2, %3 };"
: "+f"(fc[0]), "+f"(fc[1]), "+f"(fc[2]), "+f"(fc[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
"r"(b[0]), "r"(b[1]));

// CHECK: {
// CHECK-NEXT: volatile void *d_mat_frag_ct1[4] = { &d[0], &d[1], &d[2], &d[3] };
// CHECK-NEXT: sycl::vec<uint32_t, 2> a_mat_frag_ct1(a[0], a[1]);
Expand Down
Loading