From d88d1b349955aeaa69d1c86d5770286306023d20 Mon Sep 17 00:00:00 2001 From: Felix Li Date: Tue, 14 Jul 2026 11:56:51 +0000 Subject: [PATCH 1/3] fix changes --- CLAUDE.md | 6 +- docs/api/compiler.rst | 2 +- docs/api/dsl.rst | 2 +- docs/kernel_authoring_guide.md | 28 ++-- include/flydsl/Dialect/Fly/IR/FlyDialect.h | 8 ++ .../flydsl/Dialect/Fly/IR/FlyInterfaces.td | 14 +- include/flydsl/Dialect/FlyROCDL/IR/Atom.td | 28 ++-- .../flydsl/Dialect/FlyROCDL/IR/CopyAtom.td | 19 ++- include/flydsl/Dialect/FlyROCDL/IR/Dialect.td | 18 ++- lib/Dialect/Fly/Transforms/LayoutLowering.cpp | 29 ++-- lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp | 55 ++++---- python/flydsl/expr/rocdl/__init__.py | 2 + python/flydsl/expr/rocdl/cdna5.py | 115 ++++++++++++++++ python/flydsl/expr/rocdl/universal.py | 128 ------------------ tests/kernels/test_gfx1250_atoms_device.py | 6 +- tests/mlir/Conversion/tdm_gfx1250.mlir | 39 +++--- tests/unit/test_gfx1250_atoms.py | 6 +- 17 files changed, 238 insertions(+), 267 deletions(-) create mode 100644 python/flydsl/expr/rocdl/cdna5.py diff --git a/CLAUDE.md b/CLAUDE.md index d6e1afec7..1915d770b 100644 --- a/CLAUDE.md +++ b/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 diff --git a/docs/api/compiler.rst b/docs/api/compiler.rst index 62e75f596..5495841e3 100644 --- a/docs/api/compiler.rst +++ b/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 ------------ diff --git a/docs/api/dsl.rst b/docs/api/dsl.rst index ef45a98a6..b2b23ff3f 100644 --- a/docs/api/dsl.rst +++ b/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 diff --git a/docs/kernel_authoring_guide.md b/docs/kernel_authoring_guide.md index 76843b6da..7d93269ff 100644 --- a/docs/kernel_authoring_guide.md +++ b/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`) diff --git a/include/flydsl/Dialect/Fly/IR/FlyDialect.h b/include/flydsl/Dialect/Fly/IR/FlyDialect.h index 5f8f21cb0..22034a928 100644 --- a/include/flydsl/Dialect/Fly/IR/FlyDialect.h +++ b/include/flydsl/Dialect/Fly/IR/FlyDialect.h @@ -43,6 +43,14 @@ template bool isTargetAddressSpace(Attribute attr) { return llvm::isa_and_nonnull(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 +class WholeTileCopy : public TypeTrait::TraitBase {}; + 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); diff --git a/include/flydsl/Dialect/Fly/IR/FlyInterfaces.td b/include/flydsl/Dialect/Fly/IR/FlyInterfaces.td index 9c618053c..414c8dab2 100644 --- a/include/flydsl/Dialect/Fly/IR/FlyInterfaces.td +++ b/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)> ]; } diff --git a/include/flydsl/Dialect/FlyROCDL/IR/Atom.td b/include/flydsl/Dialect/FlyROCDL/IR/Atom.td index 9d1f44fe5..10d0d44fb 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/Atom.td +++ b/include/flydsl/Dialect/FlyROCDL/IR/Atom.td @@ -12,20 +12,20 @@ def FlyROCDL_AtomStateField : I32EnumAttr<"AtomStateField", "", [ 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"> + // TDM N-D descriptor carried on the copy atom (the global base comes from the + // copy_atom_call operand pointer, not state): 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<"Extent0", 5, "extent_0">, + I32EnumAttrCase<"Extent1", 6, "extent_1">, + I32EnumAttrCase<"Extent2", 7, "extent_2">, + I32EnumAttrCase<"Extent3", 8, "extent_3">, + I32EnumAttrCase<"Extent4", 9, "extent_4">, + I32EnumAttrCase<"Stride0", 10, "stride_0">, + I32EnumAttrCase<"Stride1", 11, "stride_1">, + I32EnumAttrCase<"Stride2", 12, "stride_2">, + I32EnumAttrCase<"Stride3", 13, "stride_3"> ]> { let genSpecializedAttr = 0; let cppNamespace = FlyROCDL_Dialect.cppNamespace; diff --git a/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td b/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td index 0f5961f9e..3ca656ec0 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td +++ b/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, diff --git a/include/flydsl/Dialect/FlyROCDL/IR/Dialect.td b/include/flydsl/Dialect/FlyROCDL/IR/Dialect.td index 1c2386c11..37974c96c 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/Dialect.td +++ b/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 cases> : I32EnumAttr { let genSpecializedAttr = 0; @@ -52,13 +60,9 @@ class FlyROCDL_CopyOp 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 traits = [], - list copyOverrides = []> +class FlyROCDL_StatefulCopyOp traits = []> : TypeDef, + DeclareTypeInterfaceMethods, DeclareTypeInterfaceMethods ])> { let mnemonic = typeMnemonic; diff --git a/lib/Dialect/Fly/Transforms/LayoutLowering.cpp b/lib/Dialect/Fly/Transforms/LayoutLowering.cpp index b58a710d5..1ad4842e5 100644 --- a/lib/Dialect/Fly/Transforms/LayoutLowering.cpp +++ b/lib/Dialect/Fly/Transforms/LayoutLowering.cpp @@ -2192,27 +2192,16 @@ class ExpandCopyOpLowering : public OpRewritePattern { if (srcRank != dstRank) return rewriter.notifyMatchFailure(op, "src/dst ranks mismatch"); - // Atoms that move a whole N-D tile in one call (e.g. TDM DMA) declare - // getCopyRank() == N and read the tile geometry from the rank-N memref - // layout. Emit a single call once the tile matches that rank, instead of - // decomposing it. The default rank 1 (guarded by > 1 here) leaves the - // per-element path below unchanged for every other atom. + // A whole-tile copy atom (e.g. the gfx1250 TDM DMA) moves the entire N-D tile + // in one call and reads its geometry from the operand memref layout, so emit a + // single call on the tile instead of decomposing it per element. The atom's own + // emitAtomCall verifies the operand rank matches the tile. Detected via a + // boundary-safe type trait rather than a concrete cross-dialect cast. if (auto copyAtomTy = dyn_cast(copyAtomVal.getType())) { - if (auto copyOp = dyn_cast(copyAtomTy.getCopyOp())) { - unsigned copyRank = copyOp.getCopyRank(); - if (copyRank > 1) { - // A whole-tile atom reads its geometry from the rank-N memref, so the - // tile must be presented at exactly that rank. Anything else (a - // larger, not-yet-tiled memref or a lower-rank tile) has no valid - // per-element decomposition for this atom — fail loudly here rather - // than silently peeling it apart into wrong per-element calls below. - if (srcRank != static_cast(copyRank)) - return rewriter.notifyMatchFailure( - op, "whole-tile copy atom requires a memref of its exact copy rank"); - CopyAtomCall::create(rewriter, loc, copyAtomVal, src, dst, pred); - rewriter.eraseOp(op); - return success(); - } + if (copyAtomTy.getCopyOp().hasTrait()) { + CopyAtomCall::create(rewriter, loc, copyAtomVal, src, dst, pred); + rewriter.eraseOp(op); + return success(); } } diff --git a/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp b/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp index 8cba620ea..b742f82f0 100644 --- a/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp +++ b/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp @@ -98,42 +98,40 @@ constexpr int32_t kOuterStrideUnset = static_cast(0x80000000); } // namespace -// Stateful: the atom carries the whole TDM N-D descriptor — workgroup_mask -// (MCAST) plus the global tile descriptor. Struct slots: {mask, base, -// extent_0..4 (i32, per-dim tensor extent for OOB), stride_0..3 (i64, per-dim -// tensor stride in elements; innermost stride is assumed 1), imm_offset (i64)}. -// See CopyAtom.td for field semantics. +// Stateful: the atom carries the TDM N-D descriptor — workgroup_mask (MCAST) plus +// the tile geometry, EXCEPT the global base, which comes from the copy_atom_call +// operand pointer. Struct slots: {mask, extent_0..4 (i32, per-dim tensor extent +// for OOB), stride_0..3 (i64, per-dim tensor stride in elements; innermost stride +// is assumed 1), imm_offset (i64)}. See CopyAtom.td for field semantics. static constexpr unsigned kMaxTdmRank = 5; std::optional CopyOpGFX1250TDMType::getFieldIndex(AtomStateField field) { switch (field) { case AtomStateField::WorkgroupMask: return 0; - case AtomStateField::Base: - return 1; case AtomStateField::Extent0: - return 2; + return 1; case AtomStateField::Extent1: - return 3; + return 2; case AtomStateField::Extent2: - return 4; + return 3; case AtomStateField::Extent3: - return 5; + return 4; case AtomStateField::Extent4: - return 6; + return 5; case AtomStateField::Stride0: - return 7; + return 6; case AtomStateField::Stride1: - return 8; + return 7; case AtomStateField::Stride2: - return 9; + return 8; case AtomStateField::Stride3: - return 10; + return 9; // Byte offset added to the global base at lowering (i64, carry-safe). Lets a // K-loop advance the tile by bumping one scalar (imm_offset) instead of // re-deriving the base pointer. Default 0 folds away. case AtomStateField::ImmOffset: - return 11; + return 10; default: return std::nullopt; } @@ -150,8 +148,7 @@ static AtomStateField strideField(unsigned i) { Type CopyOpGFX1250TDMType::getConvertedType(MLIRContext *ctx) const { auto i32 = IntegerType::get(ctx, 32); auto i64 = IntegerType::get(ctx, 64); - auto glbPtr = LLVM::LLVMPointerType::get(ctx, /*addrSpace=*/1); - SmallVector fields = {i32, glbPtr}; + SmallVector fields = {i32}; // workgroup_mask fields.append(kMaxTdmRank, i32); // extent_0..4 fields.append(kMaxTdmRank - 1, i64); // stride_0..3 fields.push_back(i64); // imm_offset @@ -169,11 +166,8 @@ Value CopyOpGFX1250TDMType::getDefaultState(OpBuilder &builder, Location loc) co Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32); Value noClamp = arith::ConstantIntOp::create(builder, loc, 0x7FFFFFFF, 32); Value strideUnset = arith::ConstantIntOp::create(builder, loc, kOuterStrideUnset, 64); - Value nullBase = - LLVM::ZeroOp::create(builder, loc, LLVM::LLVMPointerType::get(ctx, /*addrSpace=*/1)); Value zero64 = arith::ConstantIntOp::create(builder, loc, 0, 64); insert(AtomStateField::WorkgroupMask, zero); - insert(AtomStateField::Base, nullBase); for (unsigned i = 0; i < kMaxTdmRank; ++i) insert(extentField(i), noClamp); for (unsigned i = 0; i < kMaxTdmRank - 1; ++i) @@ -196,10 +190,6 @@ Value CopyOpGFX1250TDMType::setAtomState(OpBuilder &builder, Location loc, Value return LLVM::InsertValueOp::create(builder, loc, atomStruct, fieldValue, ArrayRef{*idx}); } -// TDM moves the whole N-D tile in one DMA; its geometry lives in the rank-N -// memref layout, so the expand-copy lowering must emit a single rank-N call. -unsigned CopyOpGFX1250TDMType::getCopyRank() const { return static_cast(getRank()); } - Attribute CopyOpGFX1250TDMType::getThrLayout() const { return FxLayout(FxC(1), FxC(1)); } Attribute CopyOpGFX1250TDMType::getThrBitLayoutSrc() const { @@ -274,10 +264,13 @@ LogicalResult CopyOpGFX1250TDMType::emitAtomCall(OpBuilder &builder, Location lo if (!isLoad && !isStore) return failure(); - // The global operand only marks the direction; its base pointer, extents, and - // strides come from the atom descriptor state. The tile shape is compile-time on - // the operand layout (tensor dim order: index 0 = outermost, rank-1 = innermost). + // The global operand supplies the base pointer and the compile-time tile shape; + // its per-dim extents (OOB) and strides come from the atom descriptor state (the + // tile view's static stride is the packed tile-internal stride, not necessarily + // the true global stride, which may be dynamic). Tensor dim order on the layout: + // index 0 = outermost, rank-1 = innermost. fly::MemRefType glbMemTy = isLoad ? srcMemTy : dstMemTy; + Value glbPtr = isLoad ? src : dst; Value ldsPtr = isLoad ? dst : src; int32_t rank = getRank(); @@ -339,8 +332,8 @@ LogicalResult CopyOpGFX1250TDMType::emitAtomCall(OpBuilder &builder, Location lo } strideElems[rank - 1] = arith::ConstantIntOp::create(builder, loc, 1, 64); // innermost contiguous - Value glbBasePtr = stateField(AtomStateField::Base); - Value glbBase = LLVM::PtrToIntOp::create(builder, loc, i64Ty, glbBasePtr); + // Global base = the copy_atom_call operand pointer (not atom state). + Value glbBase = LLVM::PtrToIntOp::create(builder, loc, i64Ty, glbPtr); // i64 byte offset (default 0) added to the base so a K-loop can advance the tile // by bumping one scalar; the carry into glb_hi is handled by the i64 split below. glbBase = arith::AddIOp::create(builder, loc, glbBase, stateField(AtomStateField::ImmOffset)); diff --git a/python/flydsl/expr/rocdl/__init__.py b/python/flydsl/expr/rocdl/__init__.py index 8d2732bec..4b18f15cc 100644 --- a/python/flydsl/expr/rocdl/__init__.py +++ b/python/flydsl/expr/rocdl/__init__.py @@ -17,6 +17,7 @@ from ..._mlir.dialects.rocdl import * # noqa: F401,F403 from ..meta import dsl_loc_tracing from . import cdna4 as cdna4 +from . import cdna5 as cdna5 from .enum import SyncScope as SyncScope # Keep references to ODS-generated builders so we can wrap them without losing access. @@ -458,6 +459,7 @@ def lds_transpose_load(result_type, lds_memref, elem_offset, elem_bytes): # ── New high-level helpers from universal.py ────────────────────────── from .universal import * # noqa: E402,F401,F403,I001 +from .cdna5 import * # noqa: E402,F401,F403,I001 from .inline_asm import * # noqa: E402,F401,F403,I001 # ── Wrappers: accept DSL Numeric args (fx.Int32, fx.Float32, etc.) ───────── diff --git a/python/flydsl/expr/rocdl/cdna5.py b/python/flydsl/expr/rocdl/cdna5.py new file mode 100644 index 000000000..b61617524 --- /dev/null +++ b/python/flydsl/expr/rocdl/cdna5.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""gfx1250-specific ROCDL atom builders (N-D TDM whole-tile Global<->LDS copy).""" + +from ..._mlir.dialects.fly_rocdl import CopyOpGFX1250TDMType +from ..typing import Int32, Int64, Tensor + + +def TDM( + rank, + num_warps, + pad_interval=0, + pad_amount=0, + cache_modifier=0, + atomic_barrier=False, + early_timeout=False, +): + """Create a gfx1250 N-D TDM (Tensor Data Mover) Global<->LDS copy atom *type*. + + ``rank`` is the tensor/tile rank (1-5). Direction is inferred at lowering from + which side is Global vs Shared; the tile shape is compile-time on the operand + layout. ``pad_interval`` / ``pad_amount`` (elements) add LDS row padding on the + load path. + + ``atomic_barrier`` (descriptor bit 18, HW auto-barrier) and ``early_timeout`` + (bit 21, multicast-load GL1 knob) set compile-time descriptor config bits. + + The global base pointer comes from the ``copy_atom_call`` global operand; the + per-dim extent (OOB), per-dim stride, ``imm_offset`` (K-loop tile bump), and the + MCAST ``workgroup_mask`` are runtime atom state set via ``fx.atom.set_value``. + :func:`make_tdm_atom` builds the atom and populates the descriptor from a tensor. + """ + return CopyOpGFX1250TDMType.get( + rank, + num_warps, + pad_interval, + pad_amount, + cache_modifier, + atomic_barrier=atomic_barrier, + early_timeout=early_timeout, + ) + + +def make_tdm_atom( + tensor: Tensor, + tensor_extents, + strides=None, + *, + num_warps, + pad_interval=0, + pad_amount=0, + cache_modifier=0, + atomic_barrier=False, + early_timeout=False, +) -> object: + """Build a gfx1250 N-D TDM copy atom carrying ``tensor``'s tile descriptor. + + The global base pointer comes from the ``copy_atom_call`` global operand (not + atom state); the atom carries the tensor's per-dim extent (for hardware + out-of-bounds handling: load zero-fill, store drop) and per-dim strides. Reuse + the atom across a tile loop; advance the tile via the ``imm_offset`` state + (``fx.copy(atom, gt, dst, imm_offset=...)``) or by advancing the global operand. + + ``tensor_extents`` is a list of the tensor's per-dim extent in tensor dim order + ``[dim0(outermost) .. dim_{rank-1}(innermost)]`` (rank = ``len(tensor_extents)``, + 1-5); each entry is a Python ``int`` or an ``i32`` / ``index`` runtime value (or + any ``fx`` integer), and ``None`` means no clamp on that axis (INT32_MAX). + ``strides`` is an optional list of per-dim strides in elements (same order); + the innermost stride is assumed 1 and ignored, so entries for dims 0..rank-2 + are used. ``None`` (or a ``None`` entry) falls back to the tile memref's static + layout stride; pass it explicitly for a tile whose true (or dynamic) outer + stride differs from the packed tile-internal stride. + + Issue the copy with ``fx.copy_atom_call(atom, global_tile, lds)``: the global + operand supplies both the copy direction (address space) and the base pointer. + """ + from ..primitive import atom_set_value, make_copy_atom + + NO_CLAMP = 0x7FFFFFFF + STRIDE_UNSET = -0x80000000 # matches kOuterStrideUnset in CopyAtom.cpp + + extents = list(tensor_extents) + rank = len(extents) + if not 1 <= rank <= 5: + raise ValueError(f"make_tdm_atom: rank must be in [1, 5], got {rank}") + strides = list(strides) if strides is not None else [None] * rank + if len(strides) != rank: + raise ValueError(f"make_tdm_atom: expected {rank} strides, got {len(strides)}") + + copy_op = CopyOpGFX1250TDMType.get( + rank, + num_warps, + pad_interval, + pad_amount, + cache_modifier, + atomic_barrier=atomic_barrier, + early_timeout=early_timeout, + ) + atom = make_copy_atom(copy_op, tensor.element_type) + for i in range(rank): + ext = ( + Int32(NO_CLAMP) + if extents[i] is None + else (extents[i] if isinstance(extents[i], Int32) else Int32(extents[i])) + ) + atom = atom_set_value(atom, f"extent_{i}", ext) + for i in range(rank - 1): # innermost stride assumed 1, not stored + st = ( + Int64(STRIDE_UNSET) + if strides[i] is None + else (strides[i] if isinstance(strides[i], Int64) else Int64(strides[i])) + ) + atom = atom_set_value(atom, f"stride_{i}", st) + return atom diff --git a/python/flydsl/expr/rocdl/universal.py b/python/flydsl/expr/rocdl/universal.py index f44afe90e..ba01ce04e 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/python/flydsl/expr/rocdl/universal.py @@ -13,7 +13,6 @@ CopyOpCDNA3BufferAtomicType, CopyOpCDNA3BufferCopyLDSType, CopyOpCDNA3BufferCopyType, - CopyOpGFX1250TDMType, MmaOpCDNA3_MFMAType, TargetAddressSpace, ) @@ -189,41 +188,6 @@ def WMMAScale( ) -def TDM( - rank, - num_warps, - pad_interval=0, - pad_amount=0, - cache_modifier=0, - atomic_barrier=False, - early_timeout=False, -): - """Create a gfx1250 N-D TDM (Tensor Data Mover) Global<->LDS copy atom *type*. - - ``rank`` is the tensor/tile rank (1-5). Direction is inferred at lowering from - which side is Global vs Shared; the tile shape is compile-time on the operand - layout. ``pad_interval`` / ``pad_amount`` (elements) add LDS row padding on the - load path. - - ``atomic_barrier`` (descriptor bit 18, HW auto-barrier) and ``early_timeout`` - (bit 21, multicast-load GL1 knob) set compile-time descriptor config bits. - - The tile descriptor (global base pointer, per-dim extent for out-of-bounds - handling, per-dim stride) plus the MCAST ``workgroup_mask`` are runtime atom - state set via ``fx.atom.set_value``. :func:`make_tdm_atom` builds the atom and - populates that descriptor from a tensor in one call. - """ - return CopyOpGFX1250TDMType.get( - rank, - num_warps, - pad_interval, - pad_amount, - cache_modifier, - atomic_barrier=atomic_barrier, - early_timeout=early_timeout, - ) - - def make_buffer_ptr(ptr: Pointer, num_records_bytes=None): """Construct a new buffer-resource (``BufferDesc``) pointer from a global pointer, for hardware OOB-checked loads / stores. @@ -299,98 +263,6 @@ def make_buffer_tensor( return make_view(buf_ptr, layout) -def make_tdm_atom( - tensor: Tensor, - tensor_extents, - strides=None, - *, - num_warps, - pad_interval=0, - pad_amount=0, - cache_modifier=0, - atomic_barrier=False, - early_timeout=False, -) -> object: - """Build a gfx1250 N-D TDM copy atom carrying ``tensor``'s tile descriptor. - - The atom holds the global tile descriptor as runtime state: base pointer, the - tensor's per-dim extent (for hardware out-of-bounds handling: load zero-fill, - store drop), and per-dim strides. Reuse the atom across a tile loop; to move to - the next tile re-set ``base`` (or bump ``imm_offset`` via - :func:`advance_tdm_atom`). - - ``tensor_extents`` is a list of the tensor's per-dim extent in tensor dim order - ``[dim0(outermost) .. dim_{rank-1}(innermost)]`` (rank = ``len(tensor_extents)``, - 1-5); each entry is a Python ``int`` or an ``i32`` / ``index`` runtime value (or - any ``fx`` integer), and ``None`` means no clamp on that axis (INT32_MAX). - ``strides`` is an optional list of per-dim strides in elements (same order); - the innermost stride is assumed 1 and ignored, so entries for dims 0..rank-2 - are used. ``None`` (or a ``None`` entry) falls back to the tile memref's static - layout stride; pass it explicitly for a tile with a dynamic outer stride (else - the descriptor stride faults rather than silently reading stride 0). - - Issue the copy with ``fx.copy_atom_call(atom, global_tile, lds)``. NOTE: the - global operand's *pointer is unused* — only its layout (tile shape) and address - space (copy direction) are read; the base comes from the ``base`` state. - """ - from ..primitive import atom_set_value, make_copy_atom - - NO_CLAMP = 0x7FFFFFFF - STRIDE_UNSET = -0x80000000 # matches kOuterStrideUnset in CopyAtom.cpp - - extents = list(tensor_extents) - rank = len(extents) - if not 1 <= rank <= 5: - raise ValueError(f"make_tdm_atom: rank must be in [1, 5], got {rank}") - strides = list(strides) if strides is not None else [None] * rank - if len(strides) != rank: - raise ValueError(f"make_tdm_atom: expected {rank} strides, got {len(strides)}") - - copy_op = CopyOpGFX1250TDMType.get( - rank, - num_warps, - pad_interval, - pad_amount, - cache_modifier, - atomic_barrier=atomic_barrier, - early_timeout=early_timeout, - ) - atom = make_copy_atom(copy_op, tensor.element_type) - atom = atom_set_value(atom, "base", get_iter(tensor)) - for i in range(rank): - ext = ( - Int32(NO_CLAMP) - if extents[i] is None - else (extents[i] if isinstance(extents[i], Int32) else Int32(extents[i])) - ) - atom = atom_set_value(atom, f"extent_{i}", ext) - for i in range(rank - 1): # innermost stride assumed 1, not stored - st = ( - Int64(STRIDE_UNSET) - if strides[i] is None - else (strides[i] if isinstance(strides[i], Int64) else Int64(strides[i])) - ) - atom = atom_set_value(atom, f"stride_{i}", st) - return atom - - -def advance_tdm_atom(atom, byte_offset) -> object: - """Return a TDM atom with its global byte offset (``imm_offset``) set. - - The offset is added to the ``base`` pointer in i64 at lowering (carry-safe), - so a K-reduction loop can advance the tile by bumping this single scalar - instead of re-deriving and re-setting the ``base`` pointer each iteration. - ``byte_offset`` is the cumulative byte delta from ``base`` (typically - ``k_tile * k_stride_bytes``); it replaces (does not accumulate onto) any - previously-set ``imm_offset``. Accepts a Python ``int`` or an ``i64`` / any - ``fx`` integer value. - """ - from ..primitive import atom_set_value - - off = byte_offset if isinstance(byte_offset, Int64) else Int64(byte_offset) - return atom_set_value(atom, "imm_offset", off) - - @dsl_loc_tracing def get_buffer_rsrc(ptr: Pointer): """Extract the raw ROCDL buffer resource (``!llvm.ptr<8>``) from a diff --git a/tests/kernels/test_gfx1250_atoms_device.py b/tests/kernels/test_gfx1250_atoms_device.py index b0ead5a52..50b4bcce2 100644 --- a/tests/kernels/test_gfx1250_atoms_device.py +++ b/tests/kernels/test_gfx1250_atoms_device.py @@ -71,9 +71,9 @@ def tdm_roundtrip_kernel(A: fx.Tensor, C: fx.Tensor): lds = fx.SharedAllocator().allocate(fx.Array[fx.Float16, M * N]).peek() lds2d = fx.make_view(lds.ptr, fx.make_layout((M, N), (N, 1))) - # Global -> LDS. tensor_extents = full tile (no OOB clamp); the innermost - # stride is assumed 1 so only the outer stride (N) is carried, and it - # falls back to the static layout stride here (passed None). + # Global -> LDS. The base pointer comes from the global operand; the tile + # extents (full tile => no OOB clamp) and strides (None => static layout + # fallback) are atom state populated by make_tdm_atom. load_atom = fx.rocdl.make_tdm_atom(a2d, [M, N], num_warps=num_warps) fx.copy_atom_call(load_atom, a2d, lds2d) tdm_ops.tensor_wait(0) diff --git a/tests/mlir/Conversion/tdm_gfx1250.mlir b/tests/mlir/Conversion/tdm_gfx1250.mlir index bd81a6f60..92364d062 100644 --- a/tests/mlir/Conversion/tdm_gfx1250.mlir +++ b/tests/mlir/Conversion/tdm_gfx1250.mlir @@ -2,18 +2,18 @@ // Copyright (c) 2025 FlyDSL Project Contributors // RUN: %fly-opt %s --convert-fly-to-rocdl | FileCheck %s -// gfx1250 N-D TDM copy atom lowering. The atom carries the tile descriptor as -// runtime state (base pointer, per-dim extent_i, per-dim stride_i, imm_offset); -// the global operand only marks the copy direction and supplies the compile-time -// tile shape. Struct: {mask, base, extent_0..4 (i32), stride_0..3 (i64), imm_offset -// (i64)}. +// gfx1250 N-D TDM copy atom lowering. The global base pointer comes from the +// copy_atom_call operand; the atom carries the rest of the descriptor as runtime +// state (per-dim extent_i, per-dim stride_i, imm_offset). The global operand marks +// the copy direction and supplies the compile-time tile shape + base pointer. +// Struct: {mask, extent_0..4 (i32), stride_0..3 (i64), imm_offset (i64)}. // Global -> Shared => rocdl.tensor.load.to.lds // Shared -> Global => rocdl.tensor.store.from.lds // ----- // CHECK-LABEL: @test_tdm_type -// CHECK-SAME: (%{{.*}}: !llvm.struct<(i32, ptr<1>, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)>) +// CHECK-SAME: (%{{.*}}: !llvm.struct<(i32, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)>) func.func @test_tdm_type( %atom: !fly.copy_atom, 0>) { return @@ -21,9 +21,9 @@ func.func @test_tdm_type( // ----- -// Load, single warp, rank 2: base / extents / stride come from atom state. The -// global operand's pointer is unused (only its layout gives the tile shape). -// extent_0 = outer (slot 2), extent_1 = inner (slot 3), stride_0 = i64 (slot 7). +// Load, single warp, rank 2: base comes from the global operand pointer; extents +// and stride come from atom state. extent_0 = outer (slot 1), extent_1 = inner +// (slot 2), stride_0 = i64 (slot 6). // CHECK-LABEL: @test_tdm_load func.func @test_tdm_load( @@ -35,13 +35,12 @@ func.func @test_tdm_load( %a2 = fly.atom.set_value(%a1, "extent_1", %ie) : (!fly.copy_atom, 0>, i32) -> !fly.copy_atom, 0> %a3 = fly.atom.set_value(%a2, "stride_0", %os) : (!fly.copy_atom, 0>, i64) -> !fly.copy_atom, 0> // stride_0 fallback: select(stride == unset-sentinel (i64), static_layout_stride, stride) - // CHECK-DAG: %[[STRIDE:.*]] = llvm.extractvalue %{{.*}}[7] : !llvm.struct<(i32, ptr<1>, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)> + // CHECK-DAG: %[[STRIDE:.*]] = llvm.extractvalue %{{.*}}[6] : !llvm.struct<(i32, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)> // CHECK-DAG: %[[SENT:.*]] = arith.constant -2147483648 : i64 // CHECK: arith.cmpi eq, %[[STRIDE]], %[[SENT]] // CHECK: arith.select - // base (slot 1) -> global address via ptrtoint. - // CHECK: %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(i32, ptr<1>, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)> - // CHECK: llvm.ptrtoint %[[BASE]] : !llvm.ptr<1> to i64 + // base -> global address via ptrtoint of the global operand pointer. + // CHECK: llvm.ptrtoint %{{.*}} : !llvm.ptr<1> to i64 // OOB clamp from the extent state fields. // CHECK: arith.subi // CHECK: arith.maxsi @@ -95,7 +94,7 @@ func.func @test_tdm_load_3d( %dst: !fly.memref) { %a1 = fly.atom.set_value(%atom, "stride_0", %s0) : (!fly.copy_atom, 0>, i64) -> !fly.copy_atom, 0> %a2 = fly.atom.set_value(%a1, "stride_1", %s1) : (!fly.copy_atom, 0>, i64) -> !fly.copy_atom, 0> - // CHECK: llvm.extractvalue %{{.*}}[8] : !llvm.struct<(i32, ptr<1>, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)> + // CHECK: llvm.extractvalue %{{.*}}[7] : !llvm.struct<(i32, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)> // CHECK: rocdl.tensor.load.to.lds fly.copy_atom_call(%a2, %src, %dst) : (!fly.copy_atom, 0>, !fly.memref, !fly.memref) -> () return @@ -180,8 +179,9 @@ func.func @test_tdm_load_barrier_timeout( // ----- -// imm_offset (state slot 11, i64): a K-loop advances the tile by bumping one -// scalar, added to the base in i64 after ptrtoint (carry into glb_hi is automatic). +// imm_offset (state slot 10, i64): a K-loop advances the tile by bumping one +// scalar, added to the base (from the operand pointer) in i64 after ptrtoint +// (carry into glb_hi is automatic). // CHECK-LABEL: @test_tdm_load_imm_offset func.func @test_tdm_load_imm_offset( @@ -189,11 +189,10 @@ func.func @test_tdm_load_imm_offset( %off: i64, %src: !fly.memref, %dst: !fly.memref) { - // CHECK: %[[A1:.*]] = llvm.insertvalue %{{.*}}, %{{.*}}[11] : !llvm.struct<(i32, ptr<1>, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)> + // CHECK: %[[A1:.*]] = llvm.insertvalue %{{.*}}, %{{.*}}[10] : !llvm.struct<(i32, i32, i32, i32, i32, i32, i64, i64, i64, i64, i64)> %a1 = fly.atom.set_value(%atom, "imm_offset", %off) : (!fly.copy_atom, 0>, i64) -> !fly.copy_atom, 0> - // CHECK-DAG: %[[BASE:.*]] = llvm.extractvalue %[[A1]][1] - // CHECK-DAG: %[[IMM:.*]] = llvm.extractvalue %[[A1]][11] - // CHECK-DAG: %[[BI:.*]] = llvm.ptrtoint %[[BASE]] : !llvm.ptr<1> to i64 + // CHECK-DAG: %[[IMM:.*]] = llvm.extractvalue %[[A1]][10] + // CHECK-DAG: %[[BI:.*]] = llvm.ptrtoint %{{.*}} : !llvm.ptr<1> to i64 // CHECK: arith.addi %[[BI]], %[[IMM]] : i64 // CHECK: rocdl.tensor.load.to.lds fly.copy_atom_call(%a1, %src, %dst) : (!fly.copy_atom, 0>, !fly.memref, !fly.memref) -> () diff --git a/tests/unit/test_gfx1250_atoms.py b/tests/unit/test_gfx1250_atoms.py index 186e03b68..29d4428ce 100644 --- a/tests/unit/test_gfx1250_atoms.py +++ b/tests/unit/test_gfx1250_atoms.py @@ -7,8 +7,8 @@ and the 2D TDM copy atom. These only build the FlyROCDL atom *types* via the Python factories in -``flydsl.expr.rocdl.universal`` and check textual round-trip + verifier -behavior; no GPU is required. +``flydsl.expr.rocdl`` (WMMAScale in ``universal``, TDM in ``cdna5``) and check +textual round-trip + verifier behavior; no GPU is required. """ import pytest @@ -64,7 +64,7 @@ def test_wmma_scale_type_roundtrip(): def test_tdm2d_type_roundtrip(): with _ctx(), ir.Location.unknown(): from flydsl._mlir.dialects import fly_rocdl # noqa: F401 - from flydsl.expr.rocdl import universal as U + from flydsl.expr.rocdl import cdna5 as U t = U.TDM(2, 1) assert "gfx1250.tdm<" in str(t) From f3cc4140bd4651c25dff1e934f3def2677025fa1 Mon Sep 17 00:00:00 2001 From: Felix Li Date: Tue, 14 Jul 2026 12:22:09 +0000 Subject: [PATCH 2/3] chenge wmma scale --- python/flydsl/expr/rocdl/cdna5.py | 56 ++++++++++++++++++++++++++- python/flydsl/expr/rocdl/universal.py | 53 ------------------------- tests/unit/test_gfx1250_atoms.py | 6 +-- 3 files changed, 58 insertions(+), 57 deletions(-) diff --git a/python/flydsl/expr/rocdl/cdna5.py b/python/flydsl/expr/rocdl/cdna5.py index b61617524..b9a7f7a88 100644 --- a/python/flydsl/expr/rocdl/cdna5.py +++ b/python/flydsl/expr/rocdl/cdna5.py @@ -1,12 +1,66 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""gfx1250-specific ROCDL atom builders (N-D TDM whole-tile Global<->LDS copy).""" +"""gfx1250-specific ROCDL atom builders (MX-scaled WMMA + N-D TDM copy).""" +from ..._mlir import ir +from ..._mlir._mlir_libs._mlirDialectsFlyROCDL import MmaOpGFX1250_WMMAScaleType from ..._mlir.dialects.fly_rocdl import CopyOpGFX1250TDMType from ..typing import Int32, Int64, Tensor +def 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, +): + """Create a gfx1250 MX-scaled WMMA atom (E8M0 block scale) for the unified + f8/f6/f4 operand format. Per-operand scales are atom state (``scale_a`` / + ``scale_b``); ``opsel_a`` / ``opsel_b`` are forwarded as the intrinsic's + ``scaleAType`` / ``scaleBType`` operands (the scale-format / lane selector, + not an output opsel). ``mod_c`` (i16 C-operand modifier) and ``reuse_a`` / + ``reuse_b`` (operand-reuse scheduler hints) are forwarded to V_WMMA_SCALE. + + ``block_size`` selects the MX block size (elements per shared E8M0 scale): + ``32`` (default) uses V_WMMA_SCALE with i32 scale state; ``16`` uses + V_WMMA_SCALE16 with i64 scale state. + """ + ty_a = elem_ty_a.ir_type if hasattr(elem_ty_a, "ir_type") else elem_ty_a + if elem_ty_b is None: + ty_b = ty_a + else: + ty_b = elem_ty_b.ir_type if hasattr(elem_ty_b, "ir_type") else elem_ty_b + ty_acc = ( + ir.F32Type.get() + if elem_ty_acc is None + else (elem_ty_acc.ir_type if hasattr(elem_ty_acc, "ir_type") else elem_ty_acc) + ) + return MmaOpGFX1250_WMMAScaleType.get( + m, + n, + k, + ty_a, + ty_b, + ty_acc, + opsel_a=opsel_a, + opsel_b=opsel_b, + mod_c=mod_c, + reuse_a=reuse_a, + reuse_b=reuse_b, + block_size=block_size, + ) + + def TDM( rank, num_warps, diff --git a/python/flydsl/expr/rocdl/universal.py b/python/flydsl/expr/rocdl/universal.py index ba01ce04e..e55579588 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/python/flydsl/expr/rocdl/universal.py @@ -4,7 +4,6 @@ from ..._mlir import ir from ..._mlir._mlir_libs._mlirDialectsFlyROCDL import ( MmaOpGFX11_WMMAType, - MmaOpGFX1250_WMMAScaleType, MmaOpGFX1250_WMMAType, ) from ..._mlir.dialects import fly_rocdl @@ -136,58 +135,6 @@ def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): ) -def 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, -): - """Create a gfx1250 MX-scaled WMMA atom (E8M0 block scale) for the unified - f8/f6/f4 operand format. Per-operand scales are atom state (``scale_a`` / - ``scale_b``); ``opsel_a`` / ``opsel_b`` are forwarded as the intrinsic's - ``scaleAType`` / ``scaleBType`` operands (the scale-format / lane selector, - not an output opsel). ``mod_c`` (i16 C-operand modifier) and ``reuse_a`` / - ``reuse_b`` (operand-reuse scheduler hints) are forwarded to V_WMMA_SCALE. - - ``block_size`` selects the MX block size (elements per shared E8M0 scale): - ``32`` (default) uses V_WMMA_SCALE with i32 scale state; ``16`` uses - V_WMMA_SCALE16 with i64 scale state. - """ - ty_a = elem_ty_a.ir_type if hasattr(elem_ty_a, "ir_type") else elem_ty_a - if elem_ty_b is None: - ty_b = ty_a - else: - ty_b = elem_ty_b.ir_type if hasattr(elem_ty_b, "ir_type") else elem_ty_b - ty_acc = ( - ir.F32Type.get() - if elem_ty_acc is None - else (elem_ty_acc.ir_type if hasattr(elem_ty_acc, "ir_type") else elem_ty_acc) - ) - return MmaOpGFX1250_WMMAScaleType.get( - m, - n, - k, - ty_a, - ty_b, - ty_acc, - opsel_a=opsel_a, - opsel_b=opsel_b, - mod_c=mod_c, - reuse_a=reuse_a, - reuse_b=reuse_b, - block_size=block_size, - ) - - def make_buffer_ptr(ptr: Pointer, num_records_bytes=None): """Construct a new buffer-resource (``BufferDesc``) pointer from a global pointer, for hardware OOB-checked loads / stores. diff --git a/tests/unit/test_gfx1250_atoms.py b/tests/unit/test_gfx1250_atoms.py index 29d4428ce..2016065df 100644 --- a/tests/unit/test_gfx1250_atoms.py +++ b/tests/unit/test_gfx1250_atoms.py @@ -7,8 +7,8 @@ and the 2D TDM copy atom. These only build the FlyROCDL atom *types* via the Python factories in -``flydsl.expr.rocdl`` (WMMAScale in ``universal``, TDM in ``cdna5``) and check -textual round-trip + verifier behavior; no GPU is required. +``flydsl.expr.rocdl.cdna5`` (gfx1250 MX-scale WMMA + TDM) and check textual +round-trip + verifier behavior; no GPU is required. """ import pytest @@ -26,7 +26,7 @@ def _ctx(): def test_wmma_scale_type_roundtrip(): with _ctx(), ir.Location.unknown(): from flydsl._mlir.dialects import fly_rocdl # noqa: F401 (register dialect) - from flydsl.expr.rocdl import universal as U + from flydsl.expr.rocdl import cdna5 as U f8 = ir.Float8E4M3FNType.get() f4 = ir.Float4E2M1FNType.get() From 482d8681a37a845a4fbe975ea7ddae5471456ad4 Mon Sep 17 00:00:00 2001 From: coderfeli Date: Tue, 21 Jul 2026 09:54:03 +0000 Subject: [PATCH 3/3] [FlyROCDL] gfx1250 TDM: make per-dim extent/stride enum-private Drop the 8 Extent0..4 / Stride0..3 cases from the shared FlyROCDL AtomStateField enum. Their struct slots are unchanged; CopyOpGFX1250TDMType now resolves the "extent_i" / "stride_i" field names locally (tdmGeomSlot) instead of via symbolizeAtomStateField, so fly.atom.set_value still accepts them but they no longer pollute the cross-atom vocabulary. make_tdm_atom's signature/behavior and all call sites (incl. aiter) are unchanged; getConvertedType and the descriptor struct layout are identical (mask=0, extent 1..5, stride 6..9, imm_offset=10). WorkgroupMask/ImmOffset stay in the shared enum. Verified: tdm_gfx1250.mlir FileCheck + test_gfx1250_atoms.py round-trip pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- include/flydsl/Dialect/FlyROCDL/IR/Atom.td | 21 ++----- lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp | 66 +++++++++++----------- 2 files changed, 39 insertions(+), 48 deletions(-) diff --git a/include/flydsl/Dialect/FlyROCDL/IR/Atom.td b/include/flydsl/Dialect/FlyROCDL/IR/Atom.td index 10d0d44fb..1daac72b5 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/Atom.td +++ b/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 (the global base comes from the - // copy_atom_call operand pointer, not state): 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<"Extent0", 5, "extent_0">, - I32EnumAttrCase<"Extent1", 6, "extent_1">, - I32EnumAttrCase<"Extent2", 7, "extent_2">, - I32EnumAttrCase<"Extent3", 8, "extent_3">, - I32EnumAttrCase<"Extent4", 9, "extent_4">, - I32EnumAttrCase<"Stride0", 10, "stride_0">, - I32EnumAttrCase<"Stride1", 11, "stride_1">, - I32EnumAttrCase<"Stride2", 12, "stride_2">, - I32EnumAttrCase<"Stride3", 13, "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; diff --git a/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp b/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp index b742f82f0..676be7155 100644 --- a/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp +++ b/lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp @@ -105,44 +105,35 @@ constexpr int32_t kOuterStrideUnset = static_cast(0x80000000); // is assumed 1), imm_offset (i64)}. See CopyAtom.td for field semantics. static constexpr unsigned kMaxTdmRank = 5; +// Struct slots for the per-dim geometry: extent_i at kExtentSlot0+i (i32, i in +// 0..kMaxTdmRank-1), stride_i at kStrideSlot0+i (i64, i in 0..kMaxTdmRank-2). +// These are TDM-private (not shared AtomStateField cases); see tdmGeomSlot. +static constexpr unsigned kExtentSlot0 = 1; +static constexpr unsigned kStrideSlot0 = kExtentSlot0 + kMaxTdmRank; // 6 + std::optional CopyOpGFX1250TDMType::getFieldIndex(AtomStateField field) { switch (field) { case AtomStateField::WorkgroupMask: return 0; - case AtomStateField::Extent0: - return 1; - case AtomStateField::Extent1: - return 2; - case AtomStateField::Extent2: - return 3; - case AtomStateField::Extent3: - return 4; - case AtomStateField::Extent4: - return 5; - case AtomStateField::Stride0: - return 6; - case AtomStateField::Stride1: - return 7; - case AtomStateField::Stride2: - return 8; - case AtomStateField::Stride3: - return 9; // Byte offset added to the global base at lowering (i64, carry-safe). Lets a // K-loop advance the tile by bumping one scalar (imm_offset) instead of // re-deriving the base pointer. Default 0 folds away. case AtomStateField::ImmOffset: - return 10; + return kStrideSlot0 + (kMaxTdmRank - 1); // 10, last slot after the strides default: return std::nullopt; } } -// The i'th extent / stride state field (Extent0.., Stride0..). -static AtomStateField extentField(unsigned i) { - return static_cast(static_cast(AtomStateField::Extent0) + i); -} -static AtomStateField strideField(unsigned i) { - return static_cast(static_cast(AtomStateField::Stride0) + i); +// TDM-private per-dim geometry fields, resolved by name instead of via the shared +// AtomStateField enum: "extent_i" -> kExtentSlot0+i, "stride_i" -> kStrideSlot0+i. +static std::optional tdmGeomSlot(StringRef name) { + unsigned i; + if (name.consume_front("extent_") && !name.getAsInteger(10, i) && i < kMaxTdmRank) + return kExtentSlot0 + i; + if (name.consume_front("stride_") && !name.getAsInteger(10, i) && i < kMaxTdmRank - 1) + return kStrideSlot0 + i; + return std::nullopt; } Type CopyOpGFX1250TDMType::getConvertedType(MLIRContext *ctx) const { @@ -163,15 +154,18 @@ Value CopyOpGFX1250TDMType::getDefaultState(OpBuilder &builder, Location loc) co state = LLVM::InsertValueOp::create(builder, loc, state, v, ArrayRef{*getFieldIndex(f)}); }; + auto insertSlot = [&](unsigned slot, Value v) { + state = LLVM::InsertValueOp::create(builder, loc, state, v, ArrayRef{slot}); + }; Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32); Value noClamp = arith::ConstantIntOp::create(builder, loc, 0x7FFFFFFF, 32); Value strideUnset = arith::ConstantIntOp::create(builder, loc, kOuterStrideUnset, 64); Value zero64 = arith::ConstantIntOp::create(builder, loc, 0, 64); insert(AtomStateField::WorkgroupMask, zero); for (unsigned i = 0; i < kMaxTdmRank; ++i) - insert(extentField(i), noClamp); + insertSlot(kExtentSlot0 + i, noClamp); for (unsigned i = 0; i < kMaxTdmRank - 1; ++i) - insert(strideField(i), strideUnset); + insertSlot(kStrideSlot0 + i, strideUnset); insert(AtomStateField::ImmOffset, zero64); return state; } @@ -181,10 +175,13 @@ Value CopyOpGFX1250TDMType::setAtomState(OpBuilder &builder, Location loc, Value auto fieldStr = dyn_cast(fieldAttr); if (!fieldStr) return nullptr; - auto field = symbolizeAtomStateField(fieldStr.getValue()); - if (!field) - return nullptr; - auto idx = getFieldIndex(*field); + // Per-dim extent/stride are TDM-private (resolved by name); everything else goes + // through the shared AtomStateField enum. + std::optional idx = tdmGeomSlot(fieldStr.getValue()); + if (!idx) { + if (auto field = symbolizeAtomStateField(fieldStr.getValue())) + idx = getFieldIndex(*field); + } if (!idx) return nullptr; return LLVM::InsertValueOp::create(builder, loc, atomStruct, fieldValue, ArrayRef{*idx}); @@ -305,15 +302,18 @@ LogicalResult CopyOpGFX1250TDMType::emitAtomCall(OpBuilder &builder, Location lo return LLVM::ExtractValueOp::create(builder, loc, atomVal, ArrayRef{*getFieldIndex(f)}); }; + auto slotField = [&](unsigned slot) { + return LLVM::ExtractValueOp::create(builder, loc, atomVal, ArrayRef{slot}); + }; // Per-dim (tensor order) tensor extent (i32) and stride in elements (i64). The // innermost stride (dim rank-1) is assumed 1; stride_i covers dims 0..rank-2, // falling back to the static layout stride when the state slot is left unset. SmallVector extent(rank), strideElems(rank); for (int32_t i = 0; i < rank; ++i) - extent[i] = stateField(extentField(i)); + extent[i] = slotField(kExtentSlot0 + i); for (int32_t i = 0; i < rank - 1; ++i) { - Value st = stateField(strideField(i)); + Value st = slotField(kStrideSlot0 + i); if (hasStaticStride) { int64_t s = layout.getStride().at(i).getLeafAsInt().getValue(); Value unset =