[Dialect][Kernel][Perf] Add gfx120x.wmma atom and port the RDNA4 GEMM to the layout API

CLAUDE.md

@@ -57,7 +57,7 @@ FlyDSL/
 │   └── mlir_flydsl/               # MLIR Python binding package source
 ├── include/flydsl/                # C++ TableGen headers for Fly / FlyROCDL dialects and passes
 ├── lib/                           # C++ dialect implementation, conversions, runtime wrappers, Python bindings
-│   └── Dialect/FlyROCDL/{CDNA3,CDNA4,GFX11,GFX1250}/  # Per-subtarget atom lowering: MmaAtom (MFMA on CDNA3/4, WMMA on GFX11/1250) + CopyAtom (Buffer/LDS, CDNA3/4 only)
+│   └── Dialect/FlyROCDL/{CDNA3,CDNA4,GFX11,GFX120X,GFX1250}/  # Per-subtarget atom lowering: MmaAtom (MFMA on CDNA3/4, WMMA on GFX11/120X/1250) + CopyAtom (Buffer/LDS, CDNA3/4 only; TDM on GFX1250)
 ├── tools/                         # fly-opt
 ├── kernels/                       # Production kernels, importable as kernels.*
 ├── tests/
@@ -169,7 +169,7 @@ helper code that is not part of the traced closure.
 | `gfx942` | MI300X / MI308X | 64 | MFMA | CDNA3 baseline; preshuffle GEMM, PA decode, CDNA BufferCopy |
 | `gfx950` / `gfx95*` | MI350 / MI355X | 64 | MFMA | CDNA4 path; FP4, MFMA scale, wider LDS copy paths, 160KB LDS |
 | `gfx11*` | RDNA3 / RDNA3.5 (Strix Halo, e.g. gfx1151) | 32 | WMMA | No MFMA; f16/bf16 (and i8/i4) WMMA GEMM; legacy v16-operand WMMA ABI; **no native FP8** (kernels fail-fast); `kernels/rdna3_f16_gemm.py`. `is_rdna_arch()` returns True. |
-| `gfx120*` | RDNA4 (gfx1201 = Radeon AI PRO R9700) | 32 | WMMA | RDNA path, wave32; new v8-operand WMMA ABI; native FP8. `is_rdna_arch()` returns True. |
+| `gfx120*` | RDNA4 (gfx1201 = Radeon AI PRO R9700) | 32 | WMMA | RDNA path, wave32; new v8-operand WMMA ABI with the gfx11 16x16x16 shapes (`gfx120x.wmma` atom, `lib/Dialect/FlyROCDL/GFX120X/`); native FP8. `is_rdna_arch()` returns True. |
 | `gfx1250` | — | 32 | WMMA / TDM | FP8/FP4 GEMM, MoE, async/TDM copy helpers, 320KB LDS. NOTE: `is_rdna_arch('gfx1250')` returns **False** and `get_warp_size` returns 64 — the gfx1250 kernels hardcode `WAVE_SIZE = 32` themselves. |
 
 Use `from flydsl.runtime.device import get_rocm_arch, is_rdna_arch` rather than

docs/api/compiler.rst

@@ -104,7 +104,7 @@ The ``flydsl.expr.rocdl`` module provides AMD-specific operations:
 - **fx.rocdl.make_buffer_tensor** -- create buffer resource descriptor from tensor (CDNA buffer copy)
 - **fx.rocdl.BufferCopy32b** / **BufferCopy128b** -- buffer copy atoms
 - **fx.rocdl.MFMA** -- MFMA instruction atoms (CDNA3/CDNA4; e.g., ``MFMA(16, 16, 4, fx.Float32)``)
-- **fx.rocdl.WMMA** / **fx.rocdl.WMMAScale** -- gfx1250 (wave32) WMMA and E8M0 MX-scaled WMMA MMA atoms
+- **fx.rocdl.WMMA** / **fx.rocdl.WMMAScale** -- wave32 WMMA MMA atoms; ``WMMA`` is arch-dispatched (gfx11 / gfx120x RDNA4 / gfx1250), ``WMMAScale`` is the gfx1250 E8M0 MX-scaled form
 - **fx.rocdl.make_tdm_atom** / **fx.rocdl.TDM** -- gfx1250 TDM async Global↔LDS whole-tile copy atom (1–5D; base from the copy operand, per-dim extent/stride/imm_offset/mask as atom state)
 
 fly-opt CLI

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 / gfx12 / gfx1250). gfx1250 supports f32(K4), f16/bf16(K32), fp8/bf8(K64/128), i8(K64), i4(K32); integer paths take ``sign_a`` / ``sign_b`` / ``clamp``
+- **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.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

@@ -333,6 +333,16 @@ mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 128, fx.Float8E4M3FN))    # fp8 → f3
 mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 32, T.i4, T.i32, sign_a=True, sign_b=True, clamp=True))
 ```
 
+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.
+
+```python
+mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.BFloat16, fx.Float32))  # RDNA4
+```
+
 **MX-scaled WMMA** — `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)` builds the E8M0 block-scaled WMMA (`V_WMMA_SCALE` /

include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td

@@ -211,4 +211,42 @@ def FlyROCDL_MmaOpGFX11_WMMA : FlyROCDL_MmaOp<"MmaOpGFX11_WMMA", "gfx11.wmma", [
   let genVerifyDecl = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// MmaOp GFX120X — WMMA wave32 (RDNA4: gfx1200, gfx1201)
+//===----------------------------------------------------------------------===//
+//
+// 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.
+
+def FlyROCDL_MmaOpGFX120X_WMMA : FlyROCDL_MmaOp<"MmaOpGFX120X_WMMA", "gfx120x.wmma", []> {
+  let parameters = (ins
+    "int32_t":$m,
+    "int32_t":$n,
+    "int32_t":$k,
+    "Type":$elemTyA,
+    "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.
+    "bool":$signA,
+    "bool":$signB,
+    "bool":$clamp
+  );
+  let assemblyFormat = "`<` custom<MNKDimensionList>($m, $n, $k) `,` `(` $elemTyA `,` $elemTyB `)` `->` $elemTyAcc `,` `signA` `=` $signA `,` `signB` `=` $signB `,` `clamp` `=` $clamp `>`";
+
+  let builders = [
+    TypeBuilderWithInferredContext<(ins "int32_t":$m, "int32_t":$n, "int32_t":$k, "Type":$elemTyA, "Type":$elemTyB, "Type":$elemTyAcc), [{
+      return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc, false, false, false);
+    }]>,
+    TypeBuilderWithInferredContext<(ins "int32_t":$m, "int32_t":$n, "int32_t":$k, "Type":$elemTyA, "Type":$elemTyB, "Type":$elemTyAcc, "bool":$signA, "bool":$signB, "bool":$clamp), [{
+      return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc, signA, signB, clamp);
+    }]>
+  ];
+  let genVerifyDecl = 1;
+}
+
 #endif // FLYROCDL_MMAATOM