add-target-atom-op

# Add a Target-Specific Mma / Copy Op to a FlyDSL Backend Dialect Step-by-step recipe for authoring a new `MmaOp*Type` or `CopyOp*Type` in a backend dialect (`fly_rocdl`, or a future `fly_ptx` / ...), plus the **inherent design contract** every Op author must understand before writing a single line of code. The examples throughout this skill draw from the `fly_rocdl` dialect (AMD ROCDL backend). The design is deliberately backend-agnostic: the generic `!fly.mma_atom` / `!fly.copy_atom` wrappers and the three type interfaces (`Fly_MmaOpTypeInterface`, `Fly_CopyOpTypeInterface`, `Fly_StatefulOpTypeInterface`) live in the target-neutral `fly` dialect and know nothing about AMD, NVIDIA, or other specifics. A new backend follows the exact same recipe — only the payload types, the final intrinsic emission, and the directory prefix change. --- ## 1. Inherent Design: How FlyDSL Atoms Work Internalize these five facts before adding any `Op`. They explain *why* the reference `CDNA3`, `CDNA4`, `GFX1250` implementations look the way they do — and the same structure applies verbatim to any new backend. ### 1.1 Two-level type design: generic wrapper + target-specific payload There are **two kinds** of related types, and they live in different dialects: | Level | Dialect | Type (example) | Role | |-------|---------|----------------|------| | generic wrapper | `fly` | `!fly.mma_atom<...>`, `!fly.copy_atom<..., bits>` | Target-agnostic. Appears everywhere in kernel IR. | | target payload | backend dialect (e.g. `fly_rocdl`) | `!fly_rocdl.cdna3.mfma<...>`, `!fly_rocdl.cdna3.buffer_copy<32>` | Knows *which* concrete instruction/intrinsic to emit. | The generic wrapper always holds a **payload** type as its first parameter. ```mlir // Using the ROCDL backend (the current reference): !fly.mma_atom<!fly_rocdl.cdna3.mfma<16x16x4, (f32, f32) -> f32>> !fly.copy_atom<!fly_rocdl.cdna3.buffer_copy<32>, 32> // A hypothetical NVIDIA backend would look like: !fly.mma_atom<!fly_ptx.sm90.wgmma<64x128x16, (f16, f16) -> f32>> !fly.copy_atom<!fly_ptx.sm80.cp_async<128>, 128> ``` Every method you see on `MmaAtomType` / `CopyAtomType` is a trampoline: ```cpp // lib/Dialect/Fly/IR/FlyTypeDefs.cpp Attribute MmaAtomType::getShapeMNK() const { return cast<MmaOpTypeInterface>(getMmaOp()).getShapeMNK(); } LogicalResult MmaAtomType::emitAtomCall(...) const { return cast<MmaOpTypeInterface>(getMmaOp()).emitAtomCall(...); } ``` **Your job** when adding a new Op is to define the *payload* type and implement the interface methods — the wrapper and the kernel-level ops (`fly.mma_atom_call`, `fly.copy_atom_call`, `fly.make_mma_atom`, ...) work automatically. ### 1.2 Three interfaces an Op type may implement `include/flydsl/Dialect/Fly/IR/FlyInterfaces.td`: | Interface | Required for... | Methods — **mandatory** / *optional* (see §1.3) | |-----------------------------------|-----------------|----------------------| | `Fly_MayStaticTypeInterface` | Stateless atoms (CopyOp with *no* mutable state; all MmaOps today) | **`isStatic`**, **`rebuildStaticValue`** | | `Fly_CopyOpTypeInterface` | All CopyOps | **`getThrLayout`**, **`getThrBitLayoutSrc/Dst/Ref`**, **`emitAtomCall`** (mem + pred), *`emitAtomCallSSA`* (mem + pred — only if `fly-convert-atom-call-to-ssa-form` is in the pipeline) | | `Fly_MmaOpTypeInterface` | All MmaOps | **`getThrLayout`**, **`getShapeMNK`**, **`getValTypeA/B/C/D`**, **`getThrValLayoutA/B/C`**, **`emitAtomCall`**, *`emitAtomCallSSA`* (only if SSA-promotion pass is active) | | `Fly_StatefulOpTypeInterface` | Atoms that carry mutable per-call state (e.g. `soffset`, `imm_offset`) | **`getConvertedType`**, **`getDefaultState`**, **`setAtomState`** | Backend dialect could provide four convenience base classes that pre-declare the right interface combinations. In the reference ROCDL backend (`include/flydsl/Dialect/FlyROCDL/IR/Dialect.td`) they are: ```tablegen class FlyROCDL_CopyOp // stateless CopyOp : MayStatic + CopyOp class FlyROCDL_StatefulCopyOp // stateful CopyOp : CopyOp + Stateful class FlyROCDL_MmaOp // stateless MmaOp : MayStatic + MmaOp class FlyROCDL_StatefulMmaOp // stateful MmaOp : MmaOp + Stateful ``` Mnemonic: **stateful => no `MayStaticTypeInterface`**; the mutable state *is* the dynamic component, so the type is never "fully static" in the canonical-rebuild sense. ### 1.3 `emitAtomCall` vs `emitAtomCallSSA` — only `emitAtomCall` is mandatory Two kernel-IR ops carry the atom invocation, and they correspond to the two interface methods: | Kernel Op | Operand form | Lowered via | Implementation status | |--------------------------|-----------------------------------------------|-----------------------|-----------------------| | `fly.copy_atom_call` | `src/dst : !fly.memref<...>` | `emitAtomCall` | **Required** | | `fly.mma_atom_call` | `a/b/c/d : !fly.memref<...>` | `emitAtomCall` | **Required** | | `fly.copy_atom_call_ssa` | `src/dst : SSA value or !fly.memref<..., addressSpace != Register>` | `emitAtomCallSSA` | **Optional** — only needed if `fly-convert-atom-call-to-ssa-form` appears in the pipeline | | `fly.mma_atom_call_ssa` | `a/b/c : SSA value or !fly.memref<..., addressSpace != Register>` | `emitAtomCallSSA` | **Optional** (same condition) | **Default path (memref / `emitAtomCall`).** Every `fly.copy_atom_call` / `fly.mma_atom_call` in the IR lowers through `emitAtomCall`. The Op receives the operand *pointers* into register memory (`!fly.memref<..., register, layout>`), is expected to issue `llvm.load` / `llvm.store` itself to read/write threads' registers, and emit the backend intrinsic in between. This is sufficient for the full compile-to-binary pipeline — no SSA version required. **Optional path (SSA / `emitAtomCallSSA`).** A pipeline may insert the `fly-convert-atom-call-to-ssa-form` pass (see `lib/Dialect/Fly/Transforms/ConvertAtomCallToSSAForm.cpp`). That pass inspects every `AtomCall` and, for operands whose `register`-address-space memref has a **coalescable** layout (`isEligibleToPromote`: stride-1 or shape-1 after coalesce), rewrites them: 1. `PtrLoadOp` pulls the whole register memref into a single SSA value of type `RegMem2SSAType(memref)` — which is `elemTy` when the layout has cosize 1, or `vector<cosize × elemTy>` otherwise (see `RegMem2SSAType` in `Fly/Utils/PointerUtils.cpp`). 2. The `AtomCall` is replaced with `AtomCallSSA`, taking those SSA values in place of pointers. 3. For output-producing cases, a `PtrStoreOp` writes the SSA result back to the original register memref. At lowering time, `AtomCallSSA` dispatches to `emitAtomCallSSA` instead of `emitAtomCall`. The Op's job there is **just the intrinsic + any required `LLVM::BitcastOp` between the SSA `vector<...>` and the intrinsic's expected packed type** — no loads or stores because the SSA values already live in registers. **Concrete differences between the two methods:** | | `emitAtomCall` | `emitAtomCallSSA` | |--------------------------|-------------------------------------------------|-------------------------------------------------| | Operand kinds | `Value`s of type `!fly.memref<..., register>` (lowered to `!llvm.ptr`) | `Value`s of scalar / `vector<Nxelem>` type | | What the method does | `LLVM::LoadOp` to fetch operands → intrinsic → `LLVM::StoreOp` to write result | (optional bitcast to intrinsic's packed type) → intrinsic → return `Value` / `failure` | | Return type | `LogicalResult` | `FailureOr<Value>` (the result SSA value, or `failure`) | | Needs layout/cosize info | No — operand type already carries it | No — caller already packed operands into `vector<N>` | | Bitcast dance | Typically unnecessary (load yields the right type) | Often necessary (SSA vector width may not match intrinsic's expected operand width) | | Backend intrinsic emitted | Same | Same | In practice every reference Op implements `emitAtomCall` as a thin shim over `emitAtomCallSSA` — load operands, call `emitAtomCallSSA`, store the result. See `MmaOpCDNA3_MFMAType::emitAtomCall` in `CDNA3/MmaAtom.cpp` for the canonical shim and `CopyOpCDNA3BufferAtomicType::emitAtomCall` in `CDNA3/CopyAtom.cpp` for a CopyOp instance. **If your downstream pipeline never runs `fly-convert-atom-call-to-ssa-form`, you may skip `emitAtomCallSSA` entirely and write a self-contained `emitAtomCall`** — but the shim pattern is strictly better because it keeps the two paths in sync for free. ### 1.4 ThrVal layouts describe the per-thread register footprint Every MmaOp / CopyOp must publish layouts that describe *which thread holds which element* of the tile. This is consumed by `TiledCopy` / `TiledMma` in the layout-lowering pass. | Method (MmaOp) | What it describes | |------------------------|-------------------| | `getThrLayout` | thread-count layout inside one thread group that issues the instruction (e.g. `(64):(1)` for an AMD wave64 MFMA, `(32):(1)` for AMD wave32 WMMA, `(1):(1)` for a single thread, `(128):(1)` for NVIDIA WGMMA issued by a warpgroup) | | `getShapeMNK` | tuple `(M, N, K)` of the instruction tile | | `getValTypeA/B/C/D` | per-operand element type | | `getThrValLayoutA/B/C` | layout mapping `(thr, val)` → element coordinate in the **reference tile** (column-major `(M,K)` for A, `(N,K)` for B, `(M,N)` for C) | | Method (CopyOp) | What it describes | |----------------------------|-------------------| | `getThrLayout` | thread count participating in one atom call | | `getThrBitLayoutSrc/Dst/Ref` | layout in **bit-granularity** — shape is `(num_threads, num_bits)` — one bit per leaf | The base `CopyAtomType::getThrValLayoutSrc()` then "recasts" the bit layout into a `valBits`-granularity layout (see `CopyAtomType::getThrValLayout{Src,Dst,Ref}` in `FlyTypeDefs.cpp`). This is why CopyOp types publish a **bit-layout** and MmaOp types publish a **value-layout**: copies carry an extra `valBits` parameter on the wrapper, and one CopyOp type can serve multiple element widths. Use the `FxLayout / FxShape / FxStride / FxThr / FxVal / FxC` macros from `flydsl/Dialect/Fly/Utils/ThrValLayoutMacro.h.inc` — they're the auxiliary way to build these `LayoutAttr`s. ### 1.5 Critical checks for ThrVal / ThrBit layouts — read before writing any A wrong ThrVal/ThrBit layout is the #1 source of silent-wrong-result bugs in FlyDSL: the compiler accepts it, the kernel runs, and the output is garbage. There are no good runtime diagnostics for this. Before you commit any new `getThrValLayout*` / `getThrBitLayout*`, verify **every** rule below on paper or in a scratch test. #### 1.5.1 Shape must be a top-level 2-tuple `((thr...), (val...))` Look at any existing example: `FxLayout(FxShape(FxThr(...), FxVal(...)), FxStride(FxThr(...), FxVal(...)))`. The top-level shape has exactly **rank 2**: outer mode 0 is the thread axes, outer