[Dialect] Add FP8 support to GFX120X WMMA atom (#1074)

* [Dialect] Add FP8 support to GFX120X WMMA atom

* [Dialect] Support BF8 combinations in GFX120X WMMA

Signed-off-by: big_yellow_duck <83417790+big-yellow-duck@users.noreply.github.com>

---------

Signed-off-by: big_yellow_duck <83417790+big-yellow-duck@users.noreply.github.com>

Author: big-yellow-duck

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 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). ``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

docs/kernel_authoring_guide.md

@@ -338,13 +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 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.
+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,

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
@@ -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

lib/Bindings/Python/FlyROCDLExtension.cpp

@@ -144,7 +144,8 @@ struct PyMmaOpGFX120X_WMMAType : PyConcreteType<PyMmaOpGFX120X_WMMAType> {
         "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, and every fp8(E4M3FN)/"
+        "bf8(E5M2) A/B combination are supported.");
   }
 };
 

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/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,26 +65,29 @@ Attribute MmaOpGFX120X_WMMAType::getThrValLayoutC() const {
   return gfx1250::getThrValLayoutCD(getContext(), getElemTyAcc());
 }
 
+static bool isFp8OrBf8(Type elemTy) { return isa<Float8E4M3FNType, Float8E5M2Type>(elemTy); }
+
 LogicalResult MmaOpGFX120X_WMMAType::verify(function_ref<InFlightDiagnostic()> 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 = isFp8OrBf8(elemTyA) && isFp8OrBf8(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 +103,10 @@ LogicalResult MmaOpGFX120X_WMMAType::verify(function_ref<InFlightDiagnostic()> e
 // per lane.
 //   fp16 -> vector<8xf16>
 //   bf16 -> vector<8xi16>   (the bf16 WMMA intrinsic takes integer operands)
+//   fp8/bf8 -> vector<2xi32>   (8 packed 8-bit values)
 static Type getWmmaABType(MLIRContext *ctx, Type elemTy) {
+  if (isFp8OrBf8(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 +155,14 @@ FailureOr<Value> 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<Float8E4M3FNType>(elemTyA) && isa<Float8E4M3FNType>(elemTyB))
+    return ROCDL::wmma_f32_16x16x16_fp8_fp8::create(builder, loc, rawAccTy, a, b, c).getResult();
+  if (isa<Float8E4M3FNType>(elemTyA) && isa<Float8E5M2Type>(elemTyB))
+    return ROCDL::wmma_f32_16x16x16_fp8_bf8::create(builder, loc, rawAccTy, a, b, c).getResult();
+  if (isa<Float8E5M2Type>(elemTyA) && isa<Float8E4M3FNType>(elemTyB))
+    return ROCDL::wmma_f32_16x16x16_bf8_fp8::create(builder, loc, rawAccTy, a, b, c).getResult();
+  if (isa<Float8E5M2Type>(elemTyA) && isa<Float8E5M2Type>(elemTyB))
+    return ROCDL::wmma_f32_16x16x16_bf8_bf8::create(builder, loc, rawAccTy, a, b, c).getResult();
 
   return failure();
 }

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)),

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,22 +57,23 @@ 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)
 
         frag_A = thr_mma.make_fragment_A(bA)
         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,25 +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.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"],
+    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()