Atom 1250 refactor

CLAUDE.md

@@ -48,7 +48,7 @@ FlyDSL/
 ├── python/
 │   ├── flydsl/                    # Python DSL core
 │   │   ├── expr/                  # DSL expression API; direct children are TARGET-NEUTRAL (typing, primitive, gpu, derived, struct, numeric, math, vector, arith, meta, extern; + utils/)
-│   │   │   └── rocdl/             # Target-specific ROCDL package (cdna4, cluster, inline_asm, tdm_ops, universal); lazy-loaded via __init__'s _LAZY_MODULES
+│   │   │   └── rocdl/             # Target-specific ROCDL package (cdna4, cdna5, cluster, inline_asm, tdm_ops, universal); lazy-loaded via __init__'s _LAZY_MODULES
 │   │   ├── compiler/              # @flyc.kernel / @flyc.jit, AST rewriting, JIT cache, backends
 │   │   ├── runtime/               # Device runtime and GPU arch detection
 │   │   ├── utils/                 # EnvManager, SmemAllocator (legacy), logger
@@ -203,8 +203,8 @@ This is routing guidance, not a complete kernel inventory. Search the current `k
 - Avoid early `return` and branch-local `return` / `yield` in traced functions. Keep a single explicit exit path so MLIR result types stay well-defined.
 - Prefer arch-specific helper modules and constants over inline scattered `gfx*` conditionals.
 - **Helper placement.** Do not scatter small helpers across unrelated modules and do not duplicate an existing one; search for and reuse an existing helper first. Shared kernel helpers belong in `kernels/kernels_common.py` (wave size via `get_warp_size`, `dtype_to_elem_type`, `validate_moe_dtypes`, the `_if_then` SCF context manager, LLVM-ptr/stream helpers); domain-specific shared helpers go in the existing topical modules (`kernels/moe_common.py`, `layout_utils.py`, `pipeline_utils.py`, `fp8_gemm_utils.py`, `dpp_utils.py`, `mfma_epilogues.py`, `mfma_preshuffle_pipeline.py`). DSL-level numeric/arith and type helpers belong in `python/flydsl/expr/utils/arith.py` / `python/flydsl/expr/numeric.py`; compiler/runtime-wide utilities (env, logger, smem allocator) in `python/flydsl/utils/`. (PR #388 extracted shared `_if_then`/`validate_moe_dtypes` into `kernels_common.py`; PR #448 removed redundant numeric wrappers in favor of existing `fx.*` type methods.)
-- **`expr/` is target-neutral.** The direct child modules of `python/flydsl/expr/` (`typing`, `primitive`, `gpu`, `derived`, `struct`, `arith`, `math`, `vector`, `numeric`, `meta`, `extern`, `utils/`) must stay backend-agnostic: they may not import ROCDL/HIP bindings (`flydsl._mlir.dialects.rocdl`, `_mlirDialectsFlyROCDL`, `fly_rocdl`). `import flydsl.expr` must succeed without the FlyROCDL bindings; `tests/unit/test_expr_optional_rocdl.py` enforces this in CI. New target-specific (ROCDL/HIP, MFMA/WMMA, buffer/TDM/cluster) expr code goes in the `expr/rocdl/` package (`cdna4`, `cluster`, `inline_asm`, `tdm_ops`, `universal`), never in a new top-level `expr/*.py`. The target-specific modules `buffer_ops`, `rocdl`, and `tdm_ops` are lazy-loaded from `expr/__init__.py` via `__getattr__` (`_LAZY_MODULES`); add new backend modules to that lazy map rather than eager-importing them (PR #521).
-- **`expr/rocdl` is a package.** `expr/rocdl/` (`__init__.py` + `cluster.py`, `tdm_ops.py`, `cdna4.py`, `universal.py`, `inline_asm.py`) holds all target-specific ROCDL/MFMA/WMMA/buffer/TDM/cluster code. `from flydsl.expr import rocdl` and `flydsl.expr.rocdl` bind to `expr/rocdl/__init__.py`. Import submodules explicitly, e.g. `from flydsl.expr.rocdl import cluster`; `flydsl.expr.tdm_ops` is a lazy alias for `flydsl.expr.rocdl.tdm_ops` (see `expr/__init__.py` `_LAZY_MODULES`).
+- **`expr/` is target-neutral.** The direct child modules of `python/flydsl/expr/` (`typing`, `primitive`, `gpu`, `derived`, `struct`, `arith`, `math`, `vector`, `numeric`, `meta`, `extern`, `utils/`) must stay backend-agnostic: they may not import ROCDL/HIP bindings (`flydsl._mlir.dialects.rocdl`, `_mlirDialectsFlyROCDL`, `fly_rocdl`). `import flydsl.expr` must succeed without the FlyROCDL bindings; `tests/unit/test_expr_optional_rocdl.py` enforces this in CI. New target-specific (ROCDL/HIP, MFMA/WMMA, buffer/TDM/cluster) expr code goes in the `expr/rocdl/` package (`cdna4`, `cdna5`, `cluster`, `inline_asm`, `tdm_ops`, `universal`), never in a new top-level `expr/*.py`. The target-specific modules `buffer_ops`, `rocdl`, and `tdm_ops` are lazy-loaded from `expr/__init__.py` via `__getattr__` (`_LAZY_MODULES`); add new backend modules to that lazy map rather than eager-importing them (PR #521).
+- **`expr/rocdl` is a package.** `expr/rocdl/` (`__init__.py` + `cluster.py`, `tdm_ops.py`, `cdna4.py`, `cdna5.py`, `universal.py`, `inline_asm.py`) holds all target-specific ROCDL/MFMA/WMMA/buffer/TDM/cluster code (`cdna5.py` holds the gfx1250 TDM copy atom, re-exported top-level like `universal`). `from flydsl.expr import rocdl` and `flydsl.expr.rocdl` bind to `expr/rocdl/__init__.py`. Import submodules explicitly, e.g. `from flydsl.expr.rocdl import cluster`; `flydsl.expr.tdm_ops` is a lazy alias for `flydsl.expr.rocdl.tdm_ops` (see `expr/__init__.py` `_LAZY_MODULES`).
 
 ## Testing Notes
 

docs/api/compiler.rst

@@ -113,7 +113,7 @@ The ``flydsl.expr.rocdl`` module provides AMD-specific operations:
 - **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.make_tdm_atom** / **fx.rocdl.TDM** / **fx.rocdl.advance_tdm_atom** -- gfx1250 TDM async Global↔LDS whole-tile copy atom (1–5D; descriptor as atom state)
+- **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

@@ -213,7 +213,7 @@ AMD-specific operations for ROCm:
 - **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.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 tile descriptor is carried as atom state. ``fx.rocdl.TDM(rank, num_warps, ...)`` builds the atom type only; ``fx.rocdl.advance_tdm_atom(atom, byte_offset)`` bumps the K-loop tile offset
+- **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
 - **fx.rocdl.sched_vmem(cnt)** -- insert VMEM scheduling barrier
 - **fx.rocdl.sched_dsrd(cnt)** -- insert DS read scheduling barrier

docs/kernel_authoring_guide.md

@@ -364,11 +364,11 @@ mma = fx.atom_set_value(mma, "scale_b", fx.Int32(scale_b))
 fx.gemm(mma, frag_C, frag_A, frag_B, frag_C)
 ```
 
-**TDM async copy atom** — the descriptor (base pointer, per-dim extent for HW
-out-of-bounds handling, per-dim stride) is carried as **atom state**; the global
-operand of `copy_atom_call` is a *shape/direction token only* (its layout gives
-the compile-time N-D tile shape, its address space picks load vs store — its
-pointer is unused). Build it with `rocdl.make_tdm_atom`:
+**TDM async copy atom** — the **base pointer comes from the `copy_atom_call` global
+operand** (its pointer); the per-dim extent (HW out-of-bounds handling), per-dim
+stride, `imm_offset` (K-loop tile bump), and MCAST `workgroup_mask` are runtime
+**atom state**. The global operand's address space also picks load vs store. Build
+it with `rocdl.make_tdm_atom`:
 
 ```python
 # make_tdm_atom(tensor, tensor_extents, strides=None, *, num_warps,
@@ -378,17 +378,19 @@ lds = fx.SharedAllocator().allocate(fx.Array[fx.Float16, M * N]).peek()
 lds2d = fx.make_view(lds.ptr, fx.make_layout((M, N), (N, 1)))       # note: lds.ptr
 g2d = fx.make_view(fx.get_iter(A), fx.make_layout((M, N), (N, 1)))  # raw VA, not make_buffer_tensor
 
-atom = rocdl.make_tdm_atom(g2d, [M, N], num_warps=4)   # rank = len(extents), 1–5D
-fx.copy_atom_call(atom, g2d, lds2d)                    # Global → LDS (direction from address spaces)
-rocdl.tdm_ops.tensor_wait(0)                           # await the async DMA (s_wait_tensorcnt)
+atom = rocdl.make_tdm_atom(g2d, [M, N], num_warps=4)  # rank = len(extents), 1–5D
+fx.copy_atom_call(atom, g2d, lds2d)                   # Global → LDS (base from g2d)
+rocdl.tdm_ops.tensor_wait(0)                          # await the async DMA (s_wait_tensorcnt)
 
-# K-loop: bump one scalar instead of re-deriving base (imm_offset, carry-safe i64)
-atom = rocdl.advance_tdm_atom(atom, k_tile * k_stride_bytes)
+# K-loop: bump one scalar (imm_offset, carry-safe i64) instead of re-deriving base:
+fx.copy(atom, g2d, lds2d, imm_offset=k_tile * k_stride_bytes)
 ```
 
-`rocdl.TDM(rank, num_warps, ...)` builds just the atom *type* when you want to set
-the descriptor state manually. Unlike the CDNA buffer copy, TDM needs a **raw VA**
-— do not wrap the global tensor in `make_buffer_tensor`.
+`rocdl.TDM(rank, num_warps, ...)` (in `rocdl.cdna5`) builds just the atom *type*
+when you want to set the descriptor state manually. Pass `tensor_extents` smaller
+than the tile for ragged edges (HW OOB clamp) and `strides=` for a dynamic/true
+global stride that differs from the packed tile stride. Unlike the CDNA buffer copy,
+TDM needs a **raw VA** — do not wrap the global tensor in `make_buffer_tensor`.
 
 ### 4.5 GPU Operations (`fx.gpu`)
 

include/flydsl/Dialect/Fly/IR/FlyDialect.h

@@ -43,6 +43,14 @@ template <class TargetAddressSpace> bool isTargetAddressSpace(Attribute attr) {
   return llvm::isa_and_nonnull<TargetAddressSpace>(attr);
 }
 
+/// Type trait marking a copy-atom op type that moves a whole N-D tile in a single
+/// copy_atom_call (e.g. the gfx1250 TDM DMA). The expand-copy lowering checks this
+/// to emit one whole-tile call instead of decomposing the tiled memref per
+/// element. Defined here (single shared TypeID) so a target-neutral pass can query
+/// it via `hasTrait` — boundary-safe, unlike a concrete cross-dialect dyn_cast.
+template <typename ConcreteType>
+class WholeTileCopy : public TypeTrait::TraitBase<ConcreteType, WholeTileCopy> {};
+
 ParseResult parseMNKDimensionList(AsmParser &parser, int32_t &m, int32_t &n, int32_t &k);
 void printMNKDimensionList(AsmPrinter &printer, int32_t m, int32_t n, int32_t k);
 

include/flydsl/Dialect/Fly/IR/FlyInterfaces.td

@@ -132,19 +132,7 @@ def Fly_CopyOpTypeInterface : TypeInterface<"CopyOpTypeInterface"> {
            "::mlir::Value":$atomVal,
            "::mlir::Value":$src,
            "::mlir::Value":$dst,
-           "::mlir::Value":$pred)>,
-    InterfaceMethod<
-      "Number of memref dimensions this atom transfers in a single "
-      "copy_atom_call. Default 1: the expand-copy lowering peels the tiled "
-      "memref down to a rank-1 leaf and emits one call per element/fragment. An "
-      "atom that moves a whole N-D tile in one instruction (e.g. the gfx1250 "
-      "TDM 2D DMA) returns N, and the lowering emits a single call on the "
-      "rank-N memref instead of decomposing it.",
-      "unsigned",
-      "getCopyRank",
-      (ins),
-      /*methodBody=*/[{}],
-      /*defaultImplementation=*/[{ return 1; }]>
+           "::mlir::Value":$pred)>
   ];
 }
 

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

