From eed78c6dd93fd297765632861587d9c3be82e0fc Mon Sep 17 00:00:00 2001 From: big_yellow_duck <83417790+big-yellow-duck@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:54:34 +0000 Subject: [PATCH 1/2] [Dialect] Add FP8 support to GFX120X WMMA atom --- docs/api/dsl.rst | 2 +- docs/kernel_authoring_guide.md | 8 +++-- include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td | 2 +- lib/Bindings/Python/FlyROCDLExtension.cpp | 2 +- lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp | 15 ++++++--- tests/kernels/test_rdna4_wmma_atom.py | 3 +- tests/mlir/Conversion/wmma_gfx120x.mlir | 33 ++++++++++++++++++- tests/mlir/Conversion/wmma_gfx120x_neg.mlir | 10 +++--- 8 files changed, 58 insertions(+), 17 deletions(-) diff --git a/docs/api/dsl.rst b/docs/api/dsl.rst index 4de35e59bc..d27ec70a26 100644 --- a/docs/api/dsl.rst +++ b/docs/api/dsl.rst @@ -191,7 +191,7 @@ AMD-specific operations for ROCm: - **fx.rocdl.make_buffer_tensor(tensor)** -- create buffer resource from tensor (CDNA buffer copy) - **fx.rocdl.BufferCopy32b** / **BufferCopy128b** -- buffer copy instruction atoms - **fx.rocdl.MFMA(m, n, k, elem_ty_ab, elem_ty_acc=None)** -- MFMA instruction atom constructor (CDNA3/CDNA4; 4th arg is the A/B element type; accumulator defaults to f32) -- **fx.rocdl.WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, \*\*kwargs)** -- WMMA MMA atom constructor (arch-dispatched: gfx11 / gfx120x / gfx1250). gfx1250 supports f32(K4), f16/bf16(K32), fp8/bf8(K64/128), i8(K64), i4(K32); integer paths take ``sign_a`` / ``sign_b`` / ``clamp``. gfx120x (RDNA4) supports 16x16x16 f16/bf16 only, on the v8 operand ABI +- **fx.rocdl.WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, \*\*kwargs)** -- WMMA MMA atom constructor (arch-dispatched: gfx11 / gfx120x / gfx1250). gfx1250 supports f32(K4), f16/bf16(K32), fp8/bf8(K64/128), i8(K64), i4(K32); integer paths take ``sign_a`` / ``sign_b`` / ``clamp``. gfx120x (RDNA4) supports 16x16x16 f16/bf16/fp8(E4M3FN) to f32, on the v8 operand ABI - **fx.rocdl.WMMAScale(m, n, k, elem_ty_a, elem_ty_b=None, elem_ty_acc=None, \*, opsel_a=0, opsel_b=0, mod_c=0, reuse_a=False, reuse_b=False, block_size=32)** -- gfx1250 MX-scaled WMMA (E8M0 block scale, f8/f6/f4; ``16x16x128`` or ``32x16x128`` fp4-only). Per-operand scales are atom state (``scale_a`` / ``scale_b``) - **fx.rocdl.make_tdm_atom(tensor, tensor_extents, strides=None, \*, num_warps, ...)** -- build a gfx1250 TDM (Tensor Data Mover) async Global↔LDS whole-tile copy atom (rank 1-5); the global base comes from the ``copy_atom_call`` operand pointer, while the per-dim extent (OOB), stride, ``imm_offset``, and MCAST ``workgroup_mask`` are atom state. ``fx.rocdl.TDM(rank, num_warps, ...)`` builds the atom type only. Advance the K-loop tile with ``fx.copy(atom, gt, dst, imm_offset=...)`` - **fx.rocdl.sched_mfma(cnt)** -- insert MFMA scheduling barrier diff --git a/docs/kernel_authoring_guide.md b/docs/kernel_authoring_guide.md index 1fb1bdb8c6..d8a31c92e2 100644 --- a/docs/kernel_authoring_guide.md +++ b/docs/kernel_authoring_guide.md @@ -339,12 +339,14 @@ mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 32, T.i4, T.i32, sign_a=True, sign_b=T On RDNA4 (`gfx1200` / `gfx1201`) the same factory builds `MmaOpGFX120X_WMMAType` instead. RDNA4 shares the v8 register ABI but keeps the gfx11 instruction -shapes, so the only valid form is `16x16x16` f16/bf16 → f32 — the 16x16x32 and -fp8 K=64/128 shapes above are gfx1250-only and are rejected by the atom's -verifier. See `kernels/gemm/rdna_f16_gemm.py` for a full pipelined example. +shapes, so the valid forms are `16x16x16` f16/bf16/fp8(E4M3FN) → f32 — the +16x16x32 and fp8 K=64/128 shapes above are gfx1250-only and are rejected by the +atom's verifier. See `kernels/gemm/rdna_f16_gemm.py` for a full pipelined f16 +example and `kernels/gemm/rdna4_fp8_blockscale.py` for raw FP8 operands. ```python mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.BFloat16, fx.Float32)) # RDNA4 +mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.Float8E4M3FN, fx.Float32)) # RDNA4 FP8 ``` **MX-scaled WMMA** — `rocdl.WMMAScale(m, n, k, elem_ty_a, elem_ty_b=None, diff --git a/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td b/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td index 01f3e13c70..40f328c96b 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td +++ b/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td @@ -231,7 +231,7 @@ def FlyROCDL_MmaOpGFX120X_WMMA : FlyROCDL_MmaOp<"MmaOpGFX120X_WMMA", "gfx120x.wm "Type":$elemTyB, "Type":$elemTyAcc, // Integer-WMMA controls, forwarded to the ROCDL iu8/iu4 intrinsic. Ignored - // (and required to be false) on the fp16/bf16 paths. Always printed. + // (and required to be false) on the floating-point paths. Always printed. "bool":$signA, "bool":$signB, "bool":$clamp diff --git a/lib/Bindings/Python/FlyROCDLExtension.cpp b/lib/Bindings/Python/FlyROCDLExtension.cpp index 033087e3cf..754ea94d5e 100644 --- a/lib/Bindings/Python/FlyROCDLExtension.cpp +++ b/lib/Bindings/Python/FlyROCDLExtension.cpp @@ -144,7 +144,7 @@ struct PyMmaOpGFX120X_WMMAType : PyConcreteType { "sign_a"_a = false, "sign_b"_a = false, "clamp"_a = false, "context"_a = nb::none(), "Create a MmaOpGFX120X_WMMAType with m, n, k dimensions and element types " "(RDNA4 gfx1200 / gfx1201 wave32 WMMA, 16x16x16 with the v8 operand ABI). " - "sign_a/sign_b/clamp must be false: only the fp16/bf16 paths are supported."); + "sign_a/sign_b/clamp must be false: fp16/bf16/fp8(E4M3FN) paths are supported."); } }; diff --git a/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp b/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp index a0f808ba23..1e459338bd 100644 --- a/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp @@ -20,7 +20,7 @@ namespace mlir::fly_rocdl { // GFX120X (RDNA4: gfx1200, gfx1201) WMMA wave32. // // RDNA4 sits between the two existing WMMA atoms: -// * Instruction shapes are the gfx11 ones (16x16x16 fp16/bf16/iu8), and the +// * Instruction shapes are the gfx11 ones (16x16x16 fp16/bf16/fp8/iu8), and the // ROCDL intrinsics take the plain 3-operand {a, b, c} form — unlike // gfx1250, whose 16x16x32 ops carry mods/reuse operands. // * The register ABI is the gfx1250 "v8" one — each lane holds K/2 = 8 A/B @@ -72,18 +72,20 @@ LogicalResult MmaOpGFX120X_WMMAType::verify(function_ref e << " (the 16x16x32 shapes are gfx1250-only; use gfx1250.wmma there)"; } + const bool isFp8 = + isa(elemTyA) && isa(elemTyB) && elemTyAcc.isF32(); const bool isFp = (elemTyA.isF16() && elemTyB.isF16() && elemTyAcc.isF32()) || - (elemTyA.isBF16() && elemTyB.isBF16() && elemTyAcc.isF32()); + (elemTyA.isBF16() && elemTyB.isBF16() && elemTyAcc.isF32()) || isFp8; if (!isFp) { return emitError() << "unsupported GFX120X WMMA configuration: " << m << "x" << n << "x" << k << " with A=" << elemTyA << ", B=" << elemTyB << ", Acc=" << elemTyAcc; } - // The fp16/bf16 intrinsics have no sign/clamp operands; refuse to build an + // The floating-point intrinsics have no sign/clamp operands; refuse to build an // atom promising something codegen cannot deliver. if (signA || signB || clamp) { - return emitError() << "GFX120X WMMA fp16/bf16 path does not accept signA/signB/clamp " + return emitError() << "GFX120X WMMA floating-point path does not accept signA/signB/clamp " "(the ROCDL fp WMMA intrinsics have no such operands); got signA=" << signA << ", signB=" << signB << ", clamp=" << clamp; } @@ -99,7 +101,10 @@ LogicalResult MmaOpGFX120X_WMMAType::verify(function_ref e // per lane. // fp16 -> vector<8xf16> // bf16 -> vector<8xi16> (the bf16 WMMA intrinsic takes integer operands) +// fp8 -> vector<2xi32> (8 packed fp8 values) static Type getWmmaABType(MLIRContext *ctx, Type elemTy) { + if (isa(elemTy)) + return VectorType::get({2}, IntegerType::get(ctx, 32)); if (elemTy.isBF16()) return VectorType::get({8}, IntegerType::get(ctx, 16)); if (elemTy.isF16()) @@ -148,6 +153,8 @@ FailureOr MmaOpGFX120X_WMMAType::emitAtomCallSSA(OpBuilder &builder, Loca return ROCDL::wmma_f32_16x16x16_f16::create(builder, loc, rawAccTy, a, b, c).getResult(); if (elemTyA.isBF16() && elemTyB.isBF16()) return ROCDL::wmma_f32_16x16x16_bf16::create(builder, loc, rawAccTy, a, b, c).getResult(); + if (isa(elemTyA) && isa(elemTyB)) + return ROCDL::wmma_f32_16x16x16_fp8_fp8::create(builder, loc, rawAccTy, a, b, c).getResult(); return failure(); } diff --git a/tests/kernels/test_rdna4_wmma_atom.py b/tests/kernels/test_rdna4_wmma_atom.py index b3c930a51b..55e81ea1c7 100644 --- a/tests/kernels/test_rdna4_wmma_atom.py +++ b/tests/kernels/test_rdna4_wmma_atom.py @@ -90,8 +90,9 @@ def launch(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, stream: fx.Stream = fx.Stre [ (fx.BFloat16, torch.bfloat16), (fx.Float16, torch.float16), + (fx.Float8E4M3FN, torch.float8_e4m3fn), ], - ids=["bf16", "f16"], + ids=["bf16", "f16", "fp8"], ) def test_single_wmma_atom(elem_cls, torch_dtype): """A single gfx120x.wmma atom call must match A @ B.T exactly for integers. diff --git a/tests/mlir/Conversion/wmma_gfx120x.mlir b/tests/mlir/Conversion/wmma_gfx120x.mlir index c2fe88b5db..348b2c83af 100644 --- a/tests/mlir/Conversion/wmma_gfx120x.mlir +++ b/tests/mlir/Conversion/wmma_gfx120x.mlir @@ -3,7 +3,7 @@ // RUN: %fly-opt %s --fly-rewrite-func-signature --fly-canonicalize --fly-layout-lowering --convert-fly-to-rocdl | FileCheck %s // GFX120X (RDNA4: gfx1200 / gfx1201) WMMA wave32 atom lowering tests: -// fly.mma_atom_call -> rocdl.wmma.f32.16x16x16.{f16,bf16} intrinsic +// fly.mma_atom_call -> rocdl.wmma.f32.16x16x16.{f16,bf16,fp8_fp8} intrinsic // // RDNA4 keeps the gfx11 16x16x16 instruction shapes but uses the gfx1250 // "v8" register ABI, so the per-lane fragment shapes are half the gfx11 ones: @@ -13,6 +13,37 @@ // The bf16 intrinsic takes integer operands, so bf16 A/B lower to // vector<8xi16>. +// CHECK-LABEL: @test_gfx120x_wmma_atom_call_fp8 +// CHECK-SAME: (%[[D:.*]]: !llvm.ptr<5>, %[[A:.*]]: !llvm.ptr<5>, %[[B:.*]]: !llvm.ptr<5>, %[[C:.*]]: !llvm.ptr<5>) +func.func @test_gfx120x_wmma_atom_call_fp8( + %d: !fly.memref, + %a: !fly.memref, + %b: !fly.memref, + %c: !fly.memref) { + %atom = fly.make_mma_atom : !fly.mma_atom f32, signA = false, signB = false, clamp = false>> + // CHECK: %[[A_VAL:.*]] = llvm.load %[[A]] : !llvm.ptr<5> -> vector<2xi32> + // CHECK: %[[B_VAL:.*]] = llvm.load %[[B]] : !llvm.ptr<5> -> vector<2xi32> + // CHECK: %[[C_VAL:.*]] = llvm.load %[[C]] : !llvm.ptr<5> -> vector<8xf32> + // CHECK: %[[RES:.*]] = rocdl.wmma.f32.16x16x16.fp8_fp8 %[[A_VAL]], %[[B_VAL]], %[[C_VAL]] + // CHECK: llvm.store %[[RES]], %[[D]] : vector<8xf32>, !llvm.ptr<5> + fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, signA = false, signB = false, clamp = false>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () + return +} + +// CHECK-LABEL: @test_gfx120x_wmma_atom_call_ssa_fp8 +// CHECK-SAME: (%[[A:.*]]: vector<8xi8>, %[[B:.*]]: vector<8xi8>, %[[C:.*]]: vector<8xf32>) +func.func @test_gfx120x_wmma_atom_call_ssa_fp8( + %a: vector<8xf8E4M3FN>, + %b: vector<8xf8E4M3FN>, + %c: vector<8xf32>) -> vector<8xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32, signA = false, signB = false, clamp = false>> + // CHECK: %[[A_CAST:.*]] = llvm.bitcast %[[A]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[B_CAST:.*]] = llvm.bitcast %[[B]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[RES:.*]] = rocdl.wmma.f32.16x16x16.fp8_fp8 %[[A_CAST]], %[[B_CAST]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : (!fly.mma_atom f32, signA = false, signB = false, clamp = false>>, vector<8xf8E4M3FN>, vector<8xf8E4M3FN>, vector<8xf32>) -> vector<8xf32> + return %res : vector<8xf32> +} + // CHECK-LABEL: @test_gfx120x_wmma_atom_call_bf16 // CHECK-SAME: (%[[D:.*]]: !llvm.ptr<5>, %[[A:.*]]: !llvm.ptr<5>, %[[B:.*]]: !llvm.ptr<5>, %[[C:.*]]: !llvm.ptr<5>) func.func @test_gfx120x_wmma_atom_call_bf16( diff --git a/tests/mlir/Conversion/wmma_gfx120x_neg.mlir b/tests/mlir/Conversion/wmma_gfx120x_neg.mlir index 0c28b040e7..27210f1174 100644 --- a/tests/mlir/Conversion/wmma_gfx120x_neg.mlir +++ b/tests/mlir/Conversion/wmma_gfx120x_neg.mlir @@ -17,10 +17,10 @@ func.func @bad_shape_k32( // ----- -// Only the fp16/bf16 -> f32 forms are wired up so far. -// CHECK: unsupported GFX120X WMMA configuration: 16x16x16 with A='f8E4M3FN' -func.func @bad_elem_ty_fp8( - %a: !fly.mma_atom f32, signA = false, signB = false, clamp = false>>) { +// RDNA4 exposes OCP E4M3 FP8 but not the E5M2/BF8 variant. +// CHECK: unsupported GFX120X WMMA configuration: 16x16x16 with A='f8E5M2' +func.func @bad_elem_ty_bf8( + %a: !fly.mma_atom f32, signA = false, signB = false, clamp = false>>) { return } @@ -28,7 +28,7 @@ func.func @bad_elem_ty_fp8( // The fp WMMA intrinsics take no sign/clamp operands, so an atom that promises // them would silently drop the request at codegen time. -// CHECK: GFX120X WMMA fp16/bf16 path does not accept signA/signB/clamp +// CHECK: GFX120X WMMA floating-point path does not accept signA/signB/clamp func.func @bad_sign_clamp( %a: !fly.mma_atom f32, signA = true, signB = true, clamp = true>>) { return From 7c4d805c2385260151a363d1530755e0b501ad1a Mon Sep 17 00:00:00 2001 From: big_yellow_duck <83417790+big-yellow-duck@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:02:48 +0000 Subject: [PATCH 2/2] [Dialect] Support BF8 combinations in GFX120X WMMA Signed-off-by: big_yellow_duck <83417790+big-yellow-duck@users.noreply.github.com> --- docs/api/dsl.rst | 2 +- docs/kernel_authoring_guide.md | 14 ++++-- include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td | 10 ++-- lib/Bindings/Python/FlyROCDLExtension.cpp | 3 +- lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp | 22 ++++++--- python/flydsl/expr/rocdl/universal.py | 19 ++++--- tests/kernels/test_rdna4_wmma_atom.py | 36 ++++++++------ tests/mlir/Conversion/wmma_gfx120x.mlir | 49 +++++++++++++++++-- tests/mlir/Conversion/wmma_gfx120x_neg.mlir | 15 +++--- 9 files changed, 117 insertions(+), 53 deletions(-) diff --git a/docs/api/dsl.rst b/docs/api/dsl.rst index 52c484036a..cc865ca102 100644 --- a/docs/api/dsl.rst +++ b/docs/api/dsl.rst @@ -192,7 +192,7 @@ AMD-specific operations for ROCm: - **fx.rocdl.make_buffer_tensor(tensor)** -- create buffer resource from tensor (CDNA buffer copy) - **fx.rocdl.BufferCopy32b** / **BufferCopy128b** -- buffer copy instruction atoms - **fx.rocdl.MFMA(m, n, k, elem_ty_ab, elem_ty_acc=None)** -- MFMA instruction atom constructor (CDNA3/CDNA4; 4th arg is the A/B element type; accumulator defaults to f32) -- **fx.rocdl.WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, \*\*kwargs)** -- WMMA MMA atom constructor (arch-dispatched: gfx11 / gfx120x / gfx1250). gfx1250 supports f32(K4), f16/bf16(K32), fp8/bf8(K64/128), i8(K64), i4(K32); integer paths take ``sign_a`` / ``sign_b`` / ``clamp``. gfx120x (RDNA4) supports 16x16x16 f16/bf16/fp8(E4M3FN) to f32, on the v8 operand ABI +- **fx.rocdl.WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, \*\*kwargs)** -- WMMA MMA atom constructor (arch-dispatched: gfx11 / gfx120x / gfx1250). ``elem_ty_b`` optionally selects a different B operand type. gfx1250 supports f32(K4), f16/bf16(K32), fp8/bf8(K64/128), i8(K64), i4(K32); integer paths take ``sign_a`` / ``sign_b`` / ``clamp``. gfx120x (RDNA4) supports 16x16x16 f16/bf16 and every fp8(E4M3FN)/bf8(E5M2) A/B combination to f32, on the v8 operand ABI - **fx.rocdl.WMMAScale(m, n, k, elem_ty_a, elem_ty_b=None, elem_ty_acc=None, \*, opsel_a=0, opsel_b=0, mod_c=0, reuse_a=False, reuse_b=False, block_size=32)** -- gfx1250 MX-scaled WMMA (E8M0 block scale, f8/f6/f4; ``16x16x128`` or ``32x16x128`` fp4-only). Per-operand scales are atom state (``scale_a`` / ``scale_b``) - **fx.rocdl.make_tdm_atom(tensor, tensor_extents, strides=None, \*, num_warps, ...)** -- build a gfx1250 TDM (Tensor Data Mover) async Global↔LDS whole-tile copy atom (rank 1-5); the global base comes from the ``copy_atom_call`` operand pointer, while the per-dim extent (OOB), stride, ``imm_offset``, and MCAST ``workgroup_mask`` are atom state. ``fx.rocdl.TDM(rank, num_warps, ...)`` builds the atom type only. Advance the K-loop tile with ``fx.copy(atom, gt, dst, imm_offset=...)`` - **fx.rocdl.sched_mfma(cnt)** -- insert MFMA scheduling barrier diff --git a/docs/kernel_authoring_guide.md b/docs/kernel_authoring_guide.md index 648023e4cd..4a042b02dc 100644 --- a/docs/kernel_authoring_guide.md +++ b/docs/kernel_authoring_guide.md @@ -338,15 +338,19 @@ mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 32, T.i4, T.i32, sign_a=True, sign_b=T ``` On RDNA4 (`gfx1200` / `gfx1201`) the same factory builds `MmaOpGFX120X_WMMAType` -instead. RDNA4 shares the v8 register ABI but keeps the gfx11 instruction -shapes, so the valid forms are `16x16x16` f16/bf16/fp8(E4M3FN) → f32 — the -16x16x32 and fp8 K=64/128 shapes above are gfx1250-only and are rejected by the -atom's verifier. See `kernels/gemm/rdna_f16_gemm.py` for a full pipelined f16 -example and `kernels/gemm/rdna4_fp8_blockscale.py` for raw FP8 operands. +instead. RDNA4 shares the v8 register ABI, while its supported floating-point +forms use `16x16x16`: f16/bf16 and every fp8(E4M3FN)/bf8(E5M2) A/B combination +accumulating to f32. Pass a different B type with the keyword-only `elem_ty_b`; +it defaults to the A type. The 16x16x32 BF16 and fp8 K=64/128 forms above are +gfx1250-only and are rejected by the atom's verifier. See +`kernels/gemm/rdna_f16_gemm.py` for a full pipelined f16 example and +`kernels/gemm/rdna4_fp8_blockscale.py` for raw FP8 operands. ```python mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.BFloat16, fx.Float32)) # RDNA4 mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.Float8E4M3FN, fx.Float32)) # RDNA4 FP8 +mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.Float8E4M3FN, fx.Float32, + elem_ty_b=fx.Float8E5M2)) # RDNA4 FP8 x BF8 ``` **MX-scaled WMMA** — `rocdl.WMMAScale(m, n, k, elem_ty_a, elem_ty_b=None, diff --git a/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td b/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td index 40f328c96b..e2db231cc1 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td +++ b/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td @@ -216,11 +216,11 @@ def FlyROCDL_MmaOpGFX11_WMMA : FlyROCDL_MmaOp<"MmaOpGFX11_WMMA", "gfx11.wmma", [ //===----------------------------------------------------------------------===// // // RDNA4 shares the v8-operand WMMA register ABI with gfx1250 (each lane holds -// K/2 A/B elements, accumulator rows blocked 8-per-lane-half) but keeps the -// gfx11 16x16x16 instruction shapes and the 3-operand ROCDL intrinsic form. -// gfx1250 is a different target: it has 16x16x32 and the mods/reuse operand -// form, so it keeps its own `gfx1250.wmma` atom. Hence `gfx120x` rather than -// `gfx12` in the name — this atom covers gfx1200 / gfx1201 only. +// K/2 A/B elements, accumulator rows blocked 8-per-lane-half) but its supported +// floating-point forms use 16x16x16 and the 3-operand ROCDL intrinsic form. +// gfx1250 uses different floating-point shapes and the mods/reuse operand form, +// so it keeps its own `gfx1250.wmma` atom. Hence `gfx120x` rather than `gfx12` +// in the name — this atom covers gfx1200 / gfx1201 only. def FlyROCDL_MmaOpGFX120X_WMMA : FlyROCDL_MmaOp<"MmaOpGFX120X_WMMA", "gfx120x.wmma", []> { let parameters = (ins diff --git a/lib/Bindings/Python/FlyROCDLExtension.cpp b/lib/Bindings/Python/FlyROCDLExtension.cpp index 754ea94d5e..5bd4d0a107 100644 --- a/lib/Bindings/Python/FlyROCDLExtension.cpp +++ b/lib/Bindings/Python/FlyROCDLExtension.cpp @@ -144,7 +144,8 @@ struct PyMmaOpGFX120X_WMMAType : PyConcreteType { "sign_a"_a = false, "sign_b"_a = false, "clamp"_a = false, "context"_a = nb::none(), "Create a MmaOpGFX120X_WMMAType with m, n, k dimensions and element types " "(RDNA4 gfx1200 / gfx1201 wave32 WMMA, 16x16x16 with the v8 operand ABI). " - "sign_a/sign_b/clamp must be false: fp16/bf16/fp8(E4M3FN) paths are supported."); + "sign_a/sign_b/clamp must be false: fp16, bf16, and every fp8(E4M3FN)/" + "bf8(E5M2) A/B combination are supported."); } }; diff --git a/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp b/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp index 1e459338bd..2e84a2a9a9 100644 --- a/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/GFX120X/MmaAtom.cpp @@ -20,7 +20,8 @@ namespace mlir::fly_rocdl { // GFX120X (RDNA4: gfx1200, gfx1201) WMMA wave32. // // RDNA4 sits between the two existing WMMA atoms: -// * Instruction shapes are the gfx11 ones (16x16x16 fp16/bf16/fp8/iu8), and the +// * Floating-point instruction shapes are the gfx11 ones (16x16x16 +// fp16/bf16/fp8/bf8), and the // ROCDL intrinsics take the plain 3-operand {a, b, c} form — unlike // gfx1250, whose 16x16x32 ops carry mods/reuse operands. // * The register ABI is the gfx1250 "v8" one — each lane holds K/2 = 8 A/B @@ -64,16 +65,17 @@ Attribute MmaOpGFX120X_WMMAType::getThrValLayoutC() const { return gfx1250::getThrValLayoutCD(getContext(), getElemTyAcc()); } +static bool isFp8OrBf8(Type elemTy) { return isa(elemTy); } + LogicalResult MmaOpGFX120X_WMMAType::verify(function_ref emitError, int32_t m, int32_t n, int32_t k, Type elemTyA, Type elemTyB, Type elemTyAcc, bool signA, bool signB, bool clamp) { if (m != 16 || n != 16 || k != 16) { - return emitError() << "GFX120X WMMA requires M=N=K=16, got " << m << "x" << n << "x" << k - << " (the 16x16x32 shapes are gfx1250-only; use gfx1250.wmma there)"; + return emitError() << "GFX120X WMMA floating-point forms require M=N=K=16, got " << m << "x" + << n << "x" << k; } - const bool isFp8 = - isa(elemTyA) && isa(elemTyB) && elemTyAcc.isF32(); + const bool isFp8 = isFp8OrBf8(elemTyA) && isFp8OrBf8(elemTyB) && elemTyAcc.isF32(); const bool isFp = (elemTyA.isF16() && elemTyB.isF16() && elemTyAcc.isF32()) || (elemTyA.isBF16() && elemTyB.isBF16() && elemTyAcc.isF32()) || isFp8; @@ -101,9 +103,9 @@ LogicalResult MmaOpGFX120X_WMMAType::verify(function_ref e // per lane. // fp16 -> vector<8xf16> // bf16 -> vector<8xi16> (the bf16 WMMA intrinsic takes integer operands) -// fp8 -> vector<2xi32> (8 packed fp8 values) +// fp8/bf8 -> vector<2xi32> (8 packed 8-bit values) static Type getWmmaABType(MLIRContext *ctx, Type elemTy) { - if (isa(elemTy)) + if (isFp8OrBf8(elemTy)) return VectorType::get({2}, IntegerType::get(ctx, 32)); if (elemTy.isBF16()) return VectorType::get({8}, IntegerType::get(ctx, 16)); @@ -155,6 +157,12 @@ FailureOr MmaOpGFX120X_WMMAType::emitAtomCallSSA(OpBuilder &builder, Loca return ROCDL::wmma_f32_16x16x16_bf16::create(builder, loc, rawAccTy, a, b, c).getResult(); if (isa(elemTyA) && isa(elemTyB)) return ROCDL::wmma_f32_16x16x16_fp8_fp8::create(builder, loc, rawAccTy, a, b, c).getResult(); + if (isa(elemTyA) && isa(elemTyB)) + return ROCDL::wmma_f32_16x16x16_fp8_bf8::create(builder, loc, rawAccTy, a, b, c).getResult(); + if (isa(elemTyA) && isa(elemTyB)) + return ROCDL::wmma_f32_16x16x16_bf8_fp8::create(builder, loc, rawAccTy, a, b, c).getResult(); + if (isa(elemTyA) && isa(elemTyB)) + return ROCDL::wmma_f32_16x16x16_bf8_bf8::create(builder, loc, rawAccTy, a, b, c).getResult(); return failure(); } diff --git a/python/flydsl/expr/rocdl/universal.py b/python/flydsl/expr/rocdl/universal.py index 9147efafd0..d7190d9b10 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/python/flydsl/expr/rocdl/universal.py @@ -165,7 +165,10 @@ def MFMA(m, n, k, elem_ty_ab, elem_ty_acc=None): def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): """Create an arch-appropriate WMMA atom. - Supported kwargs (integer paths only — iu8 / iu4): + Supported kwargs: + elem_ty_b: optional B operand type for mixed-type instructions; defaults + to ``elem_ty_ab``. RDNA4 accepts every FP8(E4M3FN)/BF8(E5M2) + combination. sign_a (bool, default False): treat A operand as signed. sign_b (bool, default False): treat B operand as signed. clamp (bool, default False): saturate integer accumulator. @@ -175,7 +178,9 @@ def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): intrinsic has no such operands. Future WMMA ops for new architectures should extend kwargs here rather than growing the positional signature. """ - ty_ab = elem_ty_ab.ir_type if hasattr(elem_ty_ab, "ir_type") else elem_ty_ab + ty_a = elem_ty_ab.ir_type if hasattr(elem_ty_ab, "ir_type") else elem_ty_ab + elem_ty_b = kwargs.pop("elem_ty_b", None) + ty_b = ty_a if elem_ty_b is None else (elem_ty_b.ir_type if hasattr(elem_ty_b, "ir_type") else elem_ty_b) if elem_ty_acc is None: ty_acc = ir.F32Type.get() else: @@ -191,14 +196,14 @@ def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): arch = get_rocm_arch() or "" if arch.startswith("gfx11"): - return MmaOpGFX11_WMMAType.get(m, n, k, ty_ab, ty_ab, ty_acc, **kwargs) + return MmaOpGFX11_WMMAType.get(m, n, k, ty_a, ty_b, ty_acc, **kwargs) if arch.startswith("gfx1250"): return MmaOpGFX1250_WMMAType.get( m, n, k, - ty_ab, - ty_ab, + ty_a, + ty_b, ty_acc, sign_a=bool(kwargs.get("sign_a", False)), sign_b=bool(kwargs.get("sign_b", False)), @@ -209,8 +214,8 @@ def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): m, n, k, - ty_ab, - ty_ab, + ty_a, + ty_b, ty_acc, sign_a=bool(kwargs.get("sign_a", False)), sign_b=bool(kwargs.get("sign_b", False)), diff --git a/tests/kernels/test_rdna4_wmma_atom.py b/tests/kernels/test_rdna4_wmma_atom.py index 55e81ea1c7..3eed96679e 100644 --- a/tests/kernels/test_rdna4_wmma_atom.py +++ b/tests/kernels/test_rdna4_wmma_atom.py @@ -45,7 +45,7 @@ M = N = K = 16 -def _compile_single_wmma(elem_cls): +def _compile_single_wmma(elem_cls_a, elem_cls_b): """One wave, one atom: C[16,16] = A[16,16] @ B[16,16].T.""" f32 = fx.Float32 @@ -57,7 +57,7 @@ def wmma_kernel(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor): bB = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(B)), fx.make_layout((N, K), (K, 1))) bC = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(C)), fx.make_layout((M, N), (N, 1))) - mma_atom = fx.make_mma_atom(fx.rocdl.WMMA(M, N, K, elem_cls, f32)) + mma_atom = fx.make_mma_atom(fx.rocdl.WMMA(M, N, K, elem_cls_a, f32, elem_ty_b=elem_cls_b)) tiled_mma = fx.make_tiled_mma(mma_atom, fx.make_layout((1, 1, 1), (0, 0, 0))) thr_mma = tiled_mma.thr_slice(tid) @@ -65,14 +65,15 @@ def wmma_kernel(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor): frag_B = thr_mma.make_fragment_B(bB) frag_C = thr_mma.make_fragment_C(bC) - copy_ab = fx.make_copy_atom(fx.rocdl.BufferCopy(elem_cls.width), elem_cls) + copy_a = fx.make_copy_atom(fx.rocdl.BufferCopy(elem_cls_a.width), elem_cls_a) + copy_b = fx.make_copy_atom(fx.rocdl.BufferCopy(elem_cls_b.width), elem_cls_b) copy_c = fx.make_copy_atom(fx.rocdl.BufferCopy(f32.width), f32) - thr_copy_A = fx.make_tiled_copy_A(copy_ab, tiled_mma).get_slice(tid) - thr_copy_B = fx.make_tiled_copy_B(copy_ab, tiled_mma).get_slice(tid) + thr_copy_A = fx.make_tiled_copy_A(copy_a, tiled_mma).get_slice(tid) + thr_copy_B = fx.make_tiled_copy_B(copy_b, tiled_mma).get_slice(tid) thr_copy_C = fx.make_tiled_copy_C(copy_c, tiled_mma).get_slice(tid) - fx.copy(copy_ab, thr_copy_A.partition_S(bA), thr_copy_A.retile(frag_A)) - fx.copy(copy_ab, thr_copy_B.partition_S(bB), thr_copy_B.retile(frag_B)) + fx.copy(copy_a, thr_copy_A.partition_S(bA), thr_copy_A.retile(frag_A)) + fx.copy(copy_b, thr_copy_B.partition_S(bB), thr_copy_B.retile(frag_B)) frag_C.fill(0) fx.gemm(mma_atom, frag_C, frag_A, frag_B, frag_C) @@ -86,26 +87,29 @@ def launch(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, stream: fx.Stream = fx.Stre @pytest.mark.parametrize( - "elem_cls, torch_dtype", + "elem_cls_a, elem_cls_b, torch_dtype_a, torch_dtype_b", [ - (fx.BFloat16, torch.bfloat16), - (fx.Float16, torch.float16), - (fx.Float8E4M3FN, torch.float8_e4m3fn), + (fx.BFloat16, fx.BFloat16, torch.bfloat16, torch.bfloat16), + (fx.Float16, fx.Float16, torch.float16, torch.float16), + (fx.Float8E4M3FN, fx.Float8E4M3FN, torch.float8_e4m3fn, torch.float8_e4m3fn), + (fx.Float8E4M3FN, fx.Float8E5M2, torch.float8_e4m3fn, torch.float8_e5m2), + (fx.Float8E5M2, fx.Float8E4M3FN, torch.float8_e5m2, torch.float8_e4m3fn), + (fx.Float8E5M2, fx.Float8E5M2, torch.float8_e5m2, torch.float8_e5m2), ], - ids=["bf16", "f16", "fp8"], + ids=["bf16", "f16", "fp8_fp8", "fp8_bf8", "bf8_fp8", "bf8_bf8"], ) -def test_single_wmma_atom(elem_cls, torch_dtype): +def test_single_wmma_atom(elem_cls_a, elem_cls_b, torch_dtype_a, torch_dtype_b): """A single gfx120x.wmma atom call must match A @ B.T exactly for integers. Integer-valued inputs keep the result exactly representable, so any mismatch is a fragment-layout bug rather than rounding. """ torch.manual_seed(0) - a = (torch.randn(M, K, device="cuda") * 4).round().to(torch_dtype) - b = (torch.randn(N, K, device="cuda") * 4).round().to(torch_dtype) + a = (torch.randn(M, K, device="cuda") * 4).round().to(torch_dtype_a) + b = (torch.randn(N, K, device="cuda") * 4).round().to(torch_dtype_b) c = torch.zeros(M, N, dtype=torch.float32, device="cuda") - launch = _compile_single_wmma(elem_cls) + launch = _compile_single_wmma(elem_cls_a, elem_cls_b) launch(a, b, c, stream=torch.cuda.current_stream()) torch.cuda.synchronize() diff --git a/tests/mlir/Conversion/wmma_gfx120x.mlir b/tests/mlir/Conversion/wmma_gfx120x.mlir index 348b2c83af..8c584a661e 100644 --- a/tests/mlir/Conversion/wmma_gfx120x.mlir +++ b/tests/mlir/Conversion/wmma_gfx120x.mlir @@ -3,10 +3,11 @@ // RUN: %fly-opt %s --fly-rewrite-func-signature --fly-canonicalize --fly-layout-lowering --convert-fly-to-rocdl | FileCheck %s // GFX120X (RDNA4: gfx1200 / gfx1201) WMMA wave32 atom lowering tests: -// fly.mma_atom_call -> rocdl.wmma.f32.16x16x16.{f16,bf16,fp8_fp8} intrinsic +// fly.mma_atom_call -> rocdl.wmma.f32.16x16x16 +// .{f16,bf16,fp8_fp8,fp8_bf8,bf8_fp8,bf8_bf8} intrinsic // -// RDNA4 keeps the gfx11 16x16x16 instruction shapes but uses the gfx1250 -// "v8" register ABI, so the per-lane fragment shapes are half the gfx11 ones: +// These RDNA4 floating-point forms use 16x16x16 and the gfx1250 "v8" +// register ABI, so the per-lane fragment shapes are half the gfx11 ones: // A, B : 8 elements (gfx11 has 16, broadcast across the two lane halves) // C, D : 8 f32 slots (vector<8xf32>, same as gfx11) // @@ -44,6 +45,48 @@ func.func @test_gfx120x_wmma_atom_call_ssa_fp8( return %res : vector<8xf32> } +// CHECK-LABEL: @test_gfx120x_wmma_atom_call_ssa_fp8_bf8 +// CHECK-SAME: (%[[A:.*]]: vector<8xi8>, %[[B:.*]]: vector<8xi8>, %[[C:.*]]: vector<8xf32>) +func.func @test_gfx120x_wmma_atom_call_ssa_fp8_bf8( + %a: vector<8xf8E4M3FN>, + %b: vector<8xf8E5M2>, + %c: vector<8xf32>) -> vector<8xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32, signA = false, signB = false, clamp = false>> + // CHECK: %[[A_CAST:.*]] = llvm.bitcast %[[A]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[B_CAST:.*]] = llvm.bitcast %[[B]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[RES:.*]] = rocdl.wmma.f32.16x16x16.fp8_bf8 %[[A_CAST]], %[[B_CAST]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : (!fly.mma_atom f32, signA = false, signB = false, clamp = false>>, vector<8xf8E4M3FN>, vector<8xf8E5M2>, vector<8xf32>) -> vector<8xf32> + return %res : vector<8xf32> +} + +// CHECK-LABEL: @test_gfx120x_wmma_atom_call_ssa_bf8_fp8 +// CHECK-SAME: (%[[A:.*]]: vector<8xi8>, %[[B:.*]]: vector<8xi8>, %[[C:.*]]: vector<8xf32>) +func.func @test_gfx120x_wmma_atom_call_ssa_bf8_fp8( + %a: vector<8xf8E5M2>, + %b: vector<8xf8E4M3FN>, + %c: vector<8xf32>) -> vector<8xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32, signA = false, signB = false, clamp = false>> + // CHECK: %[[A_CAST:.*]] = llvm.bitcast %[[A]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[B_CAST:.*]] = llvm.bitcast %[[B]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[RES:.*]] = rocdl.wmma.f32.16x16x16.bf8_fp8 %[[A_CAST]], %[[B_CAST]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : (!fly.mma_atom f32, signA = false, signB = false, clamp = false>>, vector<8xf8E5M2>, vector<8xf8E4M3FN>, vector<8xf32>) -> vector<8xf32> + return %res : vector<8xf32> +} + +// CHECK-LABEL: @test_gfx120x_wmma_atom_call_ssa_bf8_bf8 +// CHECK-SAME: (%[[A:.*]]: vector<8xi8>, %[[B:.*]]: vector<8xi8>, %[[C:.*]]: vector<8xf32>) +func.func @test_gfx120x_wmma_atom_call_ssa_bf8_bf8( + %a: vector<8xf8E5M2>, + %b: vector<8xf8E5M2>, + %c: vector<8xf32>) -> vector<8xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32, signA = false, signB = false, clamp = false>> + // CHECK: %[[A_CAST:.*]] = llvm.bitcast %[[A]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[B_CAST:.*]] = llvm.bitcast %[[B]] : vector<8xi8> to vector<2xi32> + // CHECK: %[[RES:.*]] = rocdl.wmma.f32.16x16x16.bf8_bf8 %[[A_CAST]], %[[B_CAST]], %[[C]] + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : (!fly.mma_atom f32, signA = false, signB = false, clamp = false>>, vector<8xf8E5M2>, vector<8xf8E5M2>, vector<8xf32>) -> vector<8xf32> + return %res : vector<8xf32> +} + // CHECK-LABEL: @test_gfx120x_wmma_atom_call_bf16 // CHECK-SAME: (%[[D:.*]]: !llvm.ptr<5>, %[[A:.*]]: !llvm.ptr<5>, %[[B:.*]]: !llvm.ptr<5>, %[[C:.*]]: !llvm.ptr<5>) func.func @test_gfx120x_wmma_atom_call_bf16( diff --git a/tests/mlir/Conversion/wmma_gfx120x_neg.mlir b/tests/mlir/Conversion/wmma_gfx120x_neg.mlir index 27210f1174..9db2c53cd0 100644 --- a/tests/mlir/Conversion/wmma_gfx120x_neg.mlir +++ b/tests/mlir/Conversion/wmma_gfx120x_neg.mlir @@ -6,10 +6,9 @@ // ----- -// RDNA4 only has the 16x16x16 WMMA shapes; 16x16x32 is a gfx1250 instruction -// and must not be silently accepted here, or codegen would emit an instruction -// the target cannot select. -// CHECK: GFX120X WMMA requires M=N=K=16, got 16x16x32 +// The supported floating-point RDNA4 WMMA forms have K=16. In particular, +// 16x16x32 BF16 is a gfx1250 instruction and must not be accepted here. +// CHECK: GFX120X WMMA floating-point forms require M=N=K=16, got 16x16x32 func.func @bad_shape_k32( %a: !fly.mma_atom f32, signA = false, signB = false, clamp = false>>) { return @@ -17,10 +16,10 @@ func.func @bad_shape_k32( // ----- -// RDNA4 exposes OCP E4M3 FP8 but not the E5M2/BF8 variant. -// CHECK: unsupported GFX120X WMMA configuration: 16x16x16 with A='f8E5M2' -func.func @bad_elem_ty_bf8( - %a: !fly.mma_atom f32, signA = false, signB = false, clamp = false>>) { +// RDNA4 has no mixed F16/BF16 floating-point WMMA instruction. +// CHECK: unsupported GFX120X WMMA configuration: 16x16x16 with A='f16', B='bf16' +func.func @bad_elem_ty_mixed_16bit( + %a: !fly.mma_atom f32, signA = false, signB = false, clamp = false>>) { return }