add-target-atom-op

name	add-target-atom-op
description	Add a new target-specific Mma / Copy Op type to a FlyDSL backend dialect (`lib/Dialect/Fly<TARGET>/<SUBTARGET>/` + `include/flydsl/Dialect/Fly<TARGET>/IR/`). Covers the MmaOp/CopyOp type design, the stateful-vs-stateless variants, and the `emitAtomCall` / `emitAtomCallSSA` lowering contract to the backend dialect (LLVM/ROCDL/NVVM/SPIR-V/...). Use when adding a new tensor-core / matrix instruction (MFMA, WMMA, HMMA, WGMMA, ...), a new buffer / shared-memory / global copy atom, a new stateful copy (per-atom offset or descriptor), or bringing up a new backend dialect (`FlyPTX`, `FlyCPU`, ...). Reference implementations live in `FlyROCDL`. Usage: /add-target-atom-op
allowed-tools	Read Edit Bash Grep Glob Agent

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.

// 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:

// 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:

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	Values of type !fly.memref<..., register> (lowered to !llvm.ptr)	Values 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 LayoutAttrs.

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