@@ -11,21 +11,12 @@ def FlyROCDL_AtomStateField : I32EnumAttr<"AtomStateField", "", [
   I32EnumAttrCase<"ImmOffset",          1, "imm_offset">,
   I32EnumAttrCase<"ScaleA",             2, "scale_a">,
   I32EnumAttrCase<"ScaleB",             3, "scale_b">,
-  I32EnumAttrCase<"WorkgroupMask",      4, "workgroup_mask">,
-  // TDM N-D descriptor carried on the copy atom: global base pointer, per-dim
-  // tensor extent (extent_i, i32, tensor dim order 0=outermost) for OOB handling,
-  // and per-dim tensor stride in elements (stride_i, i64; the innermost stride is
-  // assumed 1 and not stored, so stride_i covers dims 0..rank-2).
-  I32EnumAttrCase<"Base",               5, "base">,
-  I32EnumAttrCase<"Extent0",            6, "extent_0">,
-  I32EnumAttrCase<"Extent1",            7, "extent_1">,
-  I32EnumAttrCase<"Extent2",            8, "extent_2">,
-  I32EnumAttrCase<"Extent3",            9, "extent_3">,
-  I32EnumAttrCase<"Extent4",            10, "extent_4">,
-  I32EnumAttrCase<"Stride0",            11, "stride_0">,
-  I32EnumAttrCase<"Stride1",            12, "stride_1">,
-  I32EnumAttrCase<"Stride2",            13, "stride_2">,
-  I32EnumAttrCase<"Stride3",            14, "stride_3">
+  I32EnumAttrCase<"WorkgroupMask",      4, "workgroup_mask">
+  // NOTE: the gfx1250 TDM N-D descriptor's per-dim tensor extent (OOB) and per-dim
+  // tensor stride are NOT shared enum fields. They are resolved privately by
+  // CopyOpGFX1250TDMType from the field names "extent_0".."extent_4" /
+  // "stride_0".."stride_3" (set via fly.atom.set_value), keeping this cross-atom
+  // vocabulary free of TDM's per-dim sprawl. See GFX1250/CopyAtom.cpp (tdmGeomSlot).
 ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = FlyROCDL_Dialect.cppNamespace;

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

@@ -48,26 +48,25 @@ def FlyROCDL_CopyOpLdsReadTranspose : FlyROCDL_CopyOp<"CopyOpCDNA4LdsReadTranspo
 //===----------------------------------------------------------------------===//
 // CopyOp GFX1250 — N-D TDM (Tensor Data Mover) async Global<->LDS copy (1-5D).
 //
-// Whole-tile DMA. The atom carries the tile descriptor as runtime state. The
-// global operand of the copy_atom_call is a *shape/direction token only*: only
-// its layout (compile-time tile shape) and address space (load vs store) are
-// read — its pointer is unused; the base address comes from the `base` state.
-// `rank` (1-5) is the tensor/tile rank. Stateful fields (via fly.atom.set_value):
+// Whole-tile DMA. The global operand of the copy_atom_call supplies the base
+// pointer (its pointer, NOT atom state) and the compile-time tile shape (its
+// layout); its address space picks the direction (load vs store). `rank` (1-5) is
+// the tensor/tile rank. Stateful fields (via fly.atom.set_value):
 //   `workgroup_mask` (default 0): MCAST mask, ORed into GROUP1 config [15:0].
-//   `base`: global tile base pointer (advance it to move to the next tile).
 //   `extent_i` (default INT32_MAX = no clamp): per-dim tensor extent (tensor dim
 //     order, 0=outermost) for OOB handling; tensor_dim = max(0, extent - warp
 //     start) while the tile dims stay static, so a ragged tile is HW-guarded.
 //   `stride_i` (default = unset sentinel): per-dim tensor stride in elements for
-//     dims 0..rank-2 (innermost stride is assumed 1). Unset falls back to the
-//     tile memref's static layout stride.
+//     dims 0..rank-2 (innermost stride is assumed 1). Unset falls back to the tile
+//     memref's static layout stride; set it for a dynamic/true global stride that
+//     differs from the packed tile-internal stride.
 //   `imm_offset` (default 0): i64 byte offset added to base (K-loop tile bump).
 //===----------------------------------------------------------------------===//
 
-// TDM moves the whole N-D tile in one DMA, so it overrides `getCopyRank` (== rank)
+// TDM moves the whole N-D tile in one DMA, so it carries the WholeTileCopy trait
 // to keep the expand-copy lowering from decomposing the tiled memref per element.
 def FlyROCDL_CopyOpGFX1250TDM : FlyROCDL_StatefulCopyOp<"CopyOpGFX1250TDM", "gfx1250.tdm",
-    /*traits=*/[], /*copyOverrides=*/["getCopyRank"]> {
+    /*traits=*/[Fly_WholeTileCopyTrait]> {
   let parameters = (ins
     // Tensor/tile rank (1-5).
     "int32_t":$rank,

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

@@ -13,7 +13,7 @@ include "flydsl/Dialect/Fly/IR/FlyInterfaces.td"
 def FlyROCDL_Dialect : Dialect {
   let name = "fly_rocdl";
   let cppNamespace = "::mlir::fly_rocdl";
-  
+
   let dependentDialects = [
     "ROCDL::ROCDLDialect"
   ];
@@ -22,6 +22,14 @@ def FlyROCDL_Dialect : Dialect {
   let useDefaultAttributePrinterParser = 1;
 }
 
+// Marks a copy-atom op type that moves a whole N-D tile in one copy_atom_call
+// (e.g. gfx1250.tdm). The C++ trait `::mlir::fly::WholeTileCopy` lives in the Fly
+// dialect header so the target-neutral expand-copy lowering can query it via
+// `hasTrait` without a concrete cross-dialect cast.
+def Fly_WholeTileCopyTrait : NativeTypeTrait<"WholeTileCopy"> {
+  let cppNamespace = "::mlir::fly";
+}
+
 class FlyROCDL_I32EnumAttr<string enumName, string summary, list<I32EnumAttrCase> cases>
     : I32EnumAttr<enumName, summary, cases> {
   let genSpecializedAttr = 0;
@@ -52,13 +60,9 @@ class FlyROCDL_CopyOp<string typeName, string typeMnemonic, list<Trait> traits =
   let mnemonic = typeMnemonic;
 }
 
-// `copyOverrides` lists CopyOpTypeInterface methods that have a default
-// implementation but are overridden by this type (e.g. `getCopyRank`); they
-// must be named so TableGen emits their declarations.
-class FlyROCDL_StatefulCopyOp<string typeName, string typeMnemonic, list<Trait> traits = [],
-                              list<string> copyOverrides = []>
+class FlyROCDL_StatefulCopyOp<string typeName, string typeMnemonic, list<Trait> traits = []>
     : TypeDef<FlyROCDL_Dialect, typeName, !listconcat(traits, [
-        DeclareTypeInterfaceMethods<Fly_CopyOpTypeInterface, copyOverrides>,
+        DeclareTypeInterfaceMethods<Fly_CopyOpTypeInterface>,
         DeclareTypeInterfaceMethods<Fly_StatefulOpTypeInterface>
       ])> {
   let mnemonic = typeMnemonic;