From d656f52d1757123f387a321880aa948bc6716c15 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 01:43:11 +0800 Subject: [PATCH 01/19] refactor(pattern): retire legacy matching helpers --- docs/spec/core-ir.md | 16 +++++----- src/tilefoundry/ir/pattern/__init__.py | 6 ---- src/tilefoundry/ir/pattern/match.py | 35 ++-------------------- src/tilefoundry/ir/pattern/pattern.py | 17 +++++------ src/tilefoundry/ir/types/int_tuple.py | 11 ++++++- src/tilefoundry/ir/types/layout_algebra.py | 10 +++++++ 6 files changed, 40 insertions(+), 55 deletions(-) diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 288de64f..1a77c18c 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -620,18 +620,20 @@ The implementation is split by responsibility under `ir/pattern/`: `LayoutPattern`, `SwizzlePattern`, `ComposedLayoutPattern`, `MeshPattern`, `ShardLayoutPattern`, `ScalarPattern`, `TensorPattern`, and `WildcardPattern`. It also owns the `Scalar` and `Tensor` singletons. -- `match.py` owns matches, captures, symbolic resolution, layout-frame reading, - and the shared description helpers. +- `match.py` owns matches, captures, symbolic resolution, and the shared + description helpers. An unstated (`None`) pattern field admits any value. - `constraint.py` owns cross-operand `Constraint`, `DistinctConstraint`, `SameConstraint`, and `SameModesConstraint` values. - `utils.py` owns exact-layout construction plus specialization naming and dimension lookup. -`LayoutPattern` checks `forward` and `injective` over the whole flattened -arrangement by default. With `per_mode=True`, it checks each top-level mode -independently; `MeshPattern` requires this explicit form because each mesh -level uses its own numbering space. `MeshPattern` never changes the supplied -pattern implicitly. +`LayoutPattern` matches only a bare `Layout`, preserves its nested mode +structure, and checks `forward` and `injective` over the whole flattened +arrangement by default. A sliced layout must be stated explicitly with +`ComposedLayoutPattern`; callers that accept both forms use `OrPattern`. +With `per_mode=True`, `LayoutPattern` checks each top-level mode independently; +`MeshPattern` requires this explicit form because each mesh level uses its own +numbering space. `MeshPattern` never changes the supplied pattern implicitly. Two consumer surfaces: diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py index a3344174..5f4858f8 100644 --- a/src/tilefoundry/ir/pattern/__init__.py +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -13,12 +13,9 @@ OPAQUE, UNNAMED_PLACE, Match, - affine_frame, alternatives_of, between_rules, evaluated, - fits, - grouping, is_symbolic, matched, refusals_between, @@ -98,15 +95,12 @@ "WildcardPattern", "WHOLE_BYTES", "_mangle_variant_name", - "affine_frame", "affine_part", "alternatives_of", "any_threads", "arrangement_pattern", "between_rules", "evaluated", - "fits", - "grouping", "dtype_place", "is_symbolic", "locate_dim_var", diff --git a/src/tilefoundry/ir/pattern/match.py b/src/tilefoundry/ir/pattern/match.py index 27c18f83..c0c5b567 100644 --- a/src/tilefoundry/ir/pattern/match.py +++ b/src/tilefoundry/ir/pattern/match.py @@ -5,7 +5,6 @@ from dataclasses import dataclass, field from tilefoundry.ir.clause.layout import is_layout_wildcard -from tilefoundry.ir.types import ComposedLayout from tilefoundry.ir.types.dim import DimFloorDiv, DimMul, DimVar, is_dim_op_call from tilefoundry.ir.types.substitute import DimSubstitutionError, substitute_shape_dim @@ -129,6 +128,8 @@ def alternatives_of(pattern, bindings=()) -> tuple: def matched(pattern, subject, captures=None) -> Match | None: """Match a nested pattern, symbolic dimension, wildcard, or fixed value.""" + if pattern is None: + return Match(dict(captures or {})) held = Match(dict(captures or {})) if isinstance(pattern, DimVar): if pattern.name in held.captures: @@ -141,25 +142,7 @@ def matched(pattern, subject, captures=None) -> Match | None: return held if found is not None and found == subject else None if isinstance(pattern, _pattern_type()): return pattern.match(subject, held.captures) - return held if fits(pattern, subject) else None - - -def fits(wanted, held) -> bool: - """Whether one pattern position admits one concrete value.""" - if wanted is None: - return True - if isinstance(wanted, _pattern_type()) or is_symbolic(wanted): - return matched(wanted, held) is not None - if is_layout_wildcard(wanted): - return type(held) is int - return wanted == held - - -def grouping(modes): - """The tuple nesting of *modes*, with the concrete modes erased.""" - if isinstance(modes, tuple): - return tuple(grouping(mode) for mode in modes) - return None + return held if pattern == subject else None def relations_of(values) -> tuple[str, ...]: @@ -180,15 +163,6 @@ def refusals_between(op_type, operands: dict) -> tuple[str, ...]: ) -def affine_frame(layout) -> tuple | None: - """Read a bare layout or a composition through no transform as offset + layout.""" - if not isinstance(layout, ComposedLayout): - return 0, layout - if layout.inner is not None: - return None - return layout.offset, layout.outer - - __all__ = [ "ABSENT", "ARRANGEMENT", @@ -196,12 +170,9 @@ def affine_frame(layout) -> tuple | None: "OPAQUE", "UNNAMED_PLACE", "_named", - "affine_frame", "alternatives_of", "between_rules", "evaluated", - "fits", - "grouping", "is_symbolic", "matched", "refusals_between", diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index 2bf4aa27..62e60a8e 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -15,8 +15,9 @@ TensorType, make_mesh, ) +from tilefoundry.ir.types.int_tuple import congruent from tilefoundry.ir.types.layout import flatten -from tilefoundry.ir.types.layout_algebra import is_inverse_projectable +from tilefoundry.ir.types.layout_algebra import frame_of, is_inverse_projectable from tilefoundry.ir.types.mesh import separate from .match import ( @@ -25,10 +26,8 @@ UNNAMED_PLACE, Match, _named, - affine_frame, alternatives_of, evaluated, - grouping, matched, relations_of, resolved, @@ -426,13 +425,13 @@ def positions(self) -> tuple: return (*flatten(self.shape), *flatten(self.strides)) def match(self, subject, captures=None): - framed = affine_frame(subject) - layout = None if framed is None else framed[1] + layout = subject if not isinstance(layout, Layout) or layout.strides is None: return None - if grouping(tuple(layout.shape)) != grouping(tuple(self.shape)) or grouping( - tuple(layout.strides) - ) != grouping(tuple(self.strides)): + if not ( + congruent(layout.shape, self.shape) + and congruent(layout.strides, self.strides) + ): return None extents = tuple(flatten(layout.shape)) strides = tuple(flatten(layout.strides)) @@ -710,7 +709,7 @@ def match(self, subject, captures=None): or any(not isinstance(attr, Broadcast) for attr in subject.attrs[:extra]) ): return None - framed = affine_frame(subject.mesh.layout) + framed = frame_of(subject.mesh.layout) if framed is None: return None frame = framed[1] diff --git a/src/tilefoundry/ir/types/int_tuple.py b/src/tilefoundry/ir/types/int_tuple.py index ff81f27f..e9d81639 100644 --- a/src/tilefoundry/ir/types/int_tuple.py +++ b/src/tilefoundry/ir/types/int_tuple.py @@ -23,6 +23,15 @@ def flatten(t: object) -> tuple[object, ...]: return tuple(value for item in t for value in flatten(item)) +def congruent(a, b) -> bool: + """Whether two int tuples have the same nesting, ignoring their leaves.""" + if isinstance(a, tuple) != isinstance(b, tuple): + return False + if not isinstance(a, tuple): + return True + return len(a) == len(b) and all(congruent(x, y) for x, y in zip(a, b)) + + def product(t) -> "ShapeDim": from .mesh import Topology # noqa: PLC0415 @@ -74,4 +83,4 @@ def unflatten(flat: tuple, profile) -> tuple: return nested -__all__ = ["IntTuple", "flatten", "product", "repeat_like", "unflatten"] +__all__ = ["IntTuple", "congruent", "flatten", "product", "repeat_like", "unflatten"] diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index 82303e35..c1c18b1e 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -123,6 +123,15 @@ def coalesce(layout: Union[Layout, ComposedLayout]): return Layout(shape=tuple(result_shape), strides=tuple(result_stride)) +def frame_of(layout: Union[Layout, ComposedLayout]) -> tuple[int, Layout] | None: + """Read a bare layout or an affine composition as offset + outer layout.""" + if isinstance(layout, Layout): + return 0, layout + if layout.inner is not None or not isinstance(layout.outer, Layout): + return None + return layout.offset, layout.outer + + def complement(layout: Layout, max_idx: int = 1) -> Layout: """CuTe ``complement``: the modes that fill the gaps below ``max_idx``.""" result_shape: list[int] = [] @@ -404,6 +413,7 @@ def contains(scope: ComposedLayout, t: int) -> bool: "cosize", "apply", "coalesce", + "frame_of", "complement", "is_inverse_projectable", "right_inverse", From 418c6aaf151afaee15c90f80bf4e0100eea0b929 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 02:16:28 +0800 Subject: [PATCH 02/19] feat(tir): add declarative CUDA MMA atoms --- docs/spec/code-organization.md | 7 +- docs/spec/codegen.md | 2 +- docs/spec/core-ir.md | 2 +- docs/spec/runtime.md | 32 +- docs/spec/tir.md | 208 ++++--------- docs/spec/visitor-mutator.md | 2 +- docs/spec/visitor-registry.md | 2 +- include/tilefoundry/runtime/cuda/ops/mma.cuh | 15 +- .../runtime/cuda/ops/mma/mma_impl.h | 113 +------ src/tilefoundry/codegen/cuda/tir/nn/mma.py | 21 +- src/tilefoundry/dsl/T/_platforms.py | 27 +- src/tilefoundry/dsl/_stub_gen.py | 67 +++- src/tilefoundry/inspection/printer_base.py | 3 +- src/tilefoundry/ir/tir/cuda/nn/mma.py | 236 ++++++-------- src/tilefoundry/ir/tir/cuda/nn/mma_atom.py | 287 ++++++++++++++++-- src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py | 59 ++++ src/tilefoundry/ir/tir/cuda/nn/wgmma.py | 245 +++++++++++++++ src/tilefoundry/ir/tir/verify.py | 14 +- tests/fixtures/tir/mma.py | 7 +- tests/installed/smoke_spec.py | 21 -- tests/ir/test_visitor.py | 5 +- tests/ir/types/test_mma_fragment_layouts.py | 21 +- tests/ops/tir/cuda/test_mma.py | 59 +++- tests/runtime/cuda/static_asserts.cu | 26 +- tests/runtime/test_cuda_static_asserts.py | 3 +- 25 files changed, 907 insertions(+), 577 deletions(-) create mode 100644 src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py create mode 100644 src/tilefoundry/ir/tir/cuda/nn/wgmma.py diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 997b35bb..7b5d0786 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -135,9 +135,10 @@ are not IR classes, so they go through Rule 1a. nodes are target-neutral. A node or descriptor that is specific to one compilation target nests as `ir/{dialect}/{target}/{category}/.py`; target-neutral abstractions stay at `ir/{dialect}/{category}/`. For -example the whole MMA surface is target-owned — the `Mma` op, the -`MmaOpSpec` / `MmaAtom` descriptors, the CUDA SM80 instruction spec, and its -fragment layouts all live under `ir/tir/cuda/nn/` (`mma.py` + `mma_atom.py`). +example the whole MMA surface is target-owned — `mma.py` defines the `TiledMma` +op, `mma_atom.py` defines `MmaAtom` / `AtomPattern`, and `sm80_mma.py` / +`wgmma.py` define the CUDA instruction declarations. All four live under +`ir/tir/cuda/nn/`. The backend-bound construction stays in TIR: HIR is the checking reference side, and carrying the instruction name in that reference would make two GPU targets require different HIR references. (`codegen/` and `runtime/` are diff --git a/docs/spec/codegen.md b/docs/spec/codegen.md index 1b2b3a09..66282a0d 100644 --- a/docs/spec/codegen.md +++ b/docs/spec/codegen.md @@ -207,7 +207,7 @@ tuples because it does not materialize an aggregate. The codegen context records the structural tuple by its fresh SSA `Var` identity so consumers can recover its elements without target-side storage. -Effect Ops (`Copy`, `Fill`, `Mma`, `tir.nn.*`, ...) appear in Stmt +Effect Ops (`Copy`, `Fill`, `TiledMma`, `tir.nn.*`, ...) appear in Stmt position as `Evaluate(op, args)` rather than as Stmt subclasses. The walker matches `Evaluate` and dispatches on `type(callable)` through the handler registry. Handlers stay small; the runtime function they diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 1a77c18c..1d0fd80e 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -588,7 +588,7 @@ from `op.params()`. An Op is **value-form** when its `Call` produces an observable result the IR consumes — `Call.type` is then `TensorType` or `TupleType`. An Op is **effect-form** when it performs an in-place -effect (e.g. `tir.memory.Copy` / `tir.cuda.nn.Mma`) and produces no +effect (e.g. `tir.memory.Copy` / `tir.cuda.nn.TiledMma`) and produces no readable value (`UnitType`, [types §6](./types.md#6-unittype)); in Stmt position it appears as `Evaluate(op, args)` ([tir §1.4](./tir.md#14-evaluate)). diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index 20346a54..87ab2635 100644 --- a/docs/spec/runtime.md +++ b/docs/spec/runtime.md @@ -1194,7 +1194,7 @@ the one time it was written down as the rule it argued a dependency chain into | `copy` / `copy_async` | both operands' shard layouts: shape, strides, share and move width | | `reduce` | the axes the destination broadcasts that the source splits | | `dot` | the axes the operands' meshes contract | -| `mma` | rank-2 static layouts are a tile; the warp count is the accumulator's mesh | +| `mma` | one atom's per-lane operand fragments | | `rmsnorm` | the row dependency chain and the destination's shard layout | | `sync` | the mesh's scope, base and count | | `tma_copy` | both shard layouts, asserted: one contiguous run each, whole tiles, matching element types | @@ -1404,32 +1404,14 @@ __device__ void mma(TA const &a, TB const &b, TC &c); **`mma`.** -`c += a @ b`, one entry, with the tier read off the operand layouts: rank-2 -static shard layouts on `a` and `b` are a tile and the entry loops the atom over -it; anything else is a lane's already-gathered fragment and takes the single -instruction. Codegen emits this one call either way. +`c += a @ b` for one atom. The operands are the calling lane's already-gathered +fragments, and the entry issues one instruction. - constraints: - - The tile tier reads `a` as `(M, K)` and `b` as `(N, K)`. **There is no - transpose flag.** Whether the buffer behind `b` is k-major or n-major is a - stride in its layout, and the indexing picks that up, so the same call reads - both. - - `M` and `K` must be whole multiples of the atom's `16` and `16`, and every - warp must receive a whole number of `N` atoms of `8`. Violations are - `static_assert`s, not run-time checks. - - Warps split `N`. The warp count comes from the accumulator's mesh — that is - what `c` being a `ShardTensor` is for — not from `blockDim`. - - The accumulator's engine is the lane's own registers, which is what - `local_tensor` ([§2.4.1](#241-tensor_viewshard_tensorcuh)) hands back for register storage, while its - `ShardLayout` states which entries of the tile those registers are: the - fragment map is warp-split over `N` and lane-split within each atom, and - saying so is the layout's job, not an accessor's. **The runtime publishes no - fragment-coordinate function and no accumulator constructor.** A caller that - needs the map writes it as modes and attrs, the way it writes any other - layout, and the tile it moves the fragment to or from is then the same map - over a buffer — so rescaling a row of the accumulator or storing it out is - one `ops::elementwise` between two shards of one layout, with no fragment - index at the call site. + - `a`, `b`, and `c` contain exactly the atom's `(8, 4, 4)` values per lane. + Other shapes are compile-time errors. + - The fragment map is stated by each operand's `ShardLayout`; the runtime + publishes no fragment-coordinate function and no accumulator constructor. - Today's atom is `SM80_16x8x16_F32BF16BF16F32_TN`: bf16 operands, f32 accumulate. Another instruction is another atom under the same entry, not another entry. diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 2e62fe9e..86a9a5f3 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -594,28 +594,30 @@ class Cast(Op): #### NN Ops (`tir.nn.*`) -##### Mma +##### TiledMma ```python -class Mma(Op): +class TiledMma(Op): """Effect form; matrix-multiply-accumulate ``acc += lhs @ rhs``. Attributes: acc: input; accumulator fragment. lhs: input; left-hand operand fragment. rhs: input; right-hand operand fragment. - atom: attribute; optional compile-time ``MmaAtom``, absent ⇒ bare-Mma - per-target path. + atom: attribute; required compile-time ``MmaAtom`` declaration. + scope: attribute; optional warp-aligned thread scope. """ acc: Tensor lhs: Tensor rhs: Tensor - atom: MmaAtom | None = None + atom: MmaAtom + scope: Mesh | None = None ``` - constraints: - matrix-multiply-accumulate `acc += lhs @ rhs`; per-target PTX lowering lives in [target](./target.md), the atom calling convention in [§2.3](#23-tir-ops). + - `acc` declares `READ | WRITE`; `lhs` and `rhs` declare `READ`. ##### ReLU ```python @@ -874,162 +876,74 @@ block_y, block_z, *forwarded_args)`: non-grid/block launch configuration. A `cluster` / `stream` / `attrs` value the active CUDA target does not support MUST be rejected in target lowering. -#### MMA atom and the hand-written calling convention - -A hand-written kernel issues an MMA through an explicit **atom** — a -realized instruction descriptor — instead of the bare `Mma` op whose -fragment layouts the per-target lowering chooses -([hir §1.3](./hir.md#13-op), [passes](./passes.md)). An MMA atom fixes a -concrete hardware instruction, so the whole MMA surface is **target-owned**: -the `Mma` op and the `MmaOpSpec` / `MmaAtom` descriptors -(`tilefoundry.ir.tir.cuda.nn`, mirroring the CuTe `MMA_Op` → `MMA_Atom` -layering), the concrete instructions, and their fragment layouts all live -under `tilefoundry.ir.tir.cuda.nn.mma` / `mma_atom`, following IR's dialect-first -layout `ir/{dialect}/{target}/{category}`. - -##### `MmaOpSpec` +#### Declarative MMA atoms and `T.tiled_mma` -A named, fully-specified MMA instruction (the CuTe `MMA_Op` analog). - -```python -class MmaOpSpec: - name: str # uniquely identifies the instruction; the other fields mirror it - shape_mnk: tuple[int, int, int] # the instruction's static (M, N, K) - dtype_a: DType # lhs operand element type - dtype_b: DType # rhs operand element type - dtype_c: DType # accumulator element type - operand_layout: str # source operand order string (e.g. "TN") -``` - -- constraints: - - a fully-specified MMA instruction descriptor carrying no fragment-layout - knowledge. Per-field rules below. - -###### `name` - -- MUST uniquely identify the instruction. dtype / shape / source layout - are fixed by it; the remaining fields mirror the name so verify and - codegen do not re-parse the string. - -###### `shape_mnk` - -- MUST be the instruction's `(M, N, K)` tuple; every entry MUST be a - static int. - -###### `dtype_a` - -- MUST be the `lhs` operand element type (`DType`). - -###### `dtype_b` - -- MUST be the `rhs` operand element type (`DType`); it MAY differ from - `dtype_a`. - -###### `dtype_c` - -- MUST be the accumulator element type (`DType`); it MAY differ from the - operand types (e.g. `f32` accumulation over `bf16` operands). - -###### `operand_layout` - -- MUST encode the source operand order as a string, e.g. `"TN"` (A - row-major, B col-major). An `MmaOpSpec` MUST NOT carry fragment-layout - knowledge. - -##### `MmaAtom` - -The realized atom for an `op` (the CuTe `MMA_Atom` analog), built by -`T.cuda.mma.atom(op=...)` -([parser §2](./parser.md#2-syntax-and-rules)). +An MMA instruction is a target-owned `MmaAtom` declaration. The declaration +class states its authored parameters, required physical scope, target +capability, and the `TensorPattern` read for each `A`, `B`, and `C` role. An +instance binds the parameters for one call; it does not carry a second copy of +concrete fragment layouts that could drift from those patterns. ```python class MmaAtom: - op: MmaOpSpec # the MmaOpSpec this atom realizes - A: ShardLayout # lhs fragment ShardLayout contract - B: ShardLayout # rhs fragment ShardLayout contract - C: ShardLayout # accumulator fragment ShardLayout contract - required_scope: Mesh # the thread-participation contract, carried as its own Mesh + namespace: str + scope: Mesh + capability: str + A: TensorPattern | SwitchPattern + B: TensorPattern | SwitchPattern + C: TensorPattern | SwitchPattern + parameters: tuple[ParamDef, ...] + bindings: dict[str, object] + mesh: Mesh | None + + def role(self, role: str) -> TensorPattern: ... + def scope_pattern(self) -> MeshPattern: ... ``` - constraints: - - the realized atom for an `op`; fragment layouts are returned as-is and not - rebound onto the caller's mesh. Per-field rules below. - -###### `op` - -- MUST be the `MmaOpSpec` this atom realizes. - -###### `A` - -- MUST be the `lhs` operand fragment `ShardLayout` contract — the - lane→value layout the instruction reads. It MUST be returned **as-is** - at a use site and MUST NOT be rebound onto the caller's mesh. - -###### `B` - -- MUST be the `rhs` operand fragment `ShardLayout` contract; the same - as-is / no-rebind rule as `A` applies. - -###### `C` - -- MUST be the accumulator fragment `ShardLayout` contract; the same - as-is / no-rebind rule as `A` applies. - -###### `required_scope` - -- MUST be the thread-participation contract the atom needs, carried as - its own `Mesh` (for the SM80 `16x8x16` instruction, 32 lanes arranged - as a `(4, 8)` thread mesh). It MUST NOT be the caller's mesh; the - caller's enclosing scope MUST be checked against it at verify (below). + - `parameters` MUST preserve declaration order. Construction MUST reject an + unknown binding and a value refused by its `ParamDef.pattern`; an omitted + parameter MUST take the value implied by earlier bindings or its declared + default, and otherwise construction MUST fail. + - `role("A")`, `role("B")`, and `role("C")` MUST resolve the declaration's + role pattern under the instance bindings. The logical TIR orientation is + always A `(M,K)`, B `(K,N)`, C `(M,N)`; each role pattern separately states + the fragment's physical arrangement. + - `scope_pattern()` MUST require the declaration's exact participant count + at an aligned offset. `mesh`, when present, binds the atom to one concrete + frame and MUST match the active frame at verify. + - `capability` MUST be present in the active CUDA architecture's + `instruction_capabilities`. + +The public declarations are `T.cuda.sm80.Mma()` (BF16 `16x8x16`, F32 +accumulator, register A/B/C over one warp) and +`T.cuda.sm90.Wgmma(n=..., form=..., a_major=..., mesh=...)` (BF16 +`64 x n x 16` over one warpgroup). `Form` and `Major` live beside `Wgmma` +under `T.cuda.sm90`. ##### Calling convention -Load, compute, and store are **three separate** effect statements under -an enclosing `MeshScope` ([§1.2](#12-structural-stmts-tirstmts)); `T.mma` is -verify-only and MUST NOT fuse the loads or the store. +Load, compute, and store are separate effect statements under an enclosing +`MeshScope` ([§1.2](#12-structural-stmts-tirstmts)). `T.tiled_mma` has exactly +three input operands plus one required compile-time `atom` attribute: ```python -# example -atom = T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN) -with Mesh((Topology("thread", 32),), Layout(shape=(4, 8), strides=(1, 4))) as warp: - a_frag = T.alloc_tensor(TensorType(..., layout=atom.A, storage=rmem)) - acc = T.alloc_tensor(TensorType(..., layout=atom.C, storage=rmem)) - T.copy(T.tensor_view(T.ptr_of(a), layout=atom.A), a_frag) # load - T.fill(acc, 0.0) - T.mma(acc, a_frag, b_frag, atom=atom) # compute - T.copy(acc, T.tensor_view(T.ptr_of(c), layout=atom.C)) # store +T.tiled_mma(acc, lhs, rhs, atom=T.cuda.sm80.Mma()) ``` -- The author allocates each register fragment with the matching - `atom.A/B/C` layout and fills it with its own `T.copy`. The - accumulator is initialised with `Fill` and then read-modify-written. -- `atom` is a compile-time attribute on the `Mma` Op - ([parser §2](./parser.md#2-syntax-and-rules)), not a runtime - operand. When absent, lowering takes the bare-`Mma` per-target path. - -##### Verify - -A `T.mma` carrying an `atom` MUST satisfy: - -- **operand contracts**: `acc.layout == atom.C`, `lhs.layout == atom.A`, - `rhs.layout == atom.B`. -- **scope**: some mesh on the active TIR `MeshScope` traversal cache provides the - atom's required thread scope — - `mesh_scope_matches_required_scope(mesh, atom.required_scope)`. The - match is identity- and name-independent (mesh object identity, the - binding-var name, and axis names are not compared); it holds iff: - - the two meshes share the same program topology level — a `cta` - scope is never a `thread` / warp scope, even when its layout carries - the same shape; - - both topology domains (the product of the topology extents) are - statically known; - - each mesh is self-consistent (`topology domain == layout extent`), - and the enclosing mesh is inverse-projectable; - - the thread-value decomposition matches **exactly** — same layout - shape and strides. A flat lane layout cannot host the atom's - multi-axis fragment `Split` and is rejected. - -Per-target PTX emission dispatches on the atom ([target](./target.md)). +- `acc` MUST match `atom.role("C")` and is read-write. +- `lhs` MUST match `atom.role("A")` and is read-only. +- `rhs` MUST match `atom.role("B")` and is read-only. +- A, B, and C shapes MUST be `(M,K)`, `(K,N)`, and `(M,N)` respectively; + operands MUST use the atom's declared dtypes and fragment arrangements. +- The active physical mesh MUST satisfy `atom.scope_pattern()`. If the atom + carries `mesh=...`, its affine offset and ordered lanes MUST equal the active + frame. + +There is one TIR MMA op: `T.tiled_mma`. There is no optional-atom or bare-MMA +path. Per-target emission dispatches on the atom ([target](./target.md)); the +current CUDA emitter accepts the SM80 declaration and rejects WGMMA because +this stage provides no WGMMA runtime emitter. #### Async copy Ops (`tir.async.*`) diff --git a/docs/spec/visitor-mutator.md b/docs/spec/visitor-mutator.md index 624d9acc..43b21b78 100644 --- a/docs/spec/visitor-mutator.md +++ b/docs/spec/visitor-mutator.md @@ -265,7 +265,7 @@ directly to rebuild Stmts; they do not reuse `StmtExprMutator`. ## 7. Visitor entry forms for `Evaluate` -The TIR effect-form Ops (e.g. `Copy` / `Fill` / `Mma` / `ReLU` / +The TIR effect-form Ops (e.g. `Copy` / `Fill` / `TiledMma` / `ReLU` / `RMSNorm` / `Reduce`) are `Op` subclasses, not `Stmt` subclasses; in Stmt position they appear as `Evaluate(callable=op, args)` so the invocation can sit in `Sequential` body position. Passes and visitors diff --git a/docs/spec/visitor-registry.md b/docs/spec/visitor-registry.md index 8598a11d..5ea6e082 100644 --- a/docs/spec/visitor-registry.md +++ b/docs/spec/visitor-registry.md @@ -425,7 +425,7 @@ def register_verify_stmt(cls: type): ... # decorator: register a verify h through `ctx.error(node, msg)`, which raises `VerifyError`. **`Evaluate(op, args)` dispatch.** TIR effect-form Ops -(`Copy` / `Fill` / `Mma` / `ReLU` / `RMSNorm` / `Reduce`) appear in +(`Copy` / `Fill` / `TiledMma` / `ReLU` / `RMSNorm` / `Reduce`) appear in Stmt position as `Evaluate(callable=op, args)`. The verify path keys on the Op class, not on `Evaluate` itself: `register_verify_stmt` takes the **Op class**, and `VerifyVisitor.generic_visit` — diff --git a/include/tilefoundry/runtime/cuda/ops/mma.cuh b/include/tilefoundry/runtime/cuda/ops/mma.cuh index 309cbe72..6893869d 100644 --- a/include/tilefoundry/runtime/cuda/ops/mma.cuh +++ b/include/tilefoundry/runtime/cuda/ops/mma.cuh @@ -4,16 +4,11 @@ #include "mma/mma_impl.h" -/// ``c += a @ b``, one entry, the tier read off the operand layouts. +/// ``c += a @ b`` on one atom's per-lane fragments. template __device__ void mma(TA const &a, TB const &b, TC &c) { - if constexpr (mma_impl::tile_shaped_v) { - mma_impl::Tile{}(a, b, c); - } else if constexpr (mma_impl::atom_shaped_v) { - mma_impl::Atom{}(a, b, c); - } else { - static_assert(dependent_false_v, - "ops::mma: the operands are neither a rank-2 static tile " - "nor the atom's own (8, 4, 4) lane fragments"); - } + static_assert( + mma_impl::atom_shaped_v, + "ops::mma: operands must be the atom's (8, 4, 4) lane fragments"); + mma_impl::Atom{}(a, b, c); } diff --git a/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h b/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h index 89a6070a..343ed54f 100644 --- a/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h +++ b/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h @@ -14,26 +14,8 @@ template __device__ uint16_t as_u16(T const &x) { return out; } -/// Geometry of the supported atom, read from CuTe's canonical trait. +/// The supported atom. using AtomOp = cute::SM80_16x8x16_F32BF16BF16F32_TN; -using AtomTraits = cute::MMA_Traits; -using AtomCLayout = typename AtomTraits::CLayout; - -struct AtomGeometry { - using ShapeMNK = typename AtomTraits::Shape_MNK; - static constexpr int kM = int(cute::get<0>(ShapeMNK{})); - static constexpr int kN = int(cute::get<1>(ShapeMNK{})); - static constexpr int kK = int(cute::get<2>(ShapeMNK{})); - static constexpr int kAccPerLane = - int(cute::size(AtomCLayout{})) / kWarpSize; - - __device__ static int c_coord(int lane, int value = 0) { - return int(AtomCLayout{}(cute::make_coord(lane, value))); - } - - __device__ static int row(int lane) { return c_coord(lane) % kM; } - __device__ static int col(int lane) { return c_coord(lane) / kM; } -}; } @@ -73,99 +55,6 @@ struct Atom { } }; -/// A rank-2 static layout is a tile; anything else is a gathered fragment. -template -inline constexpr bool tile_v = [] { - using L = typename cute::remove_cvref_t::layout_type; - return cute::is_static::value && decltype(cute::rank(L{}))::value == 2; -}(); - -/// How many threads the accumulator's mesh spreads the tile over. -template CUTE_HOST_DEVICE constexpr int acc_threads() { - return tilefoundry::shard_mesh_instances(); -} - -/// ``acc += a @ b`` over a whole tile, the atom looped by the tile's shape. -struct Tile { - template - __device__ void operator()(TA const &a, TB const &b, TC &c) const { - using Geo = mma_detail::AtomGeometry; - auto av = tilefoundry::local_tensor(a); - auto bv = tilefoundry::local_tensor(b); - auto &&cv = tilefoundry::local_tensor(c); - using a_elem = cute::remove_cvref_t; - - constexpr int threads = acc_threads(); - static_assert(threads >= kWarpSize && threads % kWarpSize == 0, - "ops::mma (tile tier): the accumulator's mesh must be a " - "whole number of warps"); - constexpr int warps = threads / kWarpSize; - constexpr int M = int(cute::size<0>( - typename cute::remove_cvref_t::shard_layout_type::layout{})); - constexpr int K = int(cute::size<1>( - typename cute::remove_cvref_t::shard_layout_type::layout{})); - constexpr int N = int(cute::size<0>( - typename cute::remove_cvref_t::shard_layout_type::layout{})); - static_assert(M % Geo::kM == 0 && K % Geo::kK == 0, - "the A tile must be a whole number of atoms"); - static_assert(N % (Geo::kN * warps) == 0, - "every warp must get a whole number of N atoms"); - - constexpr int m_atoms = M / Geo::kM; - constexpr int n_atoms = N / (Geo::kN * warps); - - const unsigned tid = unsigned( - tilefoundry::program_id()); - const int lane = int(tid & unsigned(kWarpSize - 1)); - const int warp = int((tid / unsigned(kWarpSize)) % unsigned(warps)); - const int row = Geo::row(lane); - const int col = Geo::col(lane); - - a_elem af[8]; - a_elem bf[4]; - float acc[Geo::kAccPerLane]; - - CUTE_UNROLL - for (int mi = 0; mi < m_atoms; ++mi) { - const int m0 = mi * Geo::kM + row; - CUTE_UNROLL - for (int ni = 0; ni < n_atoms; ++ni) { - const int n0 = (warp * n_atoms + ni) * Geo::kN + row; - const int base = (mi * n_atoms + ni) * Geo::kAccPerLane; - CUTE_UNROLL - for (int v = 0; v < Geo::kAccPerLane; ++v) - acc[v] = cv(base + v); - - for (int k0 = 0; k0 < K; k0 += Geo::kK) { - const int kc = k0 + col; - CUTE_UNROLL - for (int h = 0; h < 2; ++h) { - af[0 + h] = av(m0, kc + h); - af[2 + h] = av(m0, kc + 8 + h); - af[4 + h] = av(m0 + 8, kc + h); - af[6 + h] = av(m0 + 8, kc + 8 + h); - bf[0 + h] = bv(n0, kc + h); - bf[2 + h] = bv(n0, kc + 8 + h); - } - auto at = cute::make_tensor(&af[0], cute::Int<8>{}); - auto bt = cute::make_tensor(&bf[0], cute::Int<4>{}); - auto ct = cute::make_tensor(&acc[0], - cute::Int{}); - Atom{}(at, bt, ct); - } - - CUTE_UNROLL - for (int v = 0; v < Geo::kAccPerLane; ++v) - cv(base + v) = acc[v]; - } - } - } -}; - -template -inline constexpr bool tile_shaped_v = tile_v> && - tile_v>; - /// Recognize the atom's per-lane fragment shapes. template inline constexpr bool frag_v = [] { diff --git a/src/tilefoundry/codegen/cuda/tir/nn/mma.py b/src/tilefoundry/codegen/cuda/tir/nn/mma.py index be33e898..f03b61ac 100644 --- a/src/tilefoundry/codegen/cuda/tir/nn/mma.py +++ b/src/tilefoundry/codegen/cuda/tir/nn/mma.py @@ -1,36 +1,33 @@ """Emit effect-form matrix multiply-accumulate operations. -One call, ``tilefoundry::ops::mma``, whichever tier the operand layouts pick: -a lane's gathered fragments take the single PTX instruction and a rank-2 tile -loops the atom over it. The table below is what an atom name is allowed to -emit, not a tier -- other architecture, dtype and shape combinations need their -own runtime mapping ([runtime §2.6](docs/spec/runtime.md#26-cudaops)). +The SM80 declaration emits one ``tilefoundry::ops::mma`` call. Other +declarations need their own runtime mapping; in particular, WGMMA has no +emitter in this stage ([runtime §2.6](docs/spec/runtime.md#26-cudaops)). """ from __future__ import annotations from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.ir.core import Var -from tilefoundry.ir.tir.cuda.nn.mma import Mma +from tilefoundry.ir.tir.cuda.nn.mma import TiledMma +from tilefoundry.ir.tir.cuda.nn.sm80_mma import Mma as Sm80Mma from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen -@register_codegen(CudaTarget, Role.EMIT, Mma) +@register_codegen(CudaTarget, Role.EMIT, TiledMma) def _emit(call, ctx: CudaCodegenContext) -> None: acc, lhs, rhs = call.args[0], call.args[1], call.args[2] if not isinstance(lhs, Var) or not isinstance(rhs, Var) or not isinstance(acc, Var): raise RuntimeError( - "tir.cuda.nn.Mma: codegen path expects Var operands on acc/lhs/rhs" + "tir.cuda.nn.TiledMma: codegen path expects Var operands on acc/lhs/rhs" ) a = ctx.name_for(acc) l = ctx.name_for(lhs) r = ctx.name_for(rhs) - - atom = call.target.atom - if atom is not None and atom.op.name != "SM80_16x8x16_F32BF16BF16F32_TN": + if not isinstance(atom, Sm80Mma): raise RuntimeError( - f"tir.cuda.nn.Mma: no codegen handler for MMA op {atom.op.name!r}" + f"tir.cuda.nn.TiledMma: no codegen handler for atom {atom.reference_name}" ) ctx.emit(f"tilefoundry::ops::mma({l}, {r}, {a});") diff --git a/src/tilefoundry/dsl/T/_platforms.py b/src/tilefoundry/dsl/T/_platforms.py index 677650f3..3eda812f 100644 --- a/src/tilefoundry/dsl/T/_platforms.py +++ b/src/tilefoundry/dsl/T/_platforms.py @@ -1,31 +1,30 @@ -"""``T.cuda`` — the CUDA (NVIDIA-GPU) platform sub-namespace of ``dsl.T``.""" +"""``T.cuda`` — CUDA instruction declarations.""" from __future__ import annotations -from tilefoundry.ir.tir.cuda.nn import mma as _cuda_mma -from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom, MmaOpSpec +from tilefoundry.ir.tir.cuda.nn.sm80_mma import Mma +from tilefoundry.ir.tir.cuda.nn.wgmma import Form, Major, Wgmma -class _MmaNamespace: - """``T.cuda.mma`` — named MMA instructions + the ``atom(op=...)`` builder.""" +class _Sm80Namespace: + Mma = Mma - SM80_16x8x16_F32BF16BF16F32_TN: MmaOpSpec = _cuda_mma.SM80_16x8x16_F32BF16BF16F32_TN - @staticmethod - def atom(op: MmaOpSpec) -> MmaAtom: - """Realize an :class:`MmaAtom` from a named op (CuTe ``make_tiled_mma``).""" - return _cuda_mma.make_atom(op) +class _Sm90Namespace: + Wgmma = Wgmma + Form = Form + Major = Major class _CudaNamespace: - """``T.cuda`` — the CUDA platform namespace.""" + """``T.cuda`` — compile-time CUDA instruction declarations.""" - mma = _MmaNamespace() + sm80 = _Sm80Namespace() + sm90 = _Sm90Namespace() cuda = _CudaNamespace() - PLATFORM_NAMESPACES = {"cuda": cuda} -__all__ = ["cuda", "PLATFORM_NAMESPACES"] +__all__ = ["PLATFORM_NAMESPACES", "cuda"] diff --git a/src/tilefoundry/dsl/_stub_gen.py b/src/tilefoundry/dsl/_stub_gen.py index 28f1e9a2..3a2deca1 100644 --- a/src/tilefoundry/dsl/_stub_gen.py +++ b/src/tilefoundry/dsl/_stub_gen.py @@ -207,33 +207,68 @@ def _module_header(dialect: str, types_seen: set[str]) -> str: def _platform_namespace_stub(dialect: str) -> str | None: - """Typed stubs for ``T`` platform sub-namespaces (``T.cuda.mma.*``). + """Typed stubs for declarative ``T.cuda.sm80`` / ``T.cuda.sm90`` atoms. These are compile-time descriptor surfaces, not OpSchema-backed ops, so the - schema walk never sees them. The namespace *shape* (``cuda.mma`` + the - ``atom(op)`` builder) is fixed; the op set is introspected from the live - namespace so a new ``MmaOpSpec`` shows up automatically. + schema walk never sees them. Each constructor signature is reflected from + the declaration's ``ParamDef`` fields. """ if dialect != "T": return None from tilefoundry.dsl.T._platforms import cuda # noqa: PLC0415 - from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaOpSpec # noqa: PLC0415 + from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom # noqa: PLC0415 + + def constructor(name: str, declaration: type[MmaAtom]) -> str: + params = [] + for param in declaration.parameters: + annotation = getattr(param.annotation, "__name__", "Any") + rendered = f"{param.name}: {annotation}" + if param.has_default: + rendered += " = ..." + params.append(rendered) + params.append("mesh: Mesh | None = ...") + return ( + " @staticmethod\n" + f" def {name}(*, {', '.join(params)}) -> MmaAtom: ..." + ) + + namespaces = [] + for namespace_name in ("sm80", "sm90"): + namespace = getattr(cuda, namespace_name) + declarations = sorted( + (name, value) + for name, value in vars(type(namespace)).items() + if isinstance(value, type) and issubclass(value, MmaAtom) + ) + namespaces.append((namespace_name, namespace, declarations)) - op_names = sorted(n for n, v in vars(type(cuda.mma)).items() if isinstance(v, MmaOpSpec)) lines = [ "# Platform sub-namespaces (not OpSchema-backed).", - "from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom as MmaAtom, MmaOpSpec as MmaOpSpec", + "from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom as MmaAtom", + "from tilefoundry.ir.tir.cuda.nn.wgmma import Form as Form, Major as Major", + "from tilefoundry.ir.types import Mesh as Mesh", "", - "class _CudaMma:", - *[f" {n}: MmaOpSpec" for n in op_names], - " @staticmethod", - " def atom(op: MmaOpSpec) -> MmaAtom: ...", - "", - "class _Cuda:", - " mma: _CudaMma", - "", - "cuda: _Cuda", ] + for namespace_name, namespace, declarations in namespaces: + class_name = f"_Cuda{namespace_name.title()}" + lines.append(f"class {class_name}:") + enum_names = sorted( + name + for name, value in vars(type(namespace)).items() + if isinstance(value, type) and issubclass(value, enum.Enum) + ) + lines.extend(f" {name}: type[{name}]" for name in enum_names) + lines.extend(constructor(name, declaration) for name, declaration in declarations) + lines.append("") + lines.extend( + [ + "class _Cuda:", + " sm80: _CudaSm80", + " sm90: _CudaSm90", + "", + "cuda: _Cuda", + ] + ) return "\n".join(lines) diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 591b8c35..3607d83c 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -414,7 +414,8 @@ def visit_Partial(self, value: Partial, ctx=None) -> str: return f'P("{value.reduction}")' def atom_reference(self, value: MmaAtom, ctx=None) -> str: - return f"T.cuda.mma.atom(op=T.cuda.mma.{value.op.name})" + mesh = None if value.mesh is None else self.visit(value.mesh, ctx) + return value.written(mesh) def render_value(self, value, ctx=None, indent: str = "") -> str: """Render a non-expression attribute through the same visitor when possible.""" diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index fb9775e2..d7704ccf 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma.py @@ -1,22 +1,25 @@ -r"""Define CUDA MMA effects, instructions, and fixed fragment layouts. - -Fragment constants encode the CuTe thread-value maps in row-major order; -``make_atom`` binds a named instruction to those layouts and its required mesh. -The descriptor records live in ``mma_atom.py``. - -See [tir §2.3](docs/spec/tir.md#23-tir-ops). -""" +"""CUDA tiled matrix-multiply-accumulate operation.""" from __future__ import annotations -from tilefoundry.ir.core import Op, VerifyError -from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.core import Op +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Tensor -from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology, UnitType +from tilefoundry.ir.pattern import ( + CapturePattern, + ComposedLayoutPattern, + LayoutPattern, + MeshPattern, + MultipleOfPattern, + OrPattern, +) +from tilefoundry.ir.tir.verify import input_params +from tilefoundry.ir.types import DType, Mesh, UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt -from .mma_atom import MmaAtom, MmaOpSpec +from .mma_atom import AtomPattern, FromAtom, MmaAtom, physical_frames_match, read_on +from .sm80_mma import Mma as _Sm80Mma +from .wgmma import Wgmma _FP_ACC_WIDEN = { (DType.f16, DType.f32), @@ -27,148 +30,111 @@ } -_ATOM_ROLE = {"acc": "C", "lhs": "A", "rhs": "B"} +def _warp_layout_pattern() -> LayoutPattern: + return LayoutPattern( + ((CapturePattern("n", MultipleOfPattern(32)),),), + ((1,),), + per_mode=True, + ) -@register_op(category="nn") -class Mma(Op): - """Matrix-multiply-accumulate: ``acc += lhs @ rhs``.""" - - acc = ParamDef(kind="input", pattern=Tensor) - lhs = ParamDef(kind="input", pattern=Tensor) - rhs = ParamDef(kind="input", pattern=Tensor) - atom = ParamDef(kind="attribute", annotation=MmaAtom, default=None, optional=True) +_WARP_ALIGNED = OrPattern( + ComposedLayoutPattern( + offset=CapturePattern("p0", MultipleOfPattern(32)), + outer=_warp_layout_pattern(), + ), + _warp_layout_pattern(), +) -@register_typeinfer(Mma) +@register_op(category="nn", name="tiled_mma") +class TiledMma(Op): + """Execute one tiled MMA; the atom declares its operand contracts.""" + + @property + def capability(self): + return self.atom.capability + + acc = ParamDef( + kind="input", + effect=MemoryEffect.READ | MemoryEffect.WRITE, + pattern=FromAtom("C"), + ) + lhs = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=FromAtom("A")) + rhs = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=FromAtom("B")) + atom = ParamDef( + kind="attribute", + annotation=MmaAtom, + pattern=AtomPattern(Wgmma, _Sm80Mma), + ) + scope = ParamDef( + kind="attribute", + annotation=Mesh, + pattern=MeshPattern(("thread",), _WARP_ALIGNED), + optional=True, + default=None, + ) + + +@register_typeinfer(TiledMma) def _(call: "Call", ctx: "TypeInferContext") -> UnitType: return UnitType() -@register_verify_stmt(Mma) -def _(call: "Call", ctx: "VerifyContext") -> None: - """Check the operand shapes ``ops::mma`` will read off these layouts.""" - acc_ty = ctx.type_of(call.args[0]) - lhs_ty = ctx.type_of(call.args[1]) - rhs_ty = ctx.type_of(call.args[2]) +@register_verify_stmt(TiledMma) +def verify_mma(call: "Call", ctx: "VerifyContext") -> None: + """Check each operand against its atom and the active physical frame.""" + op = call.target + atom = op.atom + held = tuple(ctx.type_of(arg) for arg in call.args) + for param, value in zip(input_params(type(op)), held): + pattern = read_on(param.pattern, op) + if pattern.match(value) is None: + ctx.error( + call, + f"MMA {param.name} is not one {atom.reference} reads; " + f"it reads {pattern.describe()}", + ) + if ctx.scope is not None and ctx.scope.module is not None: + capabilities = ctx.scope.module.target.architecture.instruction_capabilities + if op.capability not in capabilities: + ctx.error(call, f"target does not support {op.capability}") + if not ctx.mesh_scope: + ctx.error(call, "MMA requires an active physical mesh scope") + current = ctx.mesh_scope[-1] + participation = atom.scope_pattern() + if participation.match(current) is None: + ctx.error( + call, + "MMA enclosing mesh violates declared instruction participation, " + f"which is {participation.describe()}", + ) + if atom.mesh is not None and not physical_frames_match(atom.mesh, current): + ctx.error(call, "MMA atom frame differs from active mesh scope") + verify_operand_shapes(call, ctx) + +def verify_operand_shapes(call: "Call", ctx: "VerifyContext") -> None: + """Check A (M,K), B (K,N), and C (M,N) shapes and dtypes.""" + acc_ty, lhs_ty, rhs_ty = (ctx.type_of(arg) for arg in call.args[:3]) if len(lhs_ty.shape) == 2 and len(rhs_ty.shape) == 2 and len(acc_ty.shape) == 2: m, k_l = lhs_ty.shape[-2], lhs_ty.shape[-1] - rhs_sl = getattr(rhs_ty, "layout", None) - rhs_is_tile = len(getattr(getattr(rhs_sl, "layout", None), "shape", ())) == 2 - if rhs_is_tile: - n, k_r = rhs_ty.shape[-2], rhs_ty.shape[-1] - else: - k_r, n = rhs_ty.shape[-2], rhs_ty.shape[-1] + k_r, n = rhs_ty.shape[-2], rhs_ty.shape[-1] if k_l != k_r: ctx.error(call, f"Mma K-dim mismatch: {k_l} vs {k_r}") if acc_ty.shape[-2] != m or acc_ty.shape[-1] != n: ctx.error( call, - f"Mma acc shape mismatch: expected (...,{m},{n}), got (...,{acc_ty.shape[-2]},{acc_ty.shape[-1]})", + f"Mma acc shape mismatch: expected (...,{m},{n}), got " + f"(...,{acc_ty.shape[-2]},{acc_ty.shape[-1]})", ) if lhs_ty.dtype != rhs_ty.dtype: ctx.error(call, f"Mma lhs/rhs dtype mismatch: {lhs_ty.dtype} vs {rhs_ty.dtype}") if (lhs_ty.dtype, acc_ty.dtype) not in _FP_ACC_WIDEN: - ctx.error(call, f"Mma unsupported dtype combo: input {lhs_ty.dtype} acc {acc_ty.dtype}") - - atom = call.target.atom - if atom is not None: - for role, ty, want in ( - ("acc", acc_ty, atom.C), - ("lhs", lhs_ty, atom.A), - ("rhs", rhs_ty, atom.B), - ): - if getattr(ty, "layout", None) != want: - ctx.error( - call, - f"Mma {role} fragment layout does not match atom {_ATOM_ROLE[role]}", - ) - from tilefoundry.ir.mesh_scope import ( # noqa: PLC0415 - mesh_scope_matches_required_scope, - ) - - if not any( - mesh_scope_matches_required_scope(s, atom.required_scope) for s in ctx.mesh_scope - ): - raise VerifyError( - "T.mma: no enclosing mesh scope hosts the atom's required thread " - f"scope (topology {atom.required_scope.topologies[0].name!r}, " - f"{atom.required_scope.topologies[0].size} lanes)" - ) - - -_SM80_THREAD_MESH = Mesh( - topologies=(Topology("thread", 32),), - layout=Layout(shape=(4, 8), strides=(1, 4)), - names=("warp", "lane"), -) - - -A_FRAG_LAYOUT = Layout(shape=(2, 4, 2, 8, 2), strides=(1, 2, 8, 16, 128)) -_A_FRAG_SHARD = ShardLayout( - layout=A_FRAG_LAYOUT, - attrs=(Split(1), Split(3)), - mesh=_SM80_THREAD_MESH, -) - - -B_FRAG_LAYOUT = Layout(shape=(8, 2, 4, 2), strides=(1, 8, 16, 64)) -_B_FRAG_SHARD = ShardLayout( - layout=B_FRAG_LAYOUT, - attrs=(Split(2), Split(0)), - mesh=_SM80_THREAD_MESH, -) - - -C_FRAG_LAYOUT = Layout(shape=(2, 4, 8, 2), strides=(1, 2, 8, 64)) -_C_FRAG_SHARD = ShardLayout( - layout=C_FRAG_LAYOUT, - attrs=(Split(1), Split(2)), - mesh=_SM80_THREAD_MESH, -) - - -SM80_16x8x16_F32BF16BF16F32_TN = MmaOpSpec( - name="SM80_16x8x16_F32BF16BF16F32_TN", - shape_mnk=(16, 8, 16), - dtype_a=DType.bf16, - dtype_b=DType.bf16, - dtype_c=DType.f32, - operand_layout="TN", -) - - -_ATOM_TABLE: dict[MmaOpSpec, tuple[ShardLayout, ShardLayout, ShardLayout, Mesh]] = { - SM80_16x8x16_F32BF16BF16F32_TN: ( - _A_FRAG_SHARD, - _B_FRAG_SHARD, - _C_FRAG_SHARD, - _SM80_THREAD_MESH, - ), -} - - -def make_atom(op: MmaOpSpec) -> MmaAtom: - """Build the :class:`MmaAtom` for ``op`` (CuTe ``make_tiled_mma`` analog). - - Raises ``KeyError`` (with a clear message) for an instruction that has no - registered fragment layouts yet. - """ - if not isinstance(op, MmaOpSpec): - raise TypeError(f"mma atom(op=...) expects an MmaOpSpec, got {type(op).__name__}") - entry = _ATOM_TABLE.get(op) - if entry is None: - raise KeyError( - f"no fragment layouts registered for MMA op {op.name!r}; " - f"add an entry to ir.tir.cuda.nn.mma._ATOM_TABLE" + ctx.error( + call, + f"Mma unsupported dtype combo: input {lhs_ty.dtype} acc {acc_ty.dtype}", ) - a, b, c, scope = entry - return MmaAtom(op=op, A=a, B=b, C=c, required_scope=scope) -__all__ = [ - "Mma", - "make_atom", - "SM80_16x8x16_F32BF16BF16F32_TN", -] +__all__ = ["TiledMma", "verify_mma", "verify_operand_shapes"] diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py index 86aca13f..35e0a3cb 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py @@ -1,33 +1,282 @@ -"""CUDA MMA op / atom model (CuTe ``MMA_Op`` → ``MMA_Atom``).""" +"""Declarative CUDA MMA atoms and operand patterns.""" from __future__ import annotations +import itertools from dataclasses import dataclass -from tilefoundry.ir.types import DType, Mesh, ShardLayout +from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.pattern import ( + ABSENT, + CapturePattern, + ComposedLayoutPattern, + Match, + MeshPattern, + MultipleOfPattern, + OneOfPattern, + OrPattern, + Pattern, + alternatives_of, + arrangement_pattern, + matched, + resolved, +) +from tilefoundry.ir.pattern.match import written_binding, written_bindings, written_place +from tilefoundry.ir.types import Layout, Mesh +from tilefoundry.ir.types.dim import DimVar +from tilefoundry.ir.types.layout_algebra import coalesce, frame_of -@dataclass(frozen=True) -class MmaOpSpec: - """Describe a named MMA instruction, including operand layout notation.""" +class MmaAtom: + """One instruction declaration; an instance binds its authored parameters.""" + + namespace: str + scope: Mesh + capability: str + C: object + A: object + B: object + + parameters: tuple[ParamDef, ...] = () + reference_name: str = "" + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + cls.parameters = tuple(value for value in vars(cls).values() if isinstance(value, ParamDef)) + cls.reference_name = f"{cls.namespace}.{cls.__name__}" + + def __init__(self, *, mesh: Mesh | None = None, **bindings): + unknown = set(bindings) - {param.name for param in self.parameters} + if unknown: + raise ValueError( + f"{self.reference_name} takes no parameter {', '.join(sorted(unknown))}" + ) + held = {} + for param in self.parameters: + value = bindings[param.name] if param.name in bindings else self._implied(param, held) + if matched(param.pattern, value, held) is None: + raise ValueError( + f"{self.reference_name}: {param.name}={written_binding(value)} is not one " + f"it takes{self._where(held)}; it takes {param.name} " + f"{written_place(resolved(param.pattern, held), param.name)}" + ) + held[param.name] = value + self.bindings = held + self.mesh = mesh + self._roles: dict[str, object] = {} + + @classmethod + def _implied(cls, param: ParamDef, held: dict): + left = resolved(param.pattern, held) + if not isinstance(left, Pattern) and left is not ABSENT: + return left + if param.has_default: + return param.default + raise ValueError(f"{cls.reference_name} needs {param.name}") + + @staticmethod + def _where(held: dict) -> str: + return "" if not held else f" where {written_bindings(held.items())}" + + @classmethod + def bindings_for(cls, reads) -> tuple[dict, ...]: + """Every parameter binding the operand types leave possible.""" + choices = [] + for param in cls.parameters: + if isinstance(param.pattern, OneOfPattern): + choices.append([(param.name, value) for value in param.pattern.values]) + continue + if param.has_default: + continue + extent = cls._extent_of(reads, param.name) or 0 + choices.append( + [ + (param.name, value) + for value in range(1, extent + 1) + if extent % value == 0 and matched(param.pattern, value) is not None + ] + ) + return tuple(dict(one) for one in itertools.product(*choices)) + + @staticmethod + def _extent_of(reads, name: str) -> int | None: + for pattern, held in reads: + for _, tensor in alternatives_of(pattern): + for extent, value in zip( + getattr(tensor, "shape", None) or (), getattr(held, "shape", ()) + ): + if isinstance(extent, DimVar) and extent.name == name and type(value) is int: + return value + return None + + @classmethod + def role_of(cls, role: str, bindings=()): + return resolved(getattr(cls, role), dict(bindings)) + + def role(self, role: str): + held = self._roles.get(role) + if held is None: + held = self._roles[role] = self.role_of(role, self.bindings) + return held + + @property + def shape_mnk(self) -> tuple[int, int, int]: + (m, n), (_, k) = tuple(self.role("C").shape), tuple(self.role("A").shape) + return m, n, k + + @property + def dtype_a(self): + return self.role("A").dtype + + @property + def dtype_b(self): + return self.role("B").dtype + + @property + def dtype_c(self): + return self.role("C").dtype + + @property + def required_scope(self) -> Mesh: + return self.scope + + @classmethod + def scope_pattern(cls) -> MeshPattern: + topology, = cls.scope.topologies + size = topology.size + bare = arrangement_pattern(cls.scope.layout, per_mode=True) + sliced = ComposedLayoutPattern( + offset=CapturePattern("p0", MultipleOfPattern(size)), + outer=arrangement_pattern(cls.scope.layout, per_mode=True), + ) + return MeshPattern((topology.name,), OrPattern(sliced, bare)) + + def on(self, mesh: Mesh) -> MmaAtom: + return type(self)(mesh=mesh, **self.bindings) + + def written(self, mesh: str | None = None) -> str: + stated, held = [], {} + for param in self.parameters: + value = self.bindings[param.name] + try: + implied = self._implied(param, held) + except ValueError: + implied = None + if implied is None or implied != value: + stated.append(f"{param.name}={self.written_value(value)}") + held[param.name] = value + if mesh is not None: + stated.append(f"mesh={mesh}") + return f"{self.reference_name}({', '.join(stated)})" + + @classmethod + def written_value(cls, value) -> str: + if type(value) is int: + return str(value) + return f"{cls.namespace}.{type(value).__name__}.{value.name}" + + @property + def reference(self) -> str: + return self.written() + + def describe(self) -> str: + return self.reference + + def __eq__(self, other): + return ( + type(other) is type(self) + and other.bindings == self.bindings + and other.mesh == self.mesh + ) - name: str - shape_mnk: tuple[int, int, int] - dtype_a: DType - dtype_b: DType - dtype_c: DType - operand_layout: str + def __hash__(self): + return hash((type(self), tuple(self.bindings.items()), self.mesh)) + + def __repr__(self): + return self.written(None if self.mesh is None else repr(self.mesh)) + + +@dataclass(frozen=True, init=False) +class AtomPattern(Pattern): + """An instance of any declared atom class.""" + + declarations: tuple[type[MmaAtom], ...] + + def __init__(self, *declarations): + object.__setattr__(self, "declarations", tuple(declarations)) + + def match(self, subject, captures=None): + if not isinstance(subject, self.declarations): + return None + return Match({**dict(captures or {}), **subject.bindings}) + + def describe(self, name: str = "_") -> str: + return "one of " + ", ".join(held.reference_name for held in self.declarations) + + def resolve(self, bindings): + return self @dataclass(frozen=True) -class MmaAtom: - """Realized MMA atom (CuTe ``MMA_Atom``) — op + fragment layouts + required scope.""" +class FromAtom(Pattern): + """An operand pattern read from one role of the call's atom.""" + + role: str + + def __post_init__(self): + if self.role not in ("A", "B", "C"): + raise ValueError("an MMA operand reads role A, B or C of its atom") + + def read_on(self, op): + return op.atom.role_of(self.role, getattr(op.atom, "bindings", ())) + + def match(self, subject, captures=None): + raise TypeError(f"the {self.role} operand is read against a call's atom; ask read_on(op)") + + def describe(self, name: str = "_") -> str: + return f"the {self.role} operand of its atom" + + +def read_on(pattern, op): + held = getattr(pattern, "read_on", None) + return pattern if held is None else held(op) + + +def _reversed_modes(modes): + if isinstance(modes, tuple): + return tuple(_reversed_modes(mode) for mode in reversed(modes)) + return modes + + +def _reverse(layout: Layout) -> Layout: + return Layout( + _reversed_modes(tuple(layout.shape)), + _reversed_modes(tuple(layout.strides)), + ) + + +def physical_frames_match(left: Mesh, right: Mesh) -> bool: + """Compare affine offsets and ordered participant lanes.""" + if left.topologies != right.topologies: + return False + + def frame(mesh): + framed = frame_of(mesh.layout) + if framed is None: + return None + offset, layout = framed + if layout.strides is None: + return None + return offset, coalesce(_reverse(layout)) - op: MmaOpSpec - A: ShardLayout - B: ShardLayout - C: ShardLayout - required_scope: Mesh + a, b = frame(left), frame(right) + return a is not None and b is not None and a == b -__all__ = ["MmaOpSpec", "MmaAtom"] +__all__ = [ + "AtomPattern", + "FromAtom", + "MmaAtom", + "physical_frames_match", + "read_on", +] diff --git a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py new file mode 100644 index 00000000..9ea50dd7 --- /dev/null +++ b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py @@ -0,0 +1,59 @@ +"""The fixed SM80 warp MMA declaration.""" + +from __future__ import annotations + +from tilefoundry.ir.pattern import ShardLayoutPattern, TensorPattern, arrangement_pattern +from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology +from tilefoundry.ir.types.storage import StorageKind as S + +from .mma_atom import MmaAtom + +WARP = Mesh( + topologies=(Topology("thread", 32),), + layout=Layout(shape=(4, 8), strides=(1, 4)), + names=("warp", "lane"), +) + +_A_FRAGMENT = ShardLayout( + layout=Layout(shape=(2, 4, 2, 8, 2), strides=(1, 2, 8, 16, 128)), + attrs=(Split(1), Split(3)), + mesh=WARP, +) +_B_FRAGMENT = ShardLayout( + layout=Layout(shape=(8, 2, 4, 2), strides=(1, 8, 16, 64)), + attrs=(Split(2), Split(0)), + mesh=WARP, +) +_C_FRAGMENT = ShardLayout( + layout=Layout(shape=(2, 4, 8, 2), strides=(1, 2, 8, 64)), + attrs=(Split(1), Split(2)), + mesh=WARP, +) + + +def _fragment(shape: tuple, dtype, held: ShardLayout) -> TensorPattern: + return TensorPattern( + shape=shape, + dtype=dtype, + storage=S.RMEM, + layout=ShardLayoutPattern( + arrangement_pattern(held.layout), + held.attrs, + WARP, + ), + ) + + +class Mma(MmaAtom): + """A BF16 warp MMA, 16 x 8 x 16, accumulating in F32.""" + + namespace = "T.cuda.sm80" + scope = WARP + capability = "tensor_core" + + C = _fragment((16, 8), DType.f32, _C_FRAGMENT) + A = _fragment((16, 16), DType.bf16, _A_FRAGMENT) + B = _fragment((16, 8), DType.bf16, _B_FRAGMENT) + + +__all__ = ["Mma", "WARP"] diff --git a/src/tilefoundry/ir/tir/cuda/nn/wgmma.py b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py new file mode 100644 index 00000000..8544dce5 --- /dev/null +++ b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py @@ -0,0 +1,245 @@ +"""The parameterized SM90 warpgroup MMA declaration.""" + +from __future__ import annotations + +from enum import Enum + +from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.pattern import ( + AndPattern, + CapturePattern, + ComposedLayoutPattern, + GuardPattern, + LayoutPattern, + MultipleOfPattern, + OneOfPattern, + RangePattern, + ShardLayoutPattern, + SwitchPattern, + SwizzlePattern, + TensorPattern, +) +from tilefoundry.ir.pattern.match import is_symbolic +from tilefoundry.ir.types import Broadcast, DType, Layout, Mesh, Split, Topology +from tilefoundry.ir.types.dim import DimVar +from tilefoundry.ir.types.storage import StorageKind as S + +from .mma_atom import MmaAtom + +WARPGROUP = Mesh( + (Topology("thread", 128),), + Layout((4, 8, 4), (32, 4, 1)), + ("warp", "lane8", "lane4"), +) + + +class Major(Enum): + """Which of a shared descriptor's axes runs contiguously.""" + + MN = "MN-major" + K = "K-major" + + +class Form(Enum): + """Whether WGMMA reads A from shared memory or registers.""" + + SS = "SS" + RS = "RS" + + +class Swizzled(Enum): + """The SM90 descriptor's shared-memory swizzle width.""" + + INTERLEAVE = "INTERLEAVE" + SW32 = "SW32" + SW64 = "SW64" + SW128 = "SW128" + + @property + def bits(self) -> int: + return ("INTERLEAVE", "SW32", "SW64", "SW128").index(self.value) + + @property + def run(self) -> int: + return DESCRIPTOR_UNIT << self.bits + + +DESCRIPTOR_UNIT_BYTES = 16 +DESCRIPTOR_UNIT = DESCRIPTOR_UNIT_BYTES * 8 // DType.bf16.bit_width +START = "k0" + + +def Descriptor( + rows, + cols: int, + major, + *, + selects: str, + k_first: bool = False, +) -> SwitchPattern: + """Every shared BF16 descriptor arrangement for one logical tile.""" + if isinstance(major, str): + return SwitchPattern( + major, + { + read: Descriptor( + rows, + cols, + read, + selects=selects, + k_first=k_first, + ) + for read in Major + }, + ) + unit = DESCRIPTOR_UNIT + branches = {} + for mode in Swizzled: + width = mode.run + along = rows if major is Major.MN else cols + leading, stride = cols * width, width * unit + guarded, sliced = False, False + if is_symbolic(along): + if major is not Major.MN: + raise ValueError("a K-major descriptor reads a K it states") + guarded = True + elif along % width: + if major is Major.MN or width % along: + continue + sliced = True + if sliced: + shape, strides = ((rows // unit, unit), cols), ((unit * width, width), 1) + elif major is Major.MN: + shape = ((rows // width, width), (cols // unit, unit)) + strides = ((leading, 1), (stride, width)) + else: + shape = ((rows // unit, unit), (cols // width, width)) + strides = ((leading, width), (stride, 1)) + if k_first: + shape, strides = shape[::-1], strides[::-1] + held = LayoutPattern(shape, strides) + if mode.bits: + start = ( + CapturePattern(START, OneOfPattern(tuple(range(0, width, along)))) + if sliced + else 0 + ) + held = ComposedLayoutPattern( + SwizzlePattern(mode.bits, 4, 3), + start, + held, + ) + branches[mode] = ( + GuardPattern(along, MultipleOfPattern(width), held) if guarded else held + ) + return SwitchPattern(selects, branches) + + +def Fragment(rows: int, cols) -> LayoutPattern: + """CuTe ``CLayout_64xN``; at N=16, also the RS A fragment.""" + if rows != 64: + raise ValueError("the only register fragment this back end reads is 64 rows deep") + return LayoutPattern((8, 2, 4, 2, 4, cols // 8), (1, 8, 16, 64, 128, 512)) + + +SHARED_BY_ALL = (Broadcast(), Broadcast(), Broadcast()) +HELD_PER_THREAD = (Split(2), Split(0), Split(4)) + + +def shared(arrangement) -> ShardLayoutPattern: + return ShardLayoutPattern(arrangement, SHARED_BY_ALL, WARPGROUP) + + +def held(arrangement) -> ShardLayoutPattern: + return ShardLayoutPattern(arrangement, HELD_PER_THREAD, WARPGROUP) + + +N = DimVar("n", 8, 257) + + +class Wgmma(MmaAtom): + """A BF16 warpgroup MMA, 64 x n x 16, accumulating in F32.""" + + namespace = "T.cuda.sm90" + scope = WARPGROUP + capability = "wgmma" + + n = ParamDef( + kind="attribute", + annotation=int, + pattern=AndPattern((MultipleOfPattern(8), RangePattern(lo=N.lo, hi=N.hi - 1))), + ) + form = ParamDef( + kind="attribute", + annotation=Form, + pattern=OneOfPattern(tuple(Form)), + ) + a_major = ParamDef( + kind="attribute", + annotation=Major, + pattern=SwitchPattern( + "form", + { + Form.SS: OneOfPattern(tuple(Major)), + Form.RS: Major.K, + }, + ), + default=Major.MN, + ) + + C = TensorPattern( + shape=(64, N), + dtype=DType.f32, + storage=S.RMEM, + layout=held(Fragment(64, N)), + ) + A = SwitchPattern( + "form", + { + Form.SS: TensorPattern( + shape=(64, 16), + dtype=DType.bf16, + storage=S.SMEM, + layout=shared(Descriptor(64, 16, "a_major", selects="a_swizzle")), + ), + Form.RS: TensorPattern( + shape=(64, 16), + dtype=DType.bf16, + storage=S.RMEM, + layout=held(Fragment(64, 16)), + ), + }, + ) + B = TensorPattern( + shape=(16, N), + dtype=DType.bf16, + storage=S.SMEM, + layout=shared( + Descriptor( + N, + 16, + Major.MN, + selects="b_swizzle", + k_first=True, + ) + ), + ) + + +__all__ = [ + "DESCRIPTOR_UNIT", + "DESCRIPTOR_UNIT_BYTES", + "Descriptor", + "Form", + "Fragment", + "HELD_PER_THREAD", + "Major", + "N", + "SHARED_BY_ALL", + "START", + "Swizzled", + "WARPGROUP", + "Wgmma", + "held", + "shared", +] diff --git a/src/tilefoundry/ir/tir/verify.py b/src/tilefoundry/ir/tir/verify.py index 2595188f..1b9cc86d 100644 --- a/src/tilefoundry/ir/tir/verify.py +++ b/src/tilefoundry/ir/tir/verify.py @@ -50,7 +50,7 @@ _PRIM_FUNCTION = "[tir §1.3](docs/spec/tir.md#13-primfunction)" -def _input_params(op_type: type) -> tuple: +def input_params(op_type: type) -> tuple: return tuple(param for param in op_type._op_schema.signature if param.kind == "input") @@ -60,7 +60,7 @@ def verify_between(call, ctx, lead: str = "") -> None: rules = between_rules(op_type) if not rules: return - names = tuple(param.name for param in _input_params(op_type)) + names = tuple(param.name for param in input_params(op_type)) operands = dict(zip(names, (ctx.type_of(arg) for arg in call.args))) for rule in rules: if not rule.holds(operands): @@ -69,7 +69,7 @@ def verify_between(call, ctx, lead: str = "") -> None: def verify_operands(call, ctx, label: str) -> None: """Hold each operand to the pattern declared for its parameter.""" - for param, arg in zip(_input_params(type(call.target)), call.args): + for param, arg in zip(input_params(type(call.target)), call.args): if param.pattern is None: continue value = ctx.type_of(arg) @@ -602,4 +602,10 @@ def verify_module(fns) -> None: ) -__all__ = ["verify_between", "verify_module", "verify_operands", "verify_prim_function"] +__all__ = [ + "input_params", + "verify_between", + "verify_module", + "verify_operands", + "verify_prim_function", +] diff --git a/tests/fixtures/tir/mma.py b/tests/fixtures/tir/mma.py index 63a0213c..afc7666a 100644 --- a/tests/fixtures/tir/mma.py +++ b/tests/fixtures/tir/mma.py @@ -45,12 +45,7 @@ def mm_device( T.copy(a_view, a_frag) T.copy(b_view, b_frag) T.fill(acc, 0.0) - T.mma( - acc, - a_frag, - b_frag, - atom=T.cuda.mma.atom(op=T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN), - ) + T.tiled_mma(acc, a_frag, b_frag, atom=T.cuda.sm80.Mma()) c_view = T.tensor_view( T.ptr_of(c), layout=((2, 4 @ _warp.warp, 8 @ _warp.lane, 2), (1, 2, 8, 64)) ) diff --git a/tests/installed/smoke_spec.py b/tests/installed/smoke_spec.py index d69ca518..d641897e 100644 --- a/tests/installed/smoke_spec.py +++ b/tests/installed/smoke_spec.py @@ -49,27 +49,6 @@ def test_spec_lists_and_prints_cache_update(tf) -> None: assert "eval/runtime, not typeinfer" in done.stdout -def test_spec_separates_two_sections_that_would_share_a_key(tf) -> None: - """`tir.md` names a field `name` twice, under `SymbolRef` and under `MmaOpSpec`. - - `tir.md` names a field `name` twice, under `SymbolRef` and under - `MmaOpSpec`. Each is reachable by its enclosing section; the bare key is - not, because it would have to pick one. - """ - symbol = tf("spec", "tir", "symbolref/name") - assert symbol.returncode == 0, symbol.stderr - assert "canonical name of a `PrimFunction`" in symbol.stdout - assert "uniquely identify the instruction" not in symbol.stdout - - atom = tf("spec", "tir", "mmaopspec/name") - assert atom.returncode == 0, atom.stderr - assert "uniquely identify the instruction" in atom.stdout - - bare = tf("spec", "tir", "name") - assert bare.returncode == 1 - assert "no section 'name'" in bare.stderr - - def test_spec_rejects_a_section_that_does_not_exist(tf) -> None: done = tf("spec", "dsl", "9.9") assert done.returncode == 1 diff --git a/tests/ir/test_visitor.py b/tests/ir/test_visitor.py index 3b0f5f33..7f37d792 100644 --- a/tests/ir/test_visitor.py +++ b/tests/ir/test_visitor.py @@ -18,7 +18,8 @@ from tilefoundry.ir.core import Call, Constant, Expr, Op, Var from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion -from tilefoundry.ir.tir.cuda.nn.mma import Mma +from tilefoundry.ir.tir.cuda.nn.mma import TiledMma +from tilefoundry.ir.tir.cuda.nn.sm80_mma import Mma from tilefoundry.ir.tir.memory import Copy, Fill from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf @@ -360,7 +361,7 @@ def _seq(*items) -> Sequential: ), _eval_call(Copy(), _var("s4"), _var("d4")), _eval_call(Fill(), _var("t"), _const(0.0)), - _eval_call(Mma(), _var("L"), _var("R"), _var("A")), + _eval_call(TiledMma(atom=Mma()), _var("L"), _var("R"), _var("A")), Sequential(body=()), ) m = StmtMutator() diff --git a/tests/ir/types/test_mma_fragment_layouts.py b/tests/ir/types/test_mma_fragment_layouts.py index 96d2c1c8..1b107169 100644 --- a/tests/ir/types/test_mma_fragment_layouts.py +++ b/tests/ir/types/test_mma_fragment_layouts.py @@ -9,18 +9,29 @@ from __future__ import annotations +from tilefoundry.dsl import T from tilefoundry.ir.core import Call, Var from tilefoundry.ir.hir.sharding.reshard import Reshard -from tilefoundry.ir.tir.cuda.nn.mma import SM80_16x8x16_F32BF16BF16F32_TN, make_atom from tilefoundry.ir.types import DType, ShardLayout, Split, TensorType from tilefoundry.ir.types.int_tuple import flatten, product from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.typeinfer import inference_type -_ATOM = make_atom(SM80_16x8x16_F32BF16BF16F32_TN) -A_FRAG_SHARD = _ATOM.A -B_FRAG_SHARD = _ATOM.B -C_FRAG_SHARD = _ATOM.C +_ATOM = T.cuda.sm80.Mma() + + +def _realized_fragment(role: str) -> ShardLayout: + declared = _ATOM.role(role).layout + return ShardLayout( + layout=declared.arrangement.fixed(), + attrs=declared.attrs, + mesh=declared.mesh, + ) + + +A_FRAG_SHARD = _realized_fragment("A") +B_FRAG_SHARD = _realized_fragment("B") +C_FRAG_SHARD = _realized_fragment("C") def test_per_thread_element_counts_and_split_extents() -> None: diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 531ea6e7..8c928054 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -1,4 +1,4 @@ -"""``ops::mma``'s tile tier: the atom looped over a whole shared-memory tile. +"""``ops::mma`` on explicitly gathered per-lane fragments. See [runtime §2.6](docs/spec/runtime.md#26-cudaops). """ @@ -18,8 +18,6 @@ from tilefoundry.target import CpuTarget, CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") -_OP = T.cuda.mma.SM80_16x8x16_F32BF16BF16F32_TN - _MESH_LAYOUT = Layout(shape=(4, 8), strides=(1, 4)) @@ -37,7 +35,7 @@ def test_handwritten_mma_matches_torch() -> None: @module(entry="tile_host", target=_CUDA, topologies=(Topology("thread", 32),)) class MmaTile: - """A 16x16 by 16x8 product done as one atom, looped by the tile's shape.""" + """A 16x16 by 16x8 product done as one atom.""" @prim_func(target=_CUDA) def tile_device( @@ -45,7 +43,6 @@ def tile_device( b: Tensor[(128,), "bf16"], c: Tensor[(128,), "f32"], ): - atom = T.cuda.mma.atom(op=_OP) with Mesh( (Topology("thread", 32),), _MESH_LAYOUT, @@ -81,23 +78,58 @@ def tile_device( ) b_tile = T.alloc_tensor( Tensor[ - (8, 16), + (16, 8), "bf16", ShardLayout( - layout=Layout(shape=(8, 16), strides=(1, 8)), + layout=Layout(shape=(16, 8), strides=(8, 1)), attrs=(Broadcast(), Broadcast()), mesh=m, ), "smem", ] ) - acc = T.alloc_tensor(Tensor[(16, 8), "f32", atom.C, "rmem"]) + a_frag = T.alloc_tensor( + Tensor[ + (16, 16), + "bf16", + ((2, 4 @ m.warp, 2, 8 @ m.lane, 2), (1, 2, 8, 16, 128)), + "rmem", + ] + ) + b_frag = T.alloc_tensor( + Tensor[ + (16, 8), + "bf16", + ((8 @ m.lane, 2, 4 @ m.warp, 2), (1, 8, 16, 64)), + "rmem", + ] + ) + acc = T.alloc_tensor( + Tensor[ + (16, 8), + "f32", + ((2, 4 @ m.warp, 8 @ m.lane, 2), (1, 2, 8, 64)), + "rmem", + ] + ) T.copy(a_view, a_tile) T.copy(b_view, b_tile) - T.fill(acc, 0.0) T.sync(m) - T.mma(acc, a_tile, b_tile) - c_view = T.tensor_view(T.ptr_of(c), layout=atom.C) + a_fragment_view = T.tensor_view( + T.ptr_of(a_tile), + layout=((2, 4 @ m.warp, 2, 8 @ m.lane, 2), (1, 2, 8, 16, 128)), + ) + b_fragment_view = T.tensor_view( + T.ptr_of(b_tile), + layout=((8 @ m.lane, 2, 4 @ m.warp, 2), (1, 8, 16, 64)), + ) + T.copy(a_fragment_view, a_frag) + T.copy(b_fragment_view, b_frag) + T.fill(acc, 0.0) + T.tiled_mma(acc, a_frag, b_frag, atom=T.cuda.sm80.Mma()) + c_view = T.tensor_view( + T.ptr_of(c), layout=((2, 4 @ m.warp, 8 @ m.lane, 2), (1, 2, 8, 64)) + ) T.copy(acc, c_view) @prim_func(target=CpuTarget()) @@ -110,8 +142,7 @@ def tile_host( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_tile_tier_matches_torch_matmul() -> None: - """Selected by rank-2 static A and B local views over the whole tile.""" +def test_fragment_mma_matches_torch_matmul() -> None: rm = tilefoundry.compile(MmaTile, target=_CUDA) torch.manual_seed(0) a = torch.randn(256, dtype=torch.bfloat16, device="cuda") @@ -119,5 +150,5 @@ def test_tile_tier_matches_torch_matmul() -> None: c = torch.zeros(128, dtype=torch.float32, device="cuda") rm(a, b, c) torch.cuda.synchronize() - expected = a.view(16, 16).float().t() @ b.view(16, 8).float() + expected = a.view(16, 16).float() @ b.view(16, 8).float() assert torch.allclose(c.view(16, 8), expected, rtol=2e-2, atol=2e-2) diff --git a/tests/runtime/cuda/static_asserts.cu b/tests/runtime/cuda/static_asserts.cu index 57256a6d..7f327de3 100644 --- a/tests/runtime/cuda/static_asserts.cu +++ b/tests/runtime/cuda/static_asserts.cu @@ -375,7 +375,7 @@ __global__ void k(float *p) { #endif #if CASE == 18 -/// mma: operands that are neither a tile nor the atom's fragments. +/// mma: operands that are not the atom's fragments. __global__ void k(float *p) { float a[3] = {}, b[3] = {}, c[3] = {}; auto at = cute::make_tensor(cute::make_rmem_ptr(&a[0]), cute::Int<3>{}); @@ -404,30 +404,6 @@ __global__ void k() { } #endif -#if CASE == 21 -/// mma's tile tier: an accumulator mesh that is not a whole number of warps. -__global__ void k(cute::bfloat16_t *p, float *q) { - auto mesh = tmesh<16>(); - auto alay = - cute::make_layout(cute::make_shape(cute::Int<16>{}, cute::Int<16>{}), - cute::make_stride(cute::Int<16>{}, cute::Int<1>{})); - auto a = make_shard_tensor( - cute::make_tensor(cute::make_smem_ptr(p), alay), alay, - make_shard_layout(alay, mesh, cute::make_tuple(shard::B{}))); - auto blay = - cute::make_layout(cute::make_shape(cute::Int<8>{}, cute::Int<16>{}), - cute::make_stride(cute::Int<16>{}, cute::Int<1>{})); - auto b = make_shard_tensor( - cute::make_tensor(cute::make_smem_ptr(p + 256), blay), blay, - make_shard_layout(blay, mesh, cute::make_tuple(shard::B{}))); - auto clay = cute::make_layout(cute::make_shape(cute::Int<8>{})); - auto c = make_shard_tensor( - cute::make_tensor(cute::make_rmem_ptr(q), clay), clay, - make_shard_layout(clay, mesh, cute::make_tuple(shard::B{}))); - mma(a, b, c); -} -#endif - #if CASE == 22 /// A reduced axis that divides into neither whole lanes nor whole warps. __global__ void k(float *p) { diff --git a/tests/runtime/test_cuda_static_asserts.py b/tests/runtime/test_cuda_static_asserts.py index b77fc274..07d1e845 100644 --- a/tests/runtime/test_cuda_static_asserts.py +++ b/tests/runtime/test_cuda_static_asserts.py @@ -35,10 +35,9 @@ 15: "the two projected slices must hold the same number of elements", 16: "ops::dot (warp tier): the fastest axis of the operands' mesh", 17: "ops::dot (block tier): the operands' mesh must be a whole number of warps", - 18: "the operands are neither a rank-2 static tile", + 18: "operands must be the atom's (8, 4, 4) lane fragments", 19: "ops::sync: a mesh's scope must be cta or thread", 20: "ops::sync: a CTA mesh needs the module's grid-barrier counter", - 21: "ops::mma (tile tier): the accumulator's mesh must be a whole number of warps", 22: "a reduced mesh axis must divide into whole lanes and whole warps", 24: "both operands must leave the tile whole on every ", 25: "ops::sync: a mesh that skips warps names no barrier", From 85c5fcd5d2a27c0e3400d5b7d9687fb30fd187da Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 02:47:22 +0800 Subject: [PATCH 03/19] feat(cuda): add asynchronous tensor instructions --- docs/spec/codegen.md | 10 +- docs/spec/runtime.md | 41 +- docs/spec/tir.md | 73 +++- .../runtime/cuda/ops/copy_async_bulk.cuh | 12 + .../copy_async_bulk_impl.h} | 40 +- .../tilefoundry/runtime/cuda/ops/ldmatrix.cuh | 10 + .../runtime/cuda/ops/ldmatrix/ldmatrix_impl.h | 47 ++ include/tilefoundry/runtime/cuda/ops/tma.cuh | 13 - include/tilefoundry/runtime/cuda/runtime.cuh | 3 +- scripts/runtime_spec_surface.py | 3 +- .../tir/memory/{tma.py => copy_async_bulk.py} | 10 +- .../cuda/tir/memory/copy_async_tensor.py | 15 + .../codegen/cuda/tir/memory/ldmatrix.py | 20 + src/tilefoundry/ir/pattern/__init__.py | 4 + src/tilefoundry/ir/pattern/constraint.py | 26 +- src/tilefoundry/ir/pattern/pattern.py | 50 ++- src/tilefoundry/ir/pattern/utils.py | 11 + src/tilefoundry/ir/tir/async_copy.py | 47 +- .../ir/tir/cuda/memory/copy_async_bulk.py | 49 +++ .../ir/tir/cuda/memory/copy_async_tensor.py | 401 ++++++++++++++++++ .../ir/tir/cuda/memory/ldmatrix.py | 38 ++ src/tilefoundry/ir/tir/cuda/memory/tma.py | 48 --- src/tilefoundry/ir/types/layout_algebra.py | 80 ++++ tests/ops/tir/cuda/test_mma.py | 6 +- tests/ops/tir/cuda/test_tma.py | 8 +- tests/runtime/cuda/static_asserts.cu | 4 +- tests/runtime/test_cuda_static_asserts.py | 2 +- 27 files changed, 934 insertions(+), 137 deletions(-) create mode 100644 include/tilefoundry/runtime/cuda/ops/copy_async_bulk.cuh rename include/tilefoundry/runtime/cuda/ops/{tma/tma_impl.h => copy_async_bulk/copy_async_bulk_impl.h} (74%) create mode 100644 include/tilefoundry/runtime/cuda/ops/ldmatrix.cuh create mode 100644 include/tilefoundry/runtime/cuda/ops/ldmatrix/ldmatrix_impl.h delete mode 100644 include/tilefoundry/runtime/cuda/ops/tma.cuh rename src/tilefoundry/codegen/cuda/tir/memory/{tma.py => copy_async_bulk.py} (69%) create mode 100644 src/tilefoundry/codegen/cuda/tir/memory/copy_async_tensor.py create mode 100644 src/tilefoundry/codegen/cuda/tir/memory/ldmatrix.py create mode 100644 src/tilefoundry/ir/tir/cuda/memory/copy_async_bulk.py create mode 100644 src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py create mode 100644 src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py delete mode 100644 src/tilefoundry/ir/tir/cuda/memory/tma.py diff --git a/docs/spec/codegen.md b/docs/spec/codegen.md index 66282a0d..e7a824f5 100644 --- a/docs/spec/codegen.md +++ b/docs/spec/codegen.md @@ -207,12 +207,20 @@ tuples because it does not materialize an aggregate. The codegen context records the structural tuple by its fresh SSA `Var` identity so consumers can recover its elements without target-side storage. -Effect Ops (`Copy`, `Fill`, `TiledMma`, `tir.nn.*`, ...) appear in Stmt +Effect Ops (`Copy`, `Fill`, `TiledMma`, `CopyAsyncBulk`, `LdMatrix`, +`tir.nn.*`, ...) appear in Stmt position as `Evaluate(op, args)` rather than as Stmt subclasses. The walker matches `Evaluate` and dispatches on `type(callable)` through the handler registry. Handlers stay small; the runtime function they call carries the semantic load. +CUDA emits `CopyAsyncBulk` and `LdMatrix` as the uniform runtime calls +`tilefoundry::ops::copy_async_bulk(...)` and +`tilefoundry::ops::ldmatrix(...)`. `CopyAsyncTensor` is nevertheless a public +TIR declaration, but CUDA emission MUST fail explicitly until the host can +construct and pass encoded tensor maps; it MUST NOT substitute a bulk or +thread-issued copy. + ## 3. Runtime-owned op dispatch Where more than one runtime template implements an op, codegen emits **one diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index 87ab2635..9d143409 100644 --- a/docs/spec/runtime.md +++ b/docs/spec/runtime.md @@ -1197,12 +1197,13 @@ the one time it was written down as the rule it argued a dependency chain into | `mma` | one atom's per-lane operand fragments | | `rmsnorm` | the row dependency chain and the destination's shard layout | | `sync` | the mesh's scope, base and count | -| `tma_copy` | both shard layouts, asserted: one contiguous run each, whole tiles, matching element types | +| `copy_async_bulk` | both shard layouts, asserted: one contiguous run each, whole tiles, matching element types | +| `ldmatrix` | the source tile's declared shared-memory layout and the destination atom fragment | Anything that takes a raw pointer, an `int` or a type and answers a question about it is not an op but a utility, and belongs outside `ops::` — the warp-scoped primitives of [§2.7](#27-cudautility) and the `mbarrier` -instructions a caller writes around `tma_copy` are both that. In particular the +instructions a caller writes around `copy_async_bulk` are both that. In particular the runtime publishes **no predicate and no constant reporting which tier an op selected or how wide a move it chose**: a caller cannot use one to decide how to build its operands, since the answer is a function of the operand types it would @@ -1538,17 +1539,18 @@ __device__ inline void sync(Mesh const &mesh, launch, so the caller passes it; with none, a cooperative launch's grid group is used instead. -#### 2.6.9 `ops/tma.cuh` +#### 2.6.9 `ops/copy_async_bulk.cuh` - + ```cpp -// include/tilefoundry/runtime/cuda/ops/tma.cuh +// include/tilefoundry/runtime/cuda/ops/copy_async_bulk.cuh template -__device__ inline void tma_copy(Src const &src, Dst &dst, uint64_t *bar); +__device__ inline void copy_async_bulk(Src const &src, Dst &dst, + uint64_t *bar); ``` -**`tma_copy`.** +**`copy_async_bulk`.** Stage a tile into shared memory and signal an mbarrier when it is readable. It is not a tier of `ops::copy_async`: there every thread issues its own load and @@ -1598,6 +1600,31 @@ tile it did not fetch. property of the layout type, so an off-grain extent is a run-time hand-off to the element path inside the same entry — same barrier, same result. +#### 2.6.10 `ops/ldmatrix.cuh` + + +```cpp +// include/tilefoundry/runtime/cuda/ops/ldmatrix.cuh +template +__device__ inline void ldmatrix(Src const &src, Dst &dst); +``` + + +**`ldmatrix`.** + +One warp loads a dense shared-memory `(16, 16)` bf16 tile into the per-lane A +fragment consumed by `ops::mma`'s current SM80 atom. + +- constraints: + - All 32 lanes issue `ldmatrix.sync.aligned.m8n8.x4.shared.b16` together. + Each lane contributes the 16-byte-aligned address selected by the source's + declared shard layout; a full-broadcast shard still uses that declaration, + not the backing allocation's incidental layout. + - The four returned registers are stored in the destination fragment's + layout order so that `ops::mma` observes the PTX register order unchanged. + - The TIR declaration fixes the source shape/dtype and destination fragment; + the runtime entry does not choose an atom or a fragment layout. + ### 2.7 `cuda/utility/` `tilefoundry::shuffle_xor`, `tilefoundry::shuffle_elect` and diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 86a9a5f3..826ffdc2 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -957,15 +957,23 @@ class CopyAsync(Op): """Effect form; async gmem→smem copy, non-blocking. Attributes: - source: input; gmem staging source. - destination: input; smem staging destination. + src: input; gmem staging source. + dst: input; smem staging destination. + smem_layout: attribute; optional landing arrangement. """ - source: Tensor - destination: Tensor + src: Tensor + dst: Tensor + smem_layout: Layout | None = None ``` - constraints: - Lowers to `tilefoundry::ops::copy_async(src, dst)`. + - `src` is gmem and `dst` is smem, with the same dtype. Each layout MUST + admit the same 4-, 8-, or 16-byte vector width and MUST walk the same tile + mode at step 1. For a `ShardLayout`, vector width is read from the whole + tile arrangement; its other strides ensure every participant's start is + aligned. Two split layouts compare their tile modes only when their mesh + and shard attrs are identical. - A later read of `dst` is ordered by `CpAsyncCommit` followed by `CpAsyncWait`. @@ -994,7 +1002,7 @@ class CpAsyncWait(Op): - `n` is a non-negative compile-time count. - `n = 0` drains every outstanding committed group. -##### TmaCopy +##### CopyAsyncBulk A staging copy whose completion lands on an mbarrier, and not a tier of `CopyAsync`: there every thread issues its own load and a commit closes the @@ -1013,7 +1021,7 @@ place of a size, which is a different operand list rather than a different tier, and are outside this op. ```python -class TmaCopy(Op): +class CopyAsyncBulk(Op): """Effect form; gmem→smem staging copy completing on an mbarrier. Attributes: @@ -1036,11 +1044,54 @@ class TmaCopy(Op): transferred bytes is the implementation's, issued on the same instruction as the copy; a caller pairing this with its own `MBarrierArriveExpectTx` would be declaring a count the op already knows. - - Lowers to `tilefoundry::ops::tma_copy(src, dst, bar)` + - Lowers to `tilefoundry::ops::copy_async_bulk(src, dst, bar)` ([runtime §2.6](./runtime.md#26-cudaops)). `barrier` is a tensor here because that is what TIR names a piece of shared memory with, and a word to the runtime, so the emitted call hands over the word's own address. +##### CopyAsyncTensor + +`T.copy_async_tensor` declares the SM90 tensor-map form separately from +`CopyAsyncBulk`: its operands describe a tensor-map global layout and the box +landed in shared memory, rather than carrying an explicit mbarrier operand. + +```python +class CopyAsyncTensor(Op): + src: Tensor + dst: Tensor + smem_layout: Layout | None = None + scope: Mesh | None = None +``` + +- constraints: + - Exactly one end is gmem and one is smem; dtype and logical shape agree. + - The global end is a static tensor-map layout of at most five dimensions, + with one contiguous mode and every other byte stride a multiple of 16. + The shared end is an at-most-five-dimensional box, each extent at most + 256, optionally using a 32-, 64-, or 128-byte TMA swizzle. + - Both ends walk the same tile modes. The issuing scope is one aligned warp. + - The declaration requires the target's `tma` capability. CUDA codegen MUST + reject it until host-encoded tensor-map construction exists; this stage + does not silently lower it to another copy instruction. + +##### LdMatrix + +```python +class LdMatrix(Op): + src: Tensor + dst: Tensor + scope: Mesh | None = None +``` + +- constraints: + - `src` is a shared-memory `(16, 16)` bf16 tile and `dst` is exactly the + register A fragment declared by `T.cuda.sm80.Mma()`; their dtype and shape + agree under the ordinary `Copy` verifier. + - One canonical SM80 warp issues the operation. It requires the target's + `tensor_core` capability and lowers to `tilefoundry::ops::ldmatrix`. + - The destination layout is the atom's declaration, not a caller-selectable + `rmem_layout` attribute. + #### Barrier object Ops (`tir.sync.mbarrier_*`) A Hopper mbarrier is a 64-bit shared-memory word carrying an arrival count, a @@ -1059,9 +1110,9 @@ generic-to-shared conversion the instruction takes — they name `.shared::cta` explicitly rather than leaving the assembler to redo that window conversion on every use. -The group is what a `TmaCopy` ring needs and no more: arm the word, arrive on it +The group is what a `CopyAsyncBulk` ring needs and no more: arm the word, arrive on it declaring bytes, wait on its phase, release it. A bare `mbarrier.arrive` is -absent because `ops::tma_copy` issues its own for the strided tier, and a bare +absent because `ops::copy_async_bulk` issues its own for the strided tier, and a bare `mbarrier.expect_tx` because nothing pairs with it. ##### MBarrierInit @@ -1110,12 +1161,12 @@ class MBarrierArriveExpectTx(Op): - `tx_bytes` MUST equal the bytes the paired copy delivers. A phase expecting a different count never completes, and that failure presents as a hang rather than as a wrong value. - - This is not paired with a `TmaCopy`, which declares its own bytes on the + - This is not paired with a `CopyAsyncBulk`, which declares its own bytes on the instruction that issues the copy. It belongs to a producer issuing one itself. - Lowers to `mbarrier.arrive.expect_tx.shared::cta.b64`, with the arrival token discarded: consumers wait on the phase parity, not on a token handed - between threads. The runtime publishes no entry for it; `ops::tma_copy` + between threads. The runtime publishes no entry for it; `ops::copy_async_bulk` writes its own for the bulk tier. ##### MBarrierWaitParity diff --git a/include/tilefoundry/runtime/cuda/ops/copy_async_bulk.cuh b/include/tilefoundry/runtime/cuda/ops/copy_async_bulk.cuh new file mode 100644 index 00000000..ece70217 --- /dev/null +++ b/include/tilefoundry/runtime/cuda/ops/copy_async_bulk.cuh @@ -0,0 +1,12 @@ +/// CUDA bulk asynchronous copy public entry. +#pragma once + +#include "copy_async_bulk/copy_async_bulk_impl.h" + +/// Stage ``src`` into ``dst``, completing on ``bar``. +template +__device__ inline void copy_async_bulk(Src const &src, Dst &dst, + uint64_t *bar) { + copy_async_bulk_impl::check_operands(); + copy_async_bulk_impl::Bulk{}(src, dst, bar); +} diff --git a/include/tilefoundry/runtime/cuda/ops/tma/tma_impl.h b/include/tilefoundry/runtime/cuda/ops/copy_async_bulk/copy_async_bulk_impl.h similarity index 74% rename from include/tilefoundry/runtime/cuda/ops/tma/tma_impl.h rename to include/tilefoundry/runtime/cuda/ops/copy_async_bulk/copy_async_bulk_impl.h index dde4b0bb..027862c5 100644 --- a/include/tilefoundry/runtime/cuda/ops/tma/tma_impl.h +++ b/include/tilefoundry/runtime/cuda/ops/copy_async_bulk/copy_async_bulk_impl.h @@ -1,14 +1,12 @@ -/// TMA op internals. Included in-context from ``ops/tma.cuh``. +/// CUDA bulk asynchronous copy internals. #pragma once -namespace tma_impl { +namespace copy_async_bulk_impl { -/// The shared-window address of a generic pointer. __device__ inline uint32_t smem_addr(void const *ptr) { return static_cast(__cvta_generic_to_shared(ptr)); } -/// Whether an operand is one unbroken run of bytes. template inline constexpr bool one_run_v = [] { using L = typename cute::remove_cvref_t::layout_type; @@ -20,7 +18,6 @@ template using elem_t = cute::remove_cvref_t())(0))>; -/// Whether an operand leaves the tile whole on every instance of its mesh. template CUTE_HOST_DEVICE constexpr bool leaves_tile_whole() { if constexpr (tilefoundry::ShardTensorLike) return detail::shard_layout_is_full_broadcast< @@ -29,29 +26,28 @@ template CUTE_HOST_DEVICE constexpr bool leaves_tile_whole() { return true; } -/// What this op needs of its operands, asked once at the entry. template -CUTE_HOST_DEVICE constexpr void check_tma_operands() { +CUTE_HOST_DEVICE constexpr void check_operands() { using s_view = tilefoundry::local_view_t; using d_view = tilefoundry::local_view_t; - static_assert( - leaves_tile_whole() && leaves_tile_whole(), - "ops::tma_copy: both operands must leave the tile whole on every " - "instance"); + static_assert(leaves_tile_whole() && leaves_tile_whole(), + "ops::copy_async_bulk: both operands must leave the tile " + "whole on every " + "instance"); static_assert(tilefoundry::ShardTensorLike, - "ops::tma_copy: the destination must name a mesh"); - static_assert( - one_run_v && one_run_v, - "ops::tma_copy: both projected views must be one static unbroken run"); - static_assert(std::is_same_v, elem_t>, - "ops::tma_copy: both operands need the same element type"); + "ops::copy_async_bulk: the destination must name a mesh"); + static_assert(one_run_v && one_run_v, + "ops::copy_async_bulk: both projected views must be one " + "static unbroken run"); static_assert( - copy_impl::same_slice_size(), - "ops::tma_copy: the two projected views must hold the same number of " - "elements"); + std::is_same_v, elem_t>, + "ops::copy_async_bulk: both operands need the same element type"); + static_assert(copy_impl::same_slice_size(), + "ops::copy_async_bulk: the two projected views must hold the " + "same number of " + "elements"); } -/// The element loop `Bulk` hands off to: every instance strides the one run. template struct StridedCopy { template __device__ void operator()(SV const &sv, DV &dv) const { @@ -63,7 +59,6 @@ template struct StridedCopy { } }; -/// Every thread copies its share, then one arrival says the tile is readable. struct Strided { template __device__ void operator()(Src const &src, Dst &dst, uint64_t *bar) const { @@ -80,7 +75,6 @@ struct Strided { } }; -/// ``cp.async.bulk`` global to shared, completing on the barrier. struct Bulk { template __device__ void operator()(Src const &src, Dst &dst, uint64_t *bar) const { diff --git a/include/tilefoundry/runtime/cuda/ops/ldmatrix.cuh b/include/tilefoundry/runtime/cuda/ops/ldmatrix.cuh new file mode 100644 index 00000000..579a5e49 --- /dev/null +++ b/include/tilefoundry/runtime/cuda/ops/ldmatrix.cuh @@ -0,0 +1,10 @@ +/// CUDA ldmatrix public entry. +#pragma once + +#include "ldmatrix/ldmatrix_impl.h" + +/// Load one warp's shared-memory tile into its SM80 MMA A fragment. +template +__device__ inline void ldmatrix(Src const &src, Dst &dst) { + ldmatrix_impl::LdMatrix{}(src, dst); +} diff --git a/include/tilefoundry/runtime/cuda/ops/ldmatrix/ldmatrix_impl.h b/include/tilefoundry/runtime/cuda/ops/ldmatrix/ldmatrix_impl.h new file mode 100644 index 00000000..764a83a5 --- /dev/null +++ b/include/tilefoundry/runtime/cuda/ops/ldmatrix/ldmatrix_impl.h @@ -0,0 +1,47 @@ +/// CUDA ldmatrix implementation. Included in-context from ops/ldmatrix.cuh +/// inside namespace tilefoundry::ops. +#pragma once + +namespace ldmatrix_impl { + +__device__ inline uint32_t smem_addr(void const *ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +template +__device__ decltype(auto) source_at(Src const &src, int row, int col) { + if constexpr (tilefoundry::ShardTensorLike) { + auto whole = + cute::make_tensor(src.data(), src.shard_layout.layout_value); + return whole(row, col); + } else { + return src(row, col); + } +} + +struct LdMatrix { + template + __device__ void operator()(Src const &src, Dst &dst) const { + auto d = dst.data(); + const int lane = + int(tilefoundry::program_id()) & + 31; + const int row = lane & 15; + const int col = (lane >> 4) * 8; + uint32_t r0, r1, r2, r3; + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(smem_addr(&source_at(src, row, col)))); + + /// Mma's A fragment is stored in layout order. Its runtime packs that + /// order as [0, 1], [4, 5], [2, 3], [6, 7], so place the four PTX + /// registers where that packing reads them back in r0..r3 order. + __builtin_memcpy(&d[0], &r0, sizeof(r0)); + __builtin_memcpy(&d[4], &r1, sizeof(r1)); + __builtin_memcpy(&d[2], &r2, sizeof(r2)); + __builtin_memcpy(&d[6], &r3, sizeof(r3)); + } +}; + +} diff --git a/include/tilefoundry/runtime/cuda/ops/tma.cuh b/include/tilefoundry/runtime/cuda/ops/tma.cuh deleted file mode 100644 index abf2490e..00000000 --- a/include/tilefoundry/runtime/cuda/ops/tma.cuh +++ /dev/null @@ -1,13 +0,0 @@ -/// tilefoundry TMA op — one public entry, the operands checked against it. - -/// TMA derives addresses and byte counts from tensor layouts. -#pragma once - -#include "tma/tma_impl.h" - -/// Stage ``src`` into ``dst``, completing on ``bar``. -template -__device__ inline void tma_copy(Src const &src, Dst &dst, uint64_t *bar) { - tma_impl::check_tma_operands(); - tma_impl::Bulk{}(src, dst, bar); -} diff --git a/include/tilefoundry/runtime/cuda/runtime.cuh b/include/tilefoundry/runtime/cuda/runtime.cuh index e930dd2b..35d07eb1 100644 --- a/include/tilefoundry/runtime/cuda/runtime.cuh +++ b/include/tilefoundry/runtime/cuda/runtime.cuh @@ -200,7 +200,8 @@ namespace ops { #include "ops/elementwise.cuh" #include "ops/copy.cuh" -#include "ops/tma.cuh" +#include "ops/copy_async_bulk.cuh" +#include "ops/ldmatrix.cuh" #include "ops/reduce.cuh" /// dot after reduce: it reuses reduce's no-workspace tag. #include "ops/dot.cuh" diff --git a/scripts/runtime_spec_surface.py b/scripts/runtime_spec_surface.py index fbfc3924..ba99846d 100644 --- a/scripts/runtime_spec_surface.py +++ b/scripts/runtime_spec_surface.py @@ -54,7 +54,8 @@ "ops-reduce": "cuda/ops/reduce.cuh", "ops-rmsnorm": "cuda/ops/rmsnorm.cuh", "ops-sync": "cuda/ops/sync.cuh", - "ops-tma": "cuda/ops/tma.cuh", + "ops-copy-async-bulk": "cuda/ops/copy_async_bulk.cuh", + "ops-ldmatrix": "cuda/ops/ldmatrix.cuh", "utility-warp": "cuda/utility/warp.cuh", } diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tma.py b/src/tilefoundry/codegen/cuda/tir/memory/copy_async_bulk.py similarity index 69% rename from src/tilefoundry/codegen/cuda/tir/memory/tma.py rename to src/tilefoundry/codegen/cuda/tir/memory/copy_async_bulk.py index 5f414c6f..a8d744b2 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tma.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/copy_async_bulk.py @@ -1,24 +1,22 @@ -"""Emitter for ``TmaCopy`` — one line, whichever instruction ends up running.""" +"""Emit ``CopyAsyncBulk`` through its one runtime entry.""" from __future__ import annotations from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.cuda.tir.mbarrier import barrier_word -from tilefoundry.ir.tir.cuda.memory.tma import TmaCopy +from tilefoundry.ir.tir.cuda.memory.copy_async_bulk import CopyAsyncBulk from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen -_TMA_COPY = "tilefoundry::ops::tma_copy" - def _tensor_expr(var, ctx: CudaCodegenContext) -> str: base = ctx.name_for(var) return f"{base}_tensor" if ctx.is_kernel_param(var) else base -@register_codegen(CudaTarget, Role.EMIT, TmaCopy) +@register_codegen(CudaTarget, Role.EMIT, CopyAsyncBulk) def _emit(call, ctx: CudaCodegenContext) -> None: src = _tensor_expr(call.args[0], ctx) dst = _tensor_expr(call.args[1], ctx) bar = f"reinterpret_cast({barrier_word(call.args[2], ctx)})" - ctx.emit(f"{_TMA_COPY}({src}, {dst}, {bar});") + ctx.emit(f"tilefoundry::ops::copy_async_bulk({src}, {dst}, {bar});") diff --git a/src/tilefoundry/codegen/cuda/tir/memory/copy_async_tensor.py b/src/tilefoundry/codegen/cuda/tir/memory/copy_async_tensor.py new file mode 100644 index 00000000..1ee2ce2f --- /dev/null +++ b/src/tilefoundry/codegen/cuda/tir/memory/copy_async_tensor.py @@ -0,0 +1,15 @@ +"""Refuse tensor-map emission until host tensor-map construction exists.""" + +from __future__ import annotations + +from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.ir.tir.cuda.memory.copy_async_tensor import CopyAsyncTensor +from tilefoundry.target import CudaTarget +from tilefoundry.visitor_registry.registries import Role, register_codegen + + +@register_codegen(CudaTarget, Role.EMIT, CopyAsyncTensor) +def _emit(call, ctx: CudaCodegenContext) -> None: + raise RuntimeError( + "tir.cuda.memory.CopyAsyncTensor: no emitter without host-encoded tensor maps" + ) diff --git a/src/tilefoundry/codegen/cuda/tir/memory/ldmatrix.py b/src/tilefoundry/codegen/cuda/tir/memory/ldmatrix.py new file mode 100644 index 00000000..0e15931f --- /dev/null +++ b/src/tilefoundry/codegen/cuda/tir/memory/ldmatrix.py @@ -0,0 +1,20 @@ +"""Emit the SM80 warp-cooperative shared-memory matrix load.""" + +from __future__ import annotations + +from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.ir.tir.cuda.memory.ldmatrix import LdMatrix +from tilefoundry.target import CudaTarget +from tilefoundry.visitor_registry.registries import Role, register_codegen + + +def _tensor_expr(var, ctx: CudaCodegenContext) -> str: + base = ctx.name_for(var) + return f"{base}_tensor" if ctx.is_kernel_param(var) else base + + +@register_codegen(CudaTarget, Role.EMIT, LdMatrix) +def _emit(call, ctx: CudaCodegenContext) -> None: + src = _tensor_expr(call.args[0], ctx) + dst = _tensor_expr(call.args[1], ctx) + ctx.emit(f"tilefoundry::ops::ldmatrix({src}, {dst});") diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py index 5f4858f8..f8754474 100644 --- a/src/tilefoundry/ir/pattern/__init__.py +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -45,6 +45,7 @@ SwizzlePattern, Tensor, TensorPattern, + VectorPattern, WildcardPattern, ) from .utils import ( @@ -57,6 +58,7 @@ locate_dim_var, moved_tile, storage_place, + vector, ) __all__ = [ @@ -110,4 +112,6 @@ "relations_of", "resolved", "storage_place", + "VectorPattern", + "vector", ] diff --git a/src/tilefoundry/ir/pattern/constraint.py b/src/tilefoundry/ir/pattern/constraint.py index 1d41797b..732cba5f 100644 --- a/src/tilefoundry/ir/pattern/constraint.py +++ b/src/tilefoundry/ir/pattern/constraint.py @@ -91,8 +91,8 @@ def pair(self, operands: dict): return None if any(not isinstance(value, TensorType) for value in held) else held @staticmethod - def reading(tensor: TensorType): - layout = affine_part(tensor.layout) + def reading(tensor: TensorType, arrangement=None): + layout = affine_part(tensor.layout if arrangement is None else arrangement) if layout is None: return None, f"{tensor.layout!r} is no strided arrangement" shape, strides = tuple(layout.shape), tuple(layout.strides) @@ -117,11 +117,28 @@ def reading(tensor: TensorType): return None, f"{layout!r} has {len(unit)} modes at step 1, not one" return unit[0], None + def readings(self, pair: tuple[TensorType, TensorType]): + """Read below one shared shard frame only after proving it is the same.""" + left, right = (value.layout for value in pair) + if ( + isinstance(left, ShardLayout) + and isinstance(right, ShardLayout) + and left.attrs == right.attrs + and left.mesh == right.mesh + ): + arrangements = (left.layout, right.layout) + else: + arrangements = (None, None) + return tuple( + self.reading(value, arrangement) + for value, arrangement in zip(pair, arrangements) + ) + def holds(self, operands: dict) -> bool: pair = self.pair(operands) if pair is None: return True - (left, _), (right, _) = (self.reading(value) for value in pair) + (left, _), (right, _) = self.readings(pair) return ( left is not None and right is not None and left[0] == right[0] and left[2] and right[2] ) @@ -137,7 +154,8 @@ def refused(self, operands: dict) -> str: if pair is None: return self.written() readings = tuple( - (name, *self.reading(value)) for name, value in zip((self.left, self.right), pair) + (name, *reading) + for name, reading in zip((self.left, self.right), self.readings(pair)) ) for name, _, why in readings: if why is not None: diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index 62e60a8e..82a8d4cd 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -17,7 +17,12 @@ ) from tilefoundry.ir.types.int_tuple import congruent from tilefoundry.ir.types.layout import flatten -from tilefoundry.ir.types.layout_algebra import frame_of, is_inverse_projectable +from tilefoundry.ir.types.layout_algebra import ( + ASYNC_WIDTHS, + frame_of, + is_inverse_projectable, + vector_widths, +) from tilefoundry.ir.types.mesh import separate from .match import ( @@ -76,6 +81,49 @@ def describe(self, name: str = UNNAMED_PLACE) -> str: return name +VECTOR_READING = ( + "every run: each tile axis's modes walked fastest first, contiguous ones joined; " + "the run at step 1 and every other step a whole number of vectors" +) + + +@dataclass(frozen=True) +class VectorPattern(Pattern): + """An arrangement that moves whole 4-, 8-, or 16-byte vectors.""" + + width: CapturePattern + dtype: str + + def widths(self, subject, captures) -> tuple[int, ...]: + bits = getattr(dict(captures or {}).get(self.dtype), "bit_width", None) + return () if type(bits) is not int else vector_widths(subject, bits) + + def match(self, subject, captures=None): + held = dict(captures or {}) + widths = self.widths(subject, held) + if not widths: + return None + if self.width.name in held: + return Match(held) if held[self.width.name] in widths else None + return matched(self.width, widths[-1], held) + + def refusal(self, subject, captures=None) -> str | None: + if self.match(subject, captures) is not None: + return None + sizes = ", ".join(map(str, ASYNC_WIDTHS[:-1])) + f" or {ASYNC_WIDTHS[-1]}" + return ( + f"{subject!r} moves no whole vector of {sizes} bytes -- its run at step 1 " + "and every other step are no whole number of one -- so the two ends share " + "no run wide enough for cp.async" + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return f"vectors of {self.width.name} bytes" + + def relations(self) -> tuple[str, ...]: + return (VECTOR_READING, *relations_of((self.width,))) + + @dataclass(frozen=True, init=False) class OrPattern(Pattern): patterns: tuple diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py index 3e361e9e..5aa179ca 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -4,6 +4,7 @@ from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.types import ComposedLayout, Mesh, StorageKind, Swizzle +from tilefoundry.ir.types.layout_algebra import ASYNC_WIDTHS from .pattern import ( AttrPattern, @@ -18,6 +19,7 @@ RangePattern, SwizzlePattern, TensorPattern, + VectorPattern, WildcardPattern, ) @@ -49,6 +51,14 @@ def storage_place(index: int) -> str: return f"storage{index}" +def vector(index: int) -> VectorPattern: + """End *index* of a cp.async, counted in that end's element dtype.""" + return VectorPattern( + CapturePattern("width", OneOfPattern(ASYNC_WIDTHS)), + dtype_place(index), + ) + + _ANY_THREADS = OrPattern( ComposedLayoutPattern( offset=WildcardPattern(), @@ -139,4 +149,5 @@ def _mangle_variant_name(name: str, specializations: tuple[Pattern, ...]) -> str "locate_dim_var", "moved_tile", "storage_place", + "vector", ] diff --git a/src/tilefoundry/ir/tir/async_copy.py b/src/tilefoundry/ir/tir/async_copy.py index d4b204e8..b956c965 100644 --- a/src/tilefoundry/ir/tir/async_copy.py +++ b/src/tilefoundry/ir/tir/async_copy.py @@ -3,11 +3,18 @@ from __future__ import annotations from tilefoundry.ir.core import Op -from tilefoundry.ir.core.param_def import ParamDef +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Tensor -from tilefoundry.ir.types import UnitType -from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.ir.pattern import ( + DistinctConstraint, + SameModesConstraint, + any_threads, + moved_tile, + vector, +) +from tilefoundry.ir.tir.verify import verify_between, verify_operands +from tilefoundry.ir.types import Layout, UnitType +from tilefoundry.ir.types.storage import StorageKind as S from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -15,8 +22,27 @@ class CopyAsync(Op): """Async gmem→smem copy (``cp.async.cg.shared.global``); non-blocking.""" - src = ParamDef(kind="input", pattern=Tensor) - dst = ParamDef(kind="input", pattern=Tensor) + src = ParamDef( + kind="input", + effect=MemoryEffect.READ, + pattern=moved_tile(0, S.GMEM, vector(0)), + ) + dst = ParamDef( + kind="input", + effect=MemoryEffect.WRITE, + pattern=moved_tile(1, S.SMEM, vector(1)), + ) + between = ( + DistinctConstraint("storage", "src", "dst"), + SameModesConstraint("src", "dst"), + ) + smem_layout = ParamDef( + kind="attribute", + annotation=Layout, + optional=True, + default=None, + ) + scope = any_threads() @register_typeinfer(CopyAsync) @@ -25,15 +51,18 @@ def _(call: "Call", ctx: "TypeInferContext") -> UnitType: @register_verify_stmt(CopyAsync) -def _(call: "Call", ctx: "VerifyContext") -> None: +def verify_copy_async(call: "Call", ctx: "VerifyContext") -> None: + """Run native storage checks before layout relations and vector patterns.""" src = ctx.type_of(call.args[0]) dst = ctx.type_of(call.args[1]) - if dst.storage != StorageKind.SMEM: + if dst.storage != S.SMEM: ctx.error(call, f"CopyAsync destination must be smem, got {dst.storage}") - if src.storage != StorageKind.GMEM: + if src.storage != S.GMEM: ctx.error(call, f"CopyAsync source must be gmem, got {src.storage}") if src.dtype != dst.dtype: ctx.error(call, f"CopyAsync dtype mismatch: {src.dtype} vs {dst.dtype}") + verify_between(call, ctx) + verify_operands(call, ctx, "copy_async") @register_op(dialect="T", category="async", name="cp_async_commit") diff --git a/src/tilefoundry/ir/tir/cuda/memory/copy_async_bulk.py b/src/tilefoundry/ir/tir/cuda/memory/copy_async_bulk.py new file mode 100644 index 00000000..51855b52 --- /dev/null +++ b/src/tilefoundry/ir/tir/cuda/memory/copy_async_bulk.py @@ -0,0 +1,49 @@ +"""Barrier-completing gmem-to-smem bulk asynchronous copy.""" + +from __future__ import annotations + +from tilefoundry.ir.core import Op +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef +from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import Tensor +from tilefoundry.ir.types import UnitType +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt + + +@register_op(dialect="T", category="async", name="copy_async_bulk") +class CopyAsyncBulk(Op): + """Stage a tile from global to shared memory, completing on a barrier.""" + + src = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=Tensor) + dst = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=Tensor) + barrier = ParamDef( + kind="input", + effect=MemoryEffect.READ | MemoryEffect.WRITE, + pattern=Tensor, + ) + + +@register_typeinfer(CopyAsyncBulk) +def _(call: "Call", ctx: "TypeInferContext") -> UnitType: + return UnitType() + + +@register_verify_stmt(CopyAsyncBulk) +def _(call: "Call", ctx: "VerifyContext") -> None: + src = ctx.type_of(call.args[0]) + dst = ctx.type_of(call.args[1]) + bar = ctx.type_of(call.args[2]) + if src.storage != StorageKind.GMEM: + ctx.error(call, f"CopyAsyncBulk source must be gmem, got {src.storage}") + if dst.storage != StorageKind.SMEM: + ctx.error(call, f"CopyAsyncBulk destination must be smem, got {dst.storage}") + if bar.storage != StorageKind.SMEM: + ctx.error(call, f"CopyAsyncBulk barrier must be smem, got {bar.storage}") + if src.dtype != dst.dtype: + ctx.error(call, f"CopyAsyncBulk dtype mismatch: {src.dtype} vs {dst.dtype}") + if src.shape != dst.shape: + ctx.error(call, f"CopyAsyncBulk shape mismatch: {src.shape} vs {dst.shape}") + + +__all__ = ["CopyAsyncBulk"] diff --git a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py new file mode 100644 index 00000000..8958d8cb --- /dev/null +++ b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py @@ -0,0 +1,401 @@ +"""SM90 tensor-map asynchronous copy declaration and operand patterns.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + +from tilefoundry.ir.core import Op +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef +from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import ( + UNNAMED_PLACE, + AndPattern, + BitsPattern, + CapturePattern, + ComposedLayoutPattern, + DistinctConstraint, + LayoutPattern, + MeshPattern, + MultipleOfPattern, + OrPattern, + Pattern, + RangePattern, + SameModesConstraint, + SequencePattern, + SwitchPattern, + SwizzlePattern, + affine_part, + dtype_place, + matched, + moved_tile, + relations_of, + storage_place, +) +from tilefoundry.ir.pattern.match import written_place, written_tuple +from tilefoundry.ir.tir.verify import verify_between, verify_operands +from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, UnitType +from tilefoundry.ir.types.int_tuple import flatten +from tilefoundry.ir.types.layout_algebra import box_runs, frame_of +from tilefoundry.ir.types.storage import StorageKind as S +from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt + +TMA_RANK = 5 +BOX_EXTENT = 256 +TMA_UNIT_BITS = 16 * 8 +SWIZZLE_PLACE = "swizzle" +BOX_READING = ( + "every box: each tile axis's modes, contiguous ones joined up to " + f"{BOX_EXTENT} elements, one dim each, in increasing step" +) +TENSORMAP_READING = ( + "every tensormap: one dim per mode of the tile, the mode at step 1 first" +) +TMA_STORAGES = (S.GMEM, S.SMEM) + + +class TmaSwizzle(Enum): + """The shared-memory swizzle selected when the tensor map is encoded.""" + + NONE = 0 + SW32 = 32 + SW64 = 64 + SW128 = 128 + + @property + def swizzle(self) -> Swizzle | None: + return None if self is TmaSwizzle.NONE else Swizzle(self.value.bit_length() - 5, 4, 3) + + +def unframed(layout): + """Drop a shard frame only when every participant sees the whole tile.""" + if isinstance(layout, ShardLayout) and affine_part(layout) is not None: + return layout.layout + return layout + + +def _affine(subject): + framed = frame_of(unframed(subject)) + layout = framed[1] if framed is not None and framed[0] == 0 else None + if not isinstance(layout, Layout) or layout.strides is None: + return None + if any( + type(value) is not int + for group in (layout.shape, layout.strides) + for value in flatten(group) + ): + return None + return layout + + +def _missed_place(places, values, captures, first: int = 0) -> tuple | None: + held = captures + for index, (place, value) in enumerate(zip(places, values), first): + found = matched(place, value, held) + if found is None: + return index, place, value + held = found.captures + return None + + +@dataclass(frozen=True) +class BoxPattern(Pattern): + """One TMA box landed in shared memory.""" + + dims: tuple + dtype: str + span: int | None = None + + def reading(self, subject, captures) -> tuple[tuple | None, str | None]: + layout = _affine(subject) + if layout is None: + return None, f"{subject!r} is no static strided arrangement" + width = getattr(captures.get(self.dtype), "bit_width", None) + if type(width) is not int: + return None, f"the element it arranges is not bound as {self.dtype}" + runs = box_runs(layout, width, self.span) + if not runs: + return None, "it holds one element, which is no box" + if len(runs) > len(self.dims): + return None, ( + f"it is {len(runs)} runs of modes, and a box has at most " + f"{len(self.dims)} dims" + ) + extents = tuple(run.extent for run in runs) + extents += (1,) * (len(self.dims) - len(extents)) + if runs[0].step != 1: + return extents, ( + f"its smallest step is {runs[0].step}, and a box lays its dim 0 at step 1" + ) + if self.span is not None and (self.span * 8) % width: + return extents, f"a {self.span}-byte row is no whole number of {self.dtype}" + expected = runs[0].extent if self.span is None else self.span * 8 // width + for index, run in enumerate(runs[1:], 1): + if run.step != expected: + return extents, ( + f"its dim {index} steps {run.step} where a box lays it at " + f"{expected} ({self.laid()})" + ) + expected *= run.extent + return extents, None + + def laid(self) -> str: + rows = "" if self.span is None else f", rows {self.span} B apart" + return f"dim 0 fastest{rows}" + + def match(self, subject, captures=None): + held = dict(captures or {}) + extents, unlaid = self.reading(subject, held) + if extents is None or unlaid is not None: + return None + return matched(SequencePattern(*self.dims), extents, held) + + def refusal(self, subject, captures=None) -> str | None: + held = dict(captures or {}) + extents, why = self.reading(subject, held) + if extents is None: + return why + missed = _missed_place(self.dims, extents, held) + if missed is not None: + index, place, extent = missed + return ( + f"its box dim {index} holds {extent} elements, and a box reads " + f"{written_place(place.pattern, place.name)}" + ) + return why + + def describe(self, name: str = UNNAMED_PLACE) -> str: + dims = written_tuple(tuple(written_place(place) for place in self.dims)) + return f"box {dims}, {self.laid()}" + + def relations(self) -> tuple[str, ...]: + return (BOX_READING, *relations_of(self.dims)) + + +@dataclass(frozen=True, init=False) +class BoxFamily(SwitchPattern): + """Every unswizzled or swizzled shared-memory box.""" + + def match(self, subject, captures=None): + return super().match(unframed(subject), captures) + + def refusal(self, subject, captures=None) -> str | None: + layout = unframed(subject) + transform = layout.inner if isinstance(layout, ComposedLayout) else None + if transform is not None and layout.offset != 0: + return f"it is reached through {transform!r} at offset {layout.offset}, not 0" + for mode, pattern in self.branches: + if mode.swizzle == transform: + inner = pattern if transform is None else pattern.outer + return inner.refusal( + layout if transform is None else layout.outer, + captures, + ) + written = ", ".join( + repr(mode.swizzle) for mode, _ in self.branches if mode.swizzle is not None + ) + return f"it is reached through {transform!r}, and a tensormap swizzles by {written} or not at all" + + +@dataclass(frozen=True) +class TensorMapPattern(Pattern): + """The global tile described by one tensor map.""" + + steps: tuple + dtype: str + + def reading(self, subject) -> tuple[tuple | None, str | None]: + layout = _affine(subject) + if layout is None: + return None, f"{subject!r} is no static strided tensor a tensormap describes" + modes = [ + (extent, step) + for extents, steps in zip(layout.shape, layout.strides) + for extent, step in zip(flatten(extents), flatten(steps)) + if extent > 1 + ] + unit = [mode for mode in modes if mode[1] == 1] + if len(unit) != 1: + return None, ( + f"{len(unit)} of its modes step 1, and a tensormap's dim 0 is its one " + "contiguous mode" + ) + if len(modes) > len(self.steps) + 1: + return None, ( + f"it is {len(modes)} modes, and a tensormap has at most " + f"{len(self.steps) + 1} dims" + ) + others = tuple(step for _, step in modes if step != 1) + return others + (0,) * (len(self.steps) - len(others)), None + + def match(self, subject, captures=None): + steps, _ = self.reading(subject) + return None if steps is None else matched(SequencePattern(*self.steps), steps, captures) + + def refusal(self, subject, captures=None) -> str | None: + held = dict(captures or {}) + steps, why = self.reading(subject) + if steps is None: + return why + missed = _missed_place(self.steps, steps, held, first=1) + if missed is not None: + index, place, step = missed + return ( + f"its dim {index} steps {step} elements, and a tensormap reads " + f"{written_place(place.pattern, place.name)}" + ) + return None + + def describe(self, name: str = UNNAMED_PLACE) -> str: + steps = written_tuple(("1", *(written_place(place) for place in self.steps))) + return f"tensormap at {steps}" + + def relations(self) -> tuple[str, ...]: + return (TENSORMAP_READING, *relations_of(self.steps)) + + +def _dim(name: str, dtype: str, *, span: int | None = None) -> CapturePattern: + parts = [ + RangePattern(lo=1, hi=BOX_EXTENT), + BitsPattern(dtype, MultipleOfPattern(TMA_UNIT_BITS)), + ] + if span is not None: + parts.append(BitsPattern(dtype, RangePattern(hi=span * 8))) + return CapturePattern(name, AndPattern(tuple(parts))) + + +def TmaBoxPattern(dtype: str) -> BoxFamily: + rest = tuple( + CapturePattern(f"dim{index}", RangePattern(lo=1, hi=BOX_EXTENT)) + for index in range(1, TMA_RANK) + ) + boxes = {TmaSwizzle.NONE: BoxPattern((_dim("dim0", dtype), *rest), dtype)} + for index, mode in enumerate( + (mode for mode in TmaSwizzle if mode.swizzle is not None), + TMA_RANK, + ): + swizzle = mode.swizzle + boxes[mode] = ComposedLayoutPattern( + SwizzlePattern(swizzle.bits, swizzle.base, swizzle.shift), + 0, + BoxPattern((_dim(f"dim{index}", dtype, span=mode.value), *rest), dtype, mode.value), + ) + return BoxFamily(SWIZZLE_PLACE, boxes) + + +def TmaGlobalPattern(dtype: str) -> TensorMapPattern: + return TensorMapPattern( + tuple( + CapturePattern( + f"step{index}", + BitsPattern(dtype, MultipleOfPattern(TMA_UNIT_BITS)), + ) + for index in range(1, TMA_RANK) + ), + dtype, + ) + + +def TmaOperandPattern(storage: str, dtype: str) -> SwitchPattern: + return SwitchPattern( + storage, + { + S.GMEM: TmaGlobalPattern(dtype), + S.SMEM: TmaBoxPattern(dtype), + }, + ) + + +def _warp_scope() -> MeshPattern: + layout = LayoutPattern(((32,),), ((1,),), per_mode=True) + sliced = ComposedLayoutPattern( + offset=CapturePattern("p0", MultipleOfPattern(32)), + outer=layout, + ) + return MeshPattern(("thread",), OrPattern(sliced, layout)) + + +@register_op(dialect="T", category="async", name="copy_async_tensor") +class CopyAsyncTensor(Op): + """Move one tensor-map box between global and shared memory.""" + + capability = "tma" + + src = ParamDef( + kind="input", + effect=MemoryEffect.READ, + pattern=moved_tile( + 0, + TMA_STORAGES, + TmaOperandPattern(storage_place(0), dtype_place(0)), + ), + ) + dst = ParamDef( + kind="input", + effect=MemoryEffect.WRITE, + pattern=moved_tile( + 1, + TMA_STORAGES, + TmaOperandPattern(storage_place(1), dtype_place(1)), + ), + ) + between = ( + DistinctConstraint("storage", "src", "dst"), + SameModesConstraint("src", "dst"), + ) + smem_layout = ParamDef( + kind="attribute", + annotation=Layout, + optional=True, + default=None, + ) + scope = ParamDef( + kind="attribute", + annotation=Mesh, + pattern=_warp_scope(), + optional=True, + default=None, + ) + + +@register_typeinfer(CopyAsyncTensor) +def _(call: "Call", ctx: "TypeInferContext") -> UnitType: + return UnitType() + + +@register_verify_stmt(CopyAsyncTensor) +def verify_copy_async_tensor(call: "Call", ctx: "VerifyContext") -> None: + src, dst = tuple(ctx.type_of(arg) for arg in call.args) + if tuple(src.shape) != tuple(dst.shape) or src.dtype != dst.dtype: + ctx.error( + call, + f"copy_async_tensor moves one tile: src is {tuple(src.shape)} " + f"{src.dtype.name} and dst is {tuple(dst.shape)} {dst.dtype.name}", + ) + moves = " and ".join(map(str, TMA_STORAGES)) + verify_between( + call, + ctx, + f"copy_async_tensor moves a tile between {moves}, one end each: ", + ) + verify_operands(call, ctx, "copy_async_tensor") + + +__all__ = [ + "BOX_EXTENT", + "BOX_READING", + "BoxFamily", + "BoxPattern", + "CopyAsyncTensor", + "SWIZZLE_PLACE", + "TENSORMAP_READING", + "TMA_RANK", + "TMA_STORAGES", + "TMA_UNIT_BITS", + "TensorMapPattern", + "TmaBoxPattern", + "TmaGlobalPattern", + "TmaOperandPattern", + "TmaSwizzle", + "unframed", +] diff --git a/src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py b/src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py new file mode 100644 index 00000000..f7f1609e --- /dev/null +++ b/src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py @@ -0,0 +1,38 @@ +"""SM80 warp-cooperative shared-memory matrix load.""" + +from __future__ import annotations + +from tilefoundry.ir.core import Op +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef +from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import moved_tile +from tilefoundry.ir.tir.cuda.nn.sm80_mma import Mma +from tilefoundry.ir.tir.memory.copy import Copy +from tilefoundry.ir.types import Mesh +from tilefoundry.ir.types.storage import StorageKind as S +from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt +from tilefoundry.visitor_registry.registries import typeinfer_registry, verify_stmt_registry + + +@register_op(dialect="T", category="memory", name="ldmatrix") +class LdMatrix(Op): + """Load one warp's shared-memory tile into the SM80 MMA A fragment.""" + + capability = "tensor_core" + + src = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=moved_tile(0, S.SMEM)) + dst = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=Mma.A) + scope = ParamDef( + kind="attribute", + annotation=Mesh, + pattern=Mma.scope_pattern(), + optional=True, + default=None, + ) + + +register_typeinfer(LdMatrix)(typeinfer_registry.lookup(Copy)) +register_verify_stmt(LdMatrix)(verify_stmt_registry.lookup(Copy)) + + +__all__ = ["LdMatrix"] diff --git a/src/tilefoundry/ir/tir/cuda/memory/tma.py b/src/tilefoundry/ir/tir/cuda/memory/tma.py deleted file mode 100644 index 107c1e23..00000000 --- a/src/tilefoundry/ir/tir/cuda/memory/tma.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Effect-form TIR Op for a barrier-completing gmem→smem staging copy. - -See [runtime §2.6](docs/spec/runtime.md#26-cudaops). -See [tir §2.3](docs/spec/tir.md#23-tir-ops). -""" - -from __future__ import annotations - -from tilefoundry.ir.core import Op -from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Tensor -from tilefoundry.ir.types import UnitType -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt - -__all__ = ["TmaCopy"] - - -@register_op(dialect="T", category="async", name="tma_copy") -class TmaCopy(Op): - """Stage a tile from global to shared memory, completing on a barrier.""" - - src = ParamDef(kind="input", pattern=Tensor) - dst = ParamDef(kind="input", pattern=Tensor) - barrier = ParamDef(kind="input", pattern=Tensor) - - -@register_typeinfer(TmaCopy) -def _(call: "Call", ctx: "TypeInferContext") -> UnitType: - return UnitType() - - -@register_verify_stmt(TmaCopy) -def _(call: "Call", ctx: "VerifyContext") -> None: - src = ctx.type_of(call.args[0]) - dst = ctx.type_of(call.args[1]) - bar = ctx.type_of(call.args[2]) - if src.storage != StorageKind.GMEM: - ctx.error(call, f"TmaCopy source must be gmem, got {src.storage}") - if dst.storage != StorageKind.SMEM: - ctx.error(call, f"TmaCopy destination must be smem, got {dst.storage}") - if bar.storage != StorageKind.SMEM: - ctx.error(call, f"TmaCopy barrier must be smem, got {bar.storage}") - if src.dtype != dst.dtype: - ctx.error(call, f"TmaCopy dtype mismatch: {src.dtype} vs {dst.dtype}") - if src.shape != dst.shape: - ctx.error(call, f"TmaCopy shape mismatch: {src.shape} vs {dst.shape}") diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index c1c18b1e..a2f18f7b 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -10,11 +10,13 @@ from __future__ import annotations +from dataclasses import dataclass, replace from typing import Optional, Union from tilefoundry.ir.types.layout import flatten from .layout import ComposedLayout, Layout, Swizzle, size +from .shard_layout import ShardLayout from .stride import compact_col_major, idx2crd @@ -22,6 +24,19 @@ class NotProjectable(ValueError): """A layout cannot serve as a mesh execution scope (not inverse-projectable).""" +ASYNC_WIDTHS = (4, 8, 16) + + +@dataclass(frozen=True) +class Run: + """One contiguous run of modes from one logical tile axis.""" + + extent: int + step: int + axis: int + mode: int + + def _shape(layout: Layout) -> tuple[int, ...]: return flatten(layout.shape) @@ -132,6 +147,67 @@ def frame_of(layout: Union[Layout, ComposedLayout]) -> tuple[int, Layout] | None return layout.offset, layout.outer +def box_runs( + layout: Layout, + element_bits: int, + span: int | None, + limit: int | None = 256, +) -> tuple[Run, ...]: + """Read contiguous runs by tile axis, ordered by increasing step.""" + runs: list[Run] = [] + for axis, (extents, steps) in enumerate(zip(layout.shape, layout.strides)): + modes = tuple(enumerate(zip(flatten(extents), flatten(steps)))) + for mode, (extent, step) in reversed(modes): + if extent == 1: + continue + last = runs[-1] if runs and runs[-1].axis == axis else None + joined = None if last is None else last.extent * extent + if ( + last is not None + and step == last.step * last.extent + and (limit is None or joined <= limit) + and not ( + span is not None + and last.step == 1 + and joined * element_bits > span * 8 + ) + ): + runs[-1] = replace(last, extent=joined) + else: + runs.append(Run(extent, step, axis, mode)) + return tuple(sorted(runs, key=lambda run: run.step)) + + +def vector_widths(layout, element_bits: int) -> tuple[int, ...]: + """Every cp.async width that divides every run in an arrangement.""" + from tilefoundry.ir.pattern.constraint import affine_part # noqa: PLC0415 + + widest = ASYNC_WIDTHS[-1] + if isinstance(layout, ShardLayout): + layout = layout.layout + inner = getattr(layout, "inner", None) + if inner is not None and hasattr(inner, "base"): + widest = min(widest, 1 << inner.base) + held = affine_part(layout) + if held is None or any( + type(value) is not int + for group in (held.shape, held.strides) + for value in flatten(group) + ): + return () + runs = box_runs(held, element_bits, None, limit=None) + unit = [run.extent for run in runs if run.step == 1] + if len(unit) != 1: + return () + counted = (unit[0], *(run.step for run in runs if run.step != 1)) + return tuple( + width + for width in ASYNC_WIDTHS + if width <= widest + and all(value * element_bits % (width * 8) == 0 for value in counted) + ) + + def complement(layout: Layout, max_idx: int = 1) -> Layout: """CuTe ``complement``: the modes that fill the gaps below ``max_idx``.""" result_shape: list[int] = [] @@ -407,13 +483,16 @@ def contains(scope: ComposedLayout, t: int) -> bool: __all__ = [ + "ASYNC_WIDTHS", "NotProjectable", + "Run", "swizzle_of", "composition", "cosize", "apply", "coalesce", "frame_of", + "box_runs", "complement", "is_inverse_projectable", "right_inverse", @@ -421,4 +500,5 @@ def contains(scope: ComposedLayout, t: int) -> bool: "image", "project", "contains", + "vector_widths", ] diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 8c928054..9a23d44f 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -115,15 +115,11 @@ def tile_device( T.copy(a_view, a_tile) T.copy(b_view, b_tile) T.sync(m) - a_fragment_view = T.tensor_view( - T.ptr_of(a_tile), - layout=((2, 4 @ m.warp, 2, 8 @ m.lane, 2), (1, 2, 8, 16, 128)), - ) b_fragment_view = T.tensor_view( T.ptr_of(b_tile), layout=((8 @ m.lane, 2, 4 @ m.warp, 2), (1, 8, 16, 64)), ) - T.copy(a_fragment_view, a_frag) + T.ldmatrix(a_tile, a_frag) T.copy(b_fragment_view, b_frag) T.fill(acc, 0.0) T.tiled_mma(acc, a_frag, b_frag, atom=T.cuda.sm80.Mma()) diff --git a/tests/ops/tir/cuda/test_tma.py b/tests/ops/tir/cuda/test_tma.py index cc19f15f..84ce9038 100644 --- a/tests/ops/tir/cuda/test_tma.py +++ b/tests/ops/tir/cuda/test_tma.py @@ -13,7 +13,7 @@ from tilefoundry import module, prim_func from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core import Var, VerifyError -from tilefoundry.ir.tir.cuda.memory.tma import TmaCopy +from tilefoundry.ir.tir.cuda.memory.copy_async_bulk import CopyAsyncBulk from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.stmts import Evaluate, Return, Sequential from tilefoundry.ir.tir.verify import verify_prim_function @@ -29,7 +29,7 @@ def _pf(src, dst, bar=_BAR) -> PrimFunction: return PrimFunction( name="fn", params=args, - body=Sequential(body=(Evaluate(callable=TmaCopy(), args=args), Return())), + body=Sequential(body=(Evaluate(callable=CopyAsyncBulk(), args=args), Return())), ) @@ -121,7 +121,7 @@ def tma_tiers_device( bulk_bar = T.alloc_tensor(Tensor[(1,), "i64", None, "smem"]) T.mbarrier_init(bulk_bar, arrive_count=1) T.sync(m) - T.tma_copy(bulk_view, bulk_stage, bulk_bar) + T.copy_async_bulk(bulk_view, bulk_stage, bulk_bar) T.mbarrier_wait_parity(bulk_bar, 0) T.copy(bulk_stage, bulk_b) with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as mo: @@ -146,7 +146,7 @@ def tma_tiers_device( odd_bar = T.alloc_tensor(Tensor[(1,), "i64", None, "smem"]) T.mbarrier_init(odd_bar, arrive_count=1) T.sync(mo) - T.tma_copy(odd_view, odd_stage, odd_bar) + T.copy_async_bulk(odd_view, odd_stage, odd_bar) T.mbarrier_wait_parity(odd_bar, 0) T.copy(odd_stage, odd_b) diff --git a/tests/runtime/cuda/static_asserts.cu b/tests/runtime/cuda/static_asserts.cu index 7f327de3..dabceeea 100644 --- a/tests/runtime/cuda/static_asserts.cu +++ b/tests/runtime/cuda/static_asserts.cu @@ -416,7 +416,7 @@ __global__ void k(float *p) { #endif #if CASE == 24 -/// A ``tma_copy`` source a mesh really does divide. +/// A ``copy_async_bulk`` source a mesh really does divide. __global__ void k(float *p, float *q, uint64_t *bar) { auto mesh = tmesh<128>(); auto src_lay = @@ -430,7 +430,7 @@ __global__ void k(float *p, float *q, uint64_t *bar) { auto dst = make_shard_tensor( cute::make_tensor(cute::make_smem_ptr(q), dst_lay), dst_lay, make_shard_layout(dst_lay, mesh, cute::make_tuple(shard::B{}))); - tilefoundry::ops::tma_copy(src, dst, bar); + tilefoundry::ops::copy_async_bulk(src, dst, bar); } #endif diff --git a/tests/runtime/test_cuda_static_asserts.py b/tests/runtime/test_cuda_static_asserts.py index 07d1e845..27201a35 100644 --- a/tests/runtime/test_cuda_static_asserts.py +++ b/tests/runtime/test_cuda_static_asserts.py @@ -39,7 +39,7 @@ 19: "ops::sync: a mesh's scope must be cta or thread", 20: "ops::sync: a CTA mesh needs the module's grid-barrier counter", 22: "a reduced mesh axis must divide into whole lanes and whole warps", - 24: "both operands must leave the tile whole on every ", + 24: "ops::copy_async_bulk: both operands must leave the tile whole on every ", 25: "ops::sync: a mesh that skips warps names no barrier", 26: "get: this mesh does not name that level", } From d670d60904cbe738db983621d996a4b19905b1a5 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 02:52:43 +0800 Subject: [PATCH 04/19] fix(pattern): keep vector analysis above type roots --- src/tilefoundry/ir/pattern/__init__.py | 2 ++ src/tilefoundry/ir/pattern/pattern.py | 33 +++++++++++++++++++++- src/tilefoundry/ir/types/layout_algebra.py | 32 --------------------- 3 files changed, 34 insertions(+), 33 deletions(-) diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py index f8754474..0db17007 100644 --- a/src/tilefoundry/ir/pattern/__init__.py +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -47,6 +47,7 @@ TensorPattern, VectorPattern, WildcardPattern, + vector_widths, ) from .utils import ( MOVED_STORAGES, @@ -114,4 +115,5 @@ "storage_place", "VectorPattern", "vector", + "vector_widths", ] diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index 82a8d4cd..e2f2e859 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -19,12 +19,13 @@ from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.layout_algebra import ( ASYNC_WIDTHS, + box_runs, frame_of, is_inverse_projectable, - vector_widths, ) from tilefoundry.ir.types.mesh import separate +from .constraint import affine_part from .match import ( ABSENT, ARRANGEMENT, @@ -87,6 +88,34 @@ def describe(self, name: str = UNNAMED_PLACE) -> str: ) +def vector_widths(layout, element_bits: int) -> tuple[int, ...]: + """Every cp.async width that divides every run in an arrangement.""" + widest = ASYNC_WIDTHS[-1] + if isinstance(layout, ShardLayout): + layout = layout.layout + inner = getattr(layout, "inner", None) + if inner is not None and hasattr(inner, "base"): + widest = min(widest, 1 << inner.base) + held = affine_part(layout) + if held is None or any( + type(value) is not int + for group in (held.shape, held.strides) + for value in flatten(group) + ): + return () + runs = box_runs(held, element_bits, None, limit=None) + unit = [run.extent for run in runs if run.step == 1] + if len(unit) != 1: + return () + counted = (unit[0], *(run.step for run in runs if run.step != 1)) + return tuple( + width + for width in ASYNC_WIDTHS + if width <= widest + and all(value * element_bits % (width * 8) == 0 for value in counted) + ) + + @dataclass(frozen=True) class VectorPattern(Pattern): """An arrangement that moves whole 4-, 8-, or 16-byte vectors.""" @@ -831,5 +860,7 @@ def describe(self, name: str = UNNAMED_PLACE, arrangements=None) -> str: "SwitchPattern", "Tensor", "TensorPattern", + "VectorPattern", "WildcardPattern", + "vector_widths", ] diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index a2f18f7b..bac418da 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -16,7 +16,6 @@ from tilefoundry.ir.types.layout import flatten from .layout import ComposedLayout, Layout, Swizzle, size -from .shard_layout import ShardLayout from .stride import compact_col_major, idx2crd @@ -178,36 +177,6 @@ def box_runs( return tuple(sorted(runs, key=lambda run: run.step)) -def vector_widths(layout, element_bits: int) -> tuple[int, ...]: - """Every cp.async width that divides every run in an arrangement.""" - from tilefoundry.ir.pattern.constraint import affine_part # noqa: PLC0415 - - widest = ASYNC_WIDTHS[-1] - if isinstance(layout, ShardLayout): - layout = layout.layout - inner = getattr(layout, "inner", None) - if inner is not None and hasattr(inner, "base"): - widest = min(widest, 1 << inner.base) - held = affine_part(layout) - if held is None or any( - type(value) is not int - for group in (held.shape, held.strides) - for value in flatten(group) - ): - return () - runs = box_runs(held, element_bits, None, limit=None) - unit = [run.extent for run in runs if run.step == 1] - if len(unit) != 1: - return () - counted = (unit[0], *(run.step for run in runs if run.step != 1)) - return tuple( - width - for width in ASYNC_WIDTHS - if width <= widest - and all(value * element_bits % (width * 8) == 0 for value in counted) - ) - - def complement(layout: Layout, max_idx: int = 1) -> Layout: """CuTe ``complement``: the modes that fill the gaps below ``max_idx``.""" result_shape: list[int] = [] @@ -500,5 +469,4 @@ def contains(scope: ComposedLayout, t: int) -> bool: "image", "project", "contains", - "vector_widths", ] From 3bf85352ac2255182137e2da70c75e8293ab07bc Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 11:02:49 +0800 Subject: [PATCH 05/19] test(schedule): add CUDA instruction fixtures --- docs/spec/shard.md | 9 +- pyproject.toml | 6 +- src/tilefoundry/ir/types/mesh.py | 5 +- .../tir/gemm_8192x17408x5120_cta_grid.py | 467 ++++++++++++++++++ .../tir/gemm_8192x17408x5120_persistent.py | 466 +++++++++++++++++ .../gemm_8192x17408x5120_register_store.py | 447 +++++++++++++++++ .../tir/gemm_8192x17408x5120_tma_store.py | 467 ++++++++++++++++++ .../schedule/tir/sm80_mma_ldmatrix.py | 140 ++++++ .../fixtures/schedule/tir/wgmma_a_k_major.py | 132 +++++ .../fixtures/schedule/tir/wgmma_a_mn_major.py | 156 ++++++ .../tir/wgmma_cast_between_schedules.py | 136 +++++ .../schedule/tir/wgmma_cp_async_loads.py | 132 +++++ .../schedule/tir/wgmma_cta_grid_4x17.py | 177 +++++++ .../schedule/tir/wgmma_explicit_windows.py | 132 +++++ .../tir/wgmma_insert_tiles_into_output.py | 132 +++++ .../tir/wgmma_k_slices_of_wide_run.py | 177 +++++++ .../tir/wgmma_one_tile_of_larger_output.py | 131 +++++ .../schedule/tir/wgmma_repeat_along_k.py | 391 +++++++++++++++ .../schedule/tir/wgmma_repeat_along_n.py | 177 +++++++ .../tir/wgmma_rs_a_from_accumulator.py | 211 ++++++++ .../schedule/tir/wgmma_rs_a_from_smem.py | 137 +++++ .../schedule/tir/wgmma_swizzled_smem.py | 144 ++++++ .../fixtures/schedule/tir/wgmma_tma_3stage.py | 177 +++++++ .../schedule/tir/wgmma_two_schedules.py | 199 ++++++++ tests/ir/types/test_mesh.py | 28 +- tests/schedule/test_fixtures.py | 21 + 26 files changed, 4784 insertions(+), 13 deletions(-) create mode 100644 tests/fixtures/schedule/tir/gemm_8192x17408x5120_cta_grid.py create mode 100644 tests/fixtures/schedule/tir/gemm_8192x17408x5120_persistent.py create mode 100644 tests/fixtures/schedule/tir/gemm_8192x17408x5120_register_store.py create mode 100644 tests/fixtures/schedule/tir/gemm_8192x17408x5120_tma_store.py create mode 100644 tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py create mode 100644 tests/fixtures/schedule/tir/wgmma_a_k_major.py create mode 100644 tests/fixtures/schedule/tir/wgmma_a_mn_major.py create mode 100644 tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py create mode 100644 tests/fixtures/schedule/tir/wgmma_cp_async_loads.py create mode 100644 tests/fixtures/schedule/tir/wgmma_cta_grid_4x17.py create mode 100644 tests/fixtures/schedule/tir/wgmma_explicit_windows.py create mode 100644 tests/fixtures/schedule/tir/wgmma_insert_tiles_into_output.py create mode 100644 tests/fixtures/schedule/tir/wgmma_k_slices_of_wide_run.py create mode 100644 tests/fixtures/schedule/tir/wgmma_one_tile_of_larger_output.py create mode 100644 tests/fixtures/schedule/tir/wgmma_repeat_along_k.py create mode 100644 tests/fixtures/schedule/tir/wgmma_repeat_along_n.py create mode 100644 tests/fixtures/schedule/tir/wgmma_rs_a_from_accumulator.py create mode 100644 tests/fixtures/schedule/tir/wgmma_rs_a_from_smem.py create mode 100644 tests/fixtures/schedule/tir/wgmma_swizzled_smem.py create mode 100644 tests/fixtures/schedule/tir/wgmma_tma_3stage.py create mode 100644 tests/fixtures/schedule/tir/wgmma_two_schedules.py diff --git a/docs/spec/shard.md b/docs/spec/shard.md index c64a2741..9eec878d 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -335,10 +335,11 @@ Mesh composition uses the following rules: - Append concatenates the per-level arrangements. When an appended mesh is sliced, its per-level start is re-encoded in the combined device numbering and retained in the result's `ComposedLayout.offset`. -- Replacing a suffix with a sliced mesh MUST be rejected. The slice belongs on - the combined multi-level mesh, where every axis contributing to its device - offset is stated together. Replacing an unsliced suffix and replacing the - whole mesh retain their existing behavior. +- Replacing a suffix with a sliced mesh MUST preserve each retained upper + level's start and take each replaced level's start and arrangement from the + inner mesh. The combined `ComposedLayout.offset` MUST then be re-encoded in + device numbering from those per-level starts. Replacing an unsliced suffix + and replacing the whole mesh retain their existing behavior. - `make_mesh(*meshes)` invokes `check_topology` on its result. For each named level with a concrete declared extent, its position count MUST NOT exceed that extent; symbolic extents are deferred until dimensions are bound. A diff --git a/pyproject.toml b/pyproject.toml index 4c8135cd..6c90ab77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,9 +91,9 @@ markers = [ ] [tool.ruff] -# Snapshot examples, shipped model sources, and recorded parser output are evidence; -# formatting them changes that evidence. -extend-exclude = ["examples", "tests/models", "tests/parser/golden"] +# Snapshot examples, shipped model sources, and recorded parser and printer output are +# evidence; formatting them changes that evidence. +extend-exclude = ["examples", "tests/models", "tests/parser/golden", "tests/fixtures/schedule/tir"] force-exclude = true line-length = 100 target-version = "py312" diff --git a/src/tilefoundry/ir/types/mesh.py b/src/tilefoundry/ir/types/mesh.py index 0038d02a..1e38cb8b 100644 --- a/src/tilefoundry/ir/types/mesh.py +++ b/src/tilefoundry/ir/types/mesh.py @@ -324,8 +324,6 @@ def make_mesh(*meshes: Mesh) -> Mesh: elif set(here) <= set(there): result = inner elif len(there) < len(here) and here[-len(there) :] == there: - if isinstance(inner.layout, ComposedLayout): - raise ValueError("cannot replace a mesh suffix with a sliced mesh") kept = len(here) - len(there) above = _levels(result)[:kept] named = sum(len(flatten(level.shape)) for level in above) @@ -334,7 +332,8 @@ def make_mesh(*meshes: Mesh) -> Mesh: (*above, *_levels(inner)), (*_starts(result)[:kept], *_starts(inner)), (*result.names[:named], *inner.names), - sliced=isinstance(result.layout, ComposedLayout), + sliced=isinstance(result.layout, ComposedLayout) + or isinstance(inner.layout, ComposedLayout), ) else: shared = sorted(set(here) & set(there)) diff --git a/tests/fixtures/schedule/tir/gemm_8192x17408x5120_cta_grid.py b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_cta_grid.py new file mode 100644 index 00000000..76d005e6 --- /dev/null +++ b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_cta_grid.py @@ -0,0 +1,467 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(8192, 5120), "bf16"], b: Tensor[(5120, 17408), "bf16"], out: Tensor[(8192, 17408), "bf16"] +): + with Mesh((Topology("cta", 4352),), Layout((64, 68), (68, 1)), names=("d0", "d1")) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "f32", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "bf16", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_4: + for m in range(128 * cta.d0, (128 * cta.d0) + 128, 128): + for n in range(256 * cta.d1, (256 * cta.d1) + 256, 256): + lhs_stages = (T.tensor_view(65536, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(73728, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(81920, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(90112, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(16384, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(32768, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(49152, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 5120, 64): + with scope_4[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[m:m + 128, k:k + 64]), + layout=Layout((128, 64), (5120, 1)), + shape=(128, 64), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 64) % 4]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 64, n:n + 256]), + layout=Layout((64, 256), (17408, 1)), + shape=(64, 256), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 64) % 4]) + with scope_4[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 256, 256): + for o_k in range(0, 64, 64): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k:o_k + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 16:o_k + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 16:o_k + 16 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_2 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_2 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 32:o_k + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_2 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 32:o_k + 32 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_2, + lhs_view_2, + rhs_view_2, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_3 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_3 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 48:o_k + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_3 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 48:o_k + 48 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_3, + lhs_view_3, + rhs_view_3, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 256, 256): + for o_k_1 in range(0, 64, 64): + acc_view_4 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_4 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_4 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_4, + lhs_view_4, + rhs_view_4, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_5 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_5 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 16:o_k_1 + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_5 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 16:o_k_1 + 16 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_5, + lhs_view_5, + rhs_view_5, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_6 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_6 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 32:o_k_1 + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_6 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 32:o_k_1 + 32 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_6, + lhs_view_6, + rhs_view_6, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_7 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_7 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 48:o_k_1 + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_7 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 48:o_k_1 + 48 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_7, + lhs_view_7, + rhs_view_7, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + copy = T.tensor_view( + 0, + dtype='bf16', + storage=StorageKind.SMEM, + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout((128, (4, 64)), (64, (8192, 1))), + ), + shape=(128, 256), + ) + with scope_4[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.cast(acc, value_view) + src_frame = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.copy(src_frame, copy) + with scope_4[:1, :32] as scope_3: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_7: + window = T.tensor_view( + T.ptr_of(out[m:m + 128, n:n + 256]), + layout=Layout((128, 256), (17408, 1)), + shape=(128, 256), + ) + T.copy_async_tensor(copy, window) diff --git a/tests/fixtures/schedule/tir/gemm_8192x17408x5120_persistent.py b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_persistent.py new file mode 100644 index 00000000..2001b07d --- /dev/null +++ b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_persistent.py @@ -0,0 +1,466 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(8192, 5120), "bf16"], b: Tensor[(5120, 17408), "bf16"], out: Tensor[(8192, 17408), "bf16"] +): + with Mesh((Topology("cta", 132),), Layout((132,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "f32", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "bf16", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_4: + for t in range(cta.d0, 4352, 132): + lhs_stages = (T.tensor_view(65536, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(73728, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(81920, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(90112, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(16384, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(32768, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(49152, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 5120, 64): + with scope_4[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[(t // 68) * 128:(t // 68) * 128 + 128, k:k + 64]), + layout=Layout((128, 64), (5120, 1)), + shape=(128, 64), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 64) % 4]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 64, (t % 68) * 256:(t % 68) * 256 + 256]), + layout=Layout((64, 256), (17408, 1)), + shape=(64, 256), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 64) % 4]) + with scope_4[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 256, 256): + for o_k in range(0, 64, 64): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k:o_k + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 16:o_k + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 16:o_k + 16 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_2 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_2 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 32:o_k + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_2 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 32:o_k + 32 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_2, + lhs_view_2, + rhs_view_2, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_3 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_3 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 48:o_k + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_3 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 48:o_k + 48 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_3, + lhs_view_3, + rhs_view_3, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 256, 256): + for o_k_1 in range(0, 64, 64): + acc_view_4 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_4 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_4 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_4, + lhs_view_4, + rhs_view_4, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_5 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_5 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 16:o_k_1 + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_5 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 16:o_k_1 + 16 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_5, + lhs_view_5, + rhs_view_5, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_6 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_6 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 32:o_k_1 + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_6 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 32:o_k_1 + 32 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_6, + lhs_view_6, + rhs_view_6, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_7 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_7 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 48:o_k_1 + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_7 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 48:o_k_1 + 48 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_7, + lhs_view_7, + rhs_view_7, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + copy = T.tensor_view( + 0, + dtype='bf16', + storage=StorageKind.SMEM, + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout((128, (4, 64)), (64, (8192, 1))), + ), + shape=(128, 256), + ) + with scope_4[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.cast(acc, value_view) + src_frame = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.copy(src_frame, copy) + with scope_4[:1, :32] as scope_3: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_7: + window = T.tensor_view( + T.ptr_of(out[(t // 68) * 128:(t // 68) * 128 + 128, (t % 68) * 256:(t % 68) * 256 + 256]), + layout=Layout((128, 256), (17408, 1)), + shape=(128, 256), + ) + T.copy_async_tensor(copy, window) diff --git a/tests/fixtures/schedule/tir/gemm_8192x17408x5120_register_store.py b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_register_store.py new file mode 100644 index 00000000..06f2b54f --- /dev/null +++ b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_register_store.py @@ -0,0 +1,447 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(8192, 5120), "bf16"], b: Tensor[(5120, 17408), "bf16"], out: Tensor[(8192, 17408), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "f32", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "bf16", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + for m in range(0, 8192, 128): + for n in range(0, 17408, 256): + lhs_stages = (T.tensor_view(65536, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(73728, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(81920, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(90112, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(16384, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(32768, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(49152, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 5120, 64): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[m:m + 128, k:k + 64]), + layout=Layout((128, 64), (5120, 1)), + shape=(128, 64), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 64) % 4]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 64, n:n + 256]), + layout=Layout((64, 256), (17408, 1)), + shape=(64, 256), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 64) % 4]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 256, 256): + for o_k in range(0, 64, 64): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k:o_k + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 16:o_k + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 16:o_k + 16 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_2 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_2 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 32:o_k + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_2 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 32:o_k + 32 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_2, + lhs_view_2, + rhs_view_2, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_3 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_3 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 48:o_k + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_3 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 48:o_k + 48 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_3, + lhs_view_3, + rhs_view_3, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 256, 256): + for o_k_1 in range(0, 64, 64): + acc_view_4 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_4 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_4 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_4, + lhs_view_4, + rhs_view_4, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_5 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_5 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 16:o_k_1 + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_5 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 16:o_k_1 + 16 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_5, + lhs_view_5, + rhs_view_5, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_6 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_6 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 32:o_k_1 + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_6 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 32:o_k_1 + 32 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_6, + lhs_view_6, + rhs_view_6, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_7 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_7 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 48:o_k_1 + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_7 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 48:o_k_1 + 48 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_7, + lhs_view_7, + rhs_view_7, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.cast(acc, value_view) + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + window = T.tensor_view( + T.ptr_of(out[m:m + 128, n:n + 256]), + layout=Layout((128, 256), (17408, 1)), + shape=(128, 256), + ) + T.copy(value_view_1, window) diff --git a/tests/fixtures/schedule/tir/gemm_8192x17408x5120_tma_store.py b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_tma_store.py new file mode 100644 index 00000000..5c31f31d --- /dev/null +++ b/tests/fixtures/schedule/tir/gemm_8192x17408x5120_tma_store.py @@ -0,0 +1,467 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(8192, 5120), "bf16"], b: Tensor[(5120, 17408), "bf16"], out: Tensor[(8192, 17408), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "f32", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "bf16", + Layout((2, 8, 2, 4, 2, 4, 32), (16384, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_4: + for m in range(0, 8192, 128): + for n in range(0, 17408, 256): + lhs_stages = (T.tensor_view(65536, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(73728, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(81920, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(90112, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(16384, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(32768, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256)), T.tensor_view(49152, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((4, 2, 8), (4, 64)), ((4096, 512, 64), (1024, 1))), + ), shape=(64, 256))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 5120, 64): + with scope_4[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[m:m + 128, k:k + 64]), + layout=Layout((128, 64), (5120, 1)), + shape=(128, 64), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 64) % 4]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 64, n:n + 256]), + layout=Layout((64, 256), (17408, 1)), + shape=(64, 256), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 64) % 4]) + with scope_4[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 256, 256): + for o_k in range(0, 64, 64): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k:o_k + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 16:o_k + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 16:o_k + 16 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_2 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_2 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 32:o_k + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_2 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 32:o_k + 32 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_2, + lhs_view_2, + rhs_view_2, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_3 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 256]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_3 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m:o_m + 64, o_k + 48:o_k + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_3 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k + 48:o_k + 48 + 16, o_n:o_n + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_3, + lhs_view_3, + rhs_view_3, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 256, 256): + for o_k_1 in range(0, 64, 64): + acc_view_4 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_4 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_4 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_4, + lhs_view_4, + rhs_view_4, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_5 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_5 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 16:o_k_1 + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_5 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 16:o_k_1 + 16 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_5, + lhs_view_5, + rhs_view_5, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_6 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_6 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 32:o_k_1 + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_6 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 32:o_k_1 + 32 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_6, + lhs_view_6, + rhs_view_6, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_7 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 256]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 32), (1, 8, 16, 64, 128, 512)), + shape=(64, 256), + ) + lhs_view_7 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 4][o_m_1:o_m_1 + 64, o_k_1 + 48:o_k_1 + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_7 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 4][o_k_1 + 48:o_k_1 + 48 + 16, o_n_1:o_n_1 + 256]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8), (4, 64)), ((512, 64), (1024, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 256), + ) + T.tiled_mma( + acc_view_7, + lhs_view_7, + rhs_view_7, + atom=T.cuda.sm90.Wgmma(n=256, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + copy = T.tensor_view( + 0, + dtype='bf16', + storage=StorageKind.SMEM, + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout((128, (4, 64)), (64, (8192, 1))), + ), + shape=(128, 256), + ) + with scope_4[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.cast(acc, value_view) + src_frame = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 32), (16384, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.copy(src_frame, copy) + with scope_4[:1, :32] as scope_3: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_7: + window = T.tensor_view( + T.ptr_of(out[m:m + 128, n:n + 256]), + layout=Layout((128, 256), (17408, 1)), + shape=(128, 256), + ) + T.copy_async_tensor(copy, window) diff --git a/tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py b/tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py new file mode 100644 index 00000000..4e181994 --- /dev/null +++ b/tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm(a: Tensor[(16, 32), "bf16"], b: Tensor[(32, 8), "bf16"], out: Tensor[(16, 8), "bf16"]): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[(16, 8), "f32", Layout((2, 4, 8, 2), (1, 2, 8, 64)), "rmem"] + ) + value = T.alloc_tensor( + tensor_type=Tensor[(16, 8), "bf16", Layout((2, 4, 8, 2), (1, 2, 8, 64)), "rmem"] + ) + with Mesh( + (Topology("thread", 64),), Layout((2, 32), (32, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(256, dtype='bf16', storage=StorageKind.SMEM, layout=Layout((16, 16), (16, 1)), shape=(16, 16)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout((16, 16), (16, 1)), shape=(16, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout((16, 8), (8, 1)), shape=(16, 8)), T.tensor_view(128, dtype='bf16', storage=StorageKind.SMEM, layout=Layout((16, 8), (8, 1)), shape=(16, 8))) + with Mesh( + (Topology("thread", 64),), ComposedLayout( + inner=None, + offset=32, + outer=Layout((4, 8), (1, 4)), +), names=("d0", "d1") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 16, k:k + 16]), + layout=Layout((16, 16), (32, 1)), + shape=(16, 16), + ) + with Mesh( + (Topology("thread", 64),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 8]), + layout=Layout((16, 8), (8, 1)), + shape=(16, 8), + ) + with Mesh( + (Topology("thread", 64),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + ldmatrix = T.alloc_tensor( + tensor_type=Tensor[ + (16, 16), + "bf16", + Layout((2, 4, 2, 8, 2), (1, 2, 8, 16, 128)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 64),), ComposedLayout( + inner=None, + offset=32, + outer=Layout((4, 8), (1, 4)), +), names=("d0", "d1") + ) as threads_3: + T.ldmatrix(lhs_stages[(k // 16) % 2], ldmatrix) + copy = T.alloc_tensor( + tensor_type=Tensor[ + (16, 8), "bf16", Layout((8, 2, 4, 2), (1, 8, 16, 64)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 64),), ComposedLayout( + inner=None, + offset=32, + outer=Layout((4, 8), (1, 4)), +), names=("d0", "d1") + ) as threads_4: + T.copy(rhs_stages[(k // 16) % 2], copy) + for o_m in range(0, 16, 16): + for o_n in range(0, 8, 8): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 16, o_n:o_n + 8]), + layout=((2, 4 @ threads_4.d0, 8 @ threads_4.d1, 2), (1, 2, 8, 64)), + shape=(16, 8), + ) + lhs_view = T.tensor_view( + T.ptr_of(ldmatrix[o_m:o_m + 16, o_k:o_k + 16]), + layout=((2, 4 @ threads_4.d0, 2, 8 @ threads_4.d1, 2), (1, 2, 8, 16, 128)), + shape=(16, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(copy[o_k:o_k + 16, o_n:o_n + 8]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2), (1, 8, 16, 64)), + shape=(16, 8), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm80.Mma(mesh=threads_4), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 64),), ComposedLayout( + inner=None, + offset=32, + outer=Layout((4, 8), (1, 4)), +), names=("d0", "d1") + ) as threads_6: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 16, 0:0 + 8]), + layout=((2, 4 @ threads_6.d0, 8 @ threads_6.d1, 2), (1, 2, 8, 64)), + shape=(16, 8), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 64),), ComposedLayout( + inner=None, + offset=32, + outer=Layout((4, 8), (1, 4)), +), names=("d0", "d1") + ) as threads_7: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 16, 0:0 + 8]), + layout=((2, 4 @ threads_7.d0, 8 @ threads_7.d1, 2), (1, 2, 8, 64)), + shape=(16, 8), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_a_k_major.py b/tests/fixtures/schedule/tir/wgmma_a_k_major.py new file mode 100644 index 00000000..20b823a0 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_a_k_major.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b: Tensor[(32, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_a_mn_major.py b/tests/fixtures/schedule/tir/wgmma_a_mn_major.py new file mode 100644 index 00000000..61938993 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_a_mn_major.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + at: Tensor[(32, 64), "bf16"], b: Tensor[(32, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((1, 64), (2, 8)), ((1024, 1), (512, 64))), + ), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((1, 64), (2, 8)), ((1024, 1), (512, 64))), + ), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((2, 8), (1, 32)), ((256, 32), (512, 1))), + ), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((2, 8), (1, 32)), ((256, 32), (512, 1))), + ), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(at[k:k + 16, 0:0 + 64]), + layout=Layout((64, 16), (1, 64)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((1, 64), (2, 8)), ((1024, 1), (512, 64))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((2, 8), (1, 32)), ((256, 32), (512, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py b/tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py new file mode 100644 index 00000000..264aade6 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b_f32: Tensor[(32, 32), "f32"], scratch: Tensor[(1024,), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.tensor_view( + T.ptr_of(scratch[0:0 + 1024]), layout=Layout((32, 32), (32, 1)), shape=(32, 32) + ) + value_1 = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_4: + T.cast(b_f32, value) + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_4[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(value[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_4[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_4[1:] as scope_3: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_1_view = T.tensor_view( + T.ptr_of(value_1[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_1_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + value_1_view_1 = T.tensor_view( + T.ptr_of(value_1[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_1_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_cp_async_loads.py b/tests/fixtures/schedule/tir/wgmma_cp_async_loads.py new file mode 100644 index 00000000..92df8105 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_cp_async_loads.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b: Tensor[(32, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_cta_grid_4x17.py b/tests/fixtures/schedule/tir/wgmma_cta_grid_4x17.py new file mode 100644 index 00000000..99f3fcea --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_cta_grid_4x17.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(128, 64), "bf16"], b: Tensor[(64, 16), "bf16"], out: Tensor[(128, 16), "bf16"] +): + with Mesh((Topology("cta", 68),), Layout((4, 17), (17, 1)), names=("d0", "d1")) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 16), + "f32", + Layout((2, 8, 2, 4, 2, 4, 2), (1024, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 16), + "bf16", + Layout((2, 8, 2, 4, 2, 4, 2), (1024, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(768, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16)), T.tensor_view(2816, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16)), T.tensor_view(4864, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16)), T.tensor_view(256, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 64, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 128, k:k + 16]), + layout=Layout((128, 16), (64, 1)), + shape=(128, 16), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 3]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 16]), + layout=Layout((16, 16), (16, 1)), + shape=(16, 16), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 3]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 16, 16): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 3][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 3][o_k:o_k + 16, o_n:o_n + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 16, 16): + for o_k_1 in range(0, 16, 16): + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 16]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 3][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 3][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 16]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 2), (1024, 1, 8, 16, 64, 128, 512)), + shape=(128, 16), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_6: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 16]), + layout=((2 @ threads_6.d0, 8 @ threads_6.d2, 2, 4 @ threads_6.d1, 2, 4 @ threads_6.d3, 2), (1024, 1, 8, 16, 64, 128, 512)), + shape=(128, 16), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_explicit_windows.py b/tests/fixtures/schedule/tir/wgmma_explicit_windows.py new file mode 100644 index 00000000..704dbd35 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_explicit_windows.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(128, 32), "bf16"], b: Tensor[(32, 64), "bf16"], out: Tensor[(128, 64), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + for m in range(0, 128, 64): + for n in range(0, 64, 32): + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[m:m + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, n:n + 32]), + layout=Layout((16, 32), (64, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + window = T.tensor_view( + T.ptr_of(out[m:m + 64, n:n + 32]), + layout=Layout((64, 32), (64, 1)), + shape=(64, 32), + ) + T.copy(value_view_1, window) diff --git a/tests/fixtures/schedule/tir/wgmma_insert_tiles_into_output.py b/tests/fixtures/schedule/tir/wgmma_insert_tiles_into_output.py new file mode 100644 index 00000000..704dbd35 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_insert_tiles_into_output.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(128, 32), "bf16"], b: Tensor[(32, 64), "bf16"], out: Tensor[(128, 64), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + for m in range(0, 128, 64): + for n in range(0, 64, 32): + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[m:m + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, n:n + 32]), + layout=Layout((16, 32), (64, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + window = T.tensor_view( + T.ptr_of(out[m:m + 64, n:n + 32]), + layout=Layout((64, 32), (64, 1)), + shape=(64, 32), + ) + T.copy(value_view_1, window) diff --git a/tests/fixtures/schedule/tir/wgmma_k_slices_of_wide_run.py b/tests/fixtures/schedule/tir/wgmma_k_slices_of_wide_run.py new file mode 100644 index 00000000..532e0eb2 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_k_slices_of_wide_run.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 64), "bf16"], b: Tensor[(64, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((8, 8), (2, 16)), ((256, 32), (16, 1))), + ), shape=(64, 32)), T.tensor_view(4096, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((8, 8), (2, 16)), ((256, 32), (16, 1))), + ), shape=(64, 32))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 2, 8), (4, 8)), ((512, 64, 8), (128, 1))), shape=(32, 32)), T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 2, 8), (4, 8)), ((512, 64, 8), (128, 1))), shape=(32, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 64, 32): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 32]), + layout=Layout((64, 32), (64, 1)), + shape=(64, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 32) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 32, 0:0 + 32]), + layout=Layout((32, 32), (32, 1)), + shape=(32, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 32) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 32, 32): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 32) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((256, 32), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 32) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 32) % 2][o_m:o_m + 64, o_k + 16:o_k + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((256, 32), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 32) % 2][o_k + 16:o_k + 16 + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_one_tile_of_larger_output.py b/tests/fixtures/schedule/tir/wgmma_one_tile_of_larger_output.py new file mode 100644 index 00000000..3651b139 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_one_tile_of_larger_output.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b: Tensor[(32, 32), "bf16"], out: Tensor[(128, 64), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + T.fill(out, 0.0) + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + window = T.tensor_view( + T.ptr_of(out[0:0 + 64, 0:0 + 32]), + layout=Layout((64, 32), (64, 1)), + shape=(64, 32), + ) + T.copy(value_view_1, window) diff --git a/tests/fixtures/schedule/tir/wgmma_repeat_along_k.py b/tests/fixtures/schedule/tir/wgmma_repeat_along_k.py new file mode 100644 index 00000000..1ee75b91 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_repeat_along_k.py @@ -0,0 +1,391 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(128, 128), "bf16"], b: Tensor[(128, 16), "bf16"], out: Tensor[(128, 16), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 16), + "f32", + Layout((2, 8, 2, 4, 2, 4, 2), (1024, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 16), + "bf16", + Layout((2, 8, 2, 4, 2, 4, 2), (1024, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64)), T.tensor_view(10240, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((2, 8, 8), (4, 16)), ((4096, 512, 64), (16, 1))), + ), shape=(128, 64))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((4, 2, 8), (2, 8)), ((256, 64, 8), (128, 1))), shape=(64, 16)), T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((4, 2, 8), (2, 8)), ((256, 64, 8), (128, 1))), shape=(64, 16))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 128, 64): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 128, k:k + 64]), + layout=Layout((128, 64), (128, 1)), + shape=(128, 64), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 64) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 64, 0:0 + 16]), + layout=Layout((64, 16), (16, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 64) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 16, 16): + for o_k in range(0, 64, 64): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k:o_k + 16, o_n:o_n + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m:o_m + 64, o_k + 16:o_k + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k + 16:o_k + 16 + 16, o_n:o_n + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_2 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_2 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m:o_m + 64, o_k + 32:o_k + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_2 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k + 32:o_k + 32 + 16, o_n:o_n + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_2, + lhs_view_2, + rhs_view_2, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + acc_view_3 = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_3 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m:o_m + 64, o_k + 48:o_k + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view_3 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k + 48:o_k + 48 + 16, o_n:o_n + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_3, + lhs_view_3, + rhs_view_3, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 16, 16): + for o_k_1 in range(0, 64, 64): + acc_view_4 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 16]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_4 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=0, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_4 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_4, + lhs_view_4, + rhs_view_4, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_5 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 16]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_5 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m_1:o_m_1 + 64, o_k_1 + 16:o_k_1 + 16 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=16, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_5 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k_1 + 16:o_k_1 + 16 + 16, o_n_1:o_n_1 + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_5, + lhs_view_5, + rhs_view_5, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_6 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 16]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_6 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m_1:o_m_1 + 64, o_k_1 + 32:o_k_1 + 32 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=32, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_6 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k_1 + 32:o_k_1 + 32 + 16, o_n_1:o_n_1 + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_6, + lhs_view_6, + rhs_view_6, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + acc_view_7 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 16]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_7 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 64) % 2][o_m_1:o_m_1 + 64, o_k_1 + 48:o_k_1 + 48 + 16]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(3, 4, 3), + offset=48, + outer=Layout(((8, 8), 16), ((512, 64), 1)), + ), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_7 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 64) % 2][o_k_1 + 48:o_k_1 + 48 + 16, o_n_1:o_n_1 + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_7, + lhs_view_7, + rhs_view_7, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 16]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 2), (1024, 1, 8, 16, 64, 128, 512)), + shape=(128, 16), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_6: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 16]), + layout=((2 @ threads_6.d0, 8 @ threads_6.d2, 2, 4 @ threads_6.d1, 2, 4 @ threads_6.d3, 2), (1024, 1, 8, 16, 64, 128, 512)), + shape=(128, 16), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_repeat_along_n.py b/tests/fixtures/schedule/tir/wgmma_repeat_along_n.py new file mode 100644 index 00000000..35ce4f20 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_repeat_along_n.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(128, 32), "bf16"], b: Tensor[(32, 256), "bf16"], out: Tensor[(128, 256), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "f32", + Layout((2, 4, 8, 2, 4, 2, 4, 8), (16384, 4096, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 256), + "bf16", + Layout((2, 4, 8, 2, 4, 2, 4, 8), (16384, 4096, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(8192, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16)), T.tensor_view(10240, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8, 8)), ((64, 8), (1024, 128, 1))), shape=(16, 256)), T.tensor_view(4096, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8, 8)), ((64, 8), (1024, 128, 1))), shape=(16, 256))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 128, k:k + 16]), + layout=Layout((128, 16), (32, 1)), + shape=(128, 16), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 256]), + layout=Layout((16, 256), (256, 1)), + shape=(16, 256), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 256, 64): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 64]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 8), (1, 8, 16, 64, 128, 512)), + shape=(64, 64), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 64]), + layout=ShardLayout( + layout=Layout(((2, 8), (8, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 64), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=64, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 256, 64): + for o_k_1 in range(0, 16, 16): + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 64]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 8), (1, 8, 16, 64, 128, 512)), + shape=(64, 64), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 64]), + layout=ShardLayout( + layout=Layout(((2, 8), (8, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 64), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=64, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_5.d0, 4, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 8), (16384, 4096, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_6: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 256]), + layout=((2 @ threads_6.d0, 4, 8 @ threads_6.d2, 2, 4 @ threads_6.d1, 2, 4 @ threads_6.d3, 8), (16384, 4096, 1, 8, 16, 64, 128, 512)), + shape=(128, 256), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_rs_a_from_accumulator.py b/tests/fixtures/schedule/tir/wgmma_rs_a_from_accumulator.py new file mode 100644 index 00000000..417d2640 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_rs_a_from_accumulator.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b: Tensor[(32, 16), "bf16"], b1: Tensor[(16, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + p = T.alloc_tensor( + tensor_type=Tensor[ + (64, 16), "f32", Layout((8, 2, 4, 2, 4, 2), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 16), "bf16", Layout((8, 2, 4, 2, 4, 2), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value_1 = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_6: + lhs_stages = (T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(1536, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16)), T.tensor_view(256, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(p, 0.0) + for k in range(0, 32, 16): + with scope_6[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 16]), + layout=Layout((16, 16), (16, 1)), + shape=(16, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_6[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 16, 16): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(p[o_m:o_m + 64, o_n:o_n + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_6[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + T.cast(p, value) + rhs_stages_1 = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + T.fill(acc, 0.0) + for j in range(0, 16, 16): + with scope_6[:1, :32] as scope_3: + tile_2 = T.tensor_view( + T.ptr_of(b1[j:j + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_6: + T.copy_async_tensor(tile_2, rhs_stages_1[(j // 16) % 2]) + with scope_6[1:] as scope_4: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_7: + for o_m_1 in range(0, 64, 64): + for o_n_1 in range(0, 32, 32): + for o_k_1 in range(0, 16, 16): + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 32]), + layout=((8 @ threads_7.d1, 2, 4 @ threads_7.d0, 2, 4 @ threads_7.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(value[o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=((8 @ threads_7.d1, 2, 4 @ threads_7.d0, 2, 4 @ threads_7.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages_1[(j // 16) % 2][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_7, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.RS, mesh=threads_7), + ) + with scope_6[1:] as scope_5: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_8: + value_1_view = T.tensor_view( + T.ptr_of(value_1[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_8.d1, 2, 4 @ threads_8.d0, 2, 4 @ threads_8.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_1_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_9: + value_1_view_1 = T.tensor_view( + T.ptr_of(value_1[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_9.d1, 2, 4 @ threads_9.d0, 2, 4 @ threads_9.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_1_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_rs_a_from_smem.py b/tests/fixtures/schedule/tir/wgmma_rs_a_from_smem.py new file mode 100644 index 00000000..0d84fdbc --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_rs_a_from_smem.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b: Tensor[(32, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + copy = T.alloc_tensor( + tensor_type=Tensor[ + (64, 16), + "bf16", + Layout((8, 2, 4, 2, 4, 2), (1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + T.copy(lhs_stages[(k // 16) % 2], copy) + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(copy[o_m:o_m + 64, o_k:o_k + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.RS, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_6: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_6.d1, 2, 4 @ threads_6.d0, 2, 4 @ threads_6.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_swizzled_smem.py b/tests/fixtures/schedule/tir/wgmma_swizzled_smem.py new file mode 100644 index 00000000..d0d435bb --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_swizzled_smem.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b: Tensor[(32, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((2, 8), (1, 32)), ((256, 32), (512, 1))), + ), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((2, 8), (1, 32)), ((256, 32), (512, 1))), + ), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 2]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=ComposedLayout( + inner=Swizzle(2, 4, 3), + offset=0, + outer=Layout(((2, 8), (1, 32)), ((256, 32), (512, 1))), + ), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_tma_3stage.py b/tests/fixtures/schedule/tir/wgmma_tma_3stage.py new file mode 100644 index 00000000..04fbf2a7 --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_tma_3stage.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(128, 64), "bf16"], b: Tensor[(64, 16), "bf16"], out: Tensor[(128, 16), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (128, 16), + "f32", + Layout((2, 8, 2, 4, 2, 4, 2), (1024, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (128, 16), + "bf16", + Layout((2, 8, 2, 4, 2, 4, 2), (1024, 1, 8, 16, 64, 128, 512)), + "rmem", + ] + ) + with Mesh( + (Topology("thread", 384),), Layout((3, 128), (128, 1)), names=("d0", "d1") + ) as scope_3: + lhs_stages = (T.tensor_view(768, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16)), T.tensor_view(2816, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16)), T.tensor_view(4864, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8, 8), (2, 8)), ((1024, 128, 8), (64, 1))), shape=(128, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16)), T.tensor_view(256, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), shape=(16, 16))) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 64, 16): + with scope_3[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(a[0:0 + 128, k:k + 16]), + layout=Layout((128, 16), (64, 1)), + shape=(128, 16), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_stages[(k // 16) % 3]) + tile_1 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 16]), + layout=Layout((16, 16), (16, 1)), + shape=(16, 16), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_stages[(k // 16) % 3]) + with scope_3[1:] as scope_1: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_3: + for o_m in range(0, 64, 64): + for o_n in range(0, 16, 16): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 16]), + layout=((8 @ threads_3.d1, 2, 4 @ threads_3.d0, 2, 4 @ threads_3.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 3][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 3][o_k:o_k + 16, o_n:o_n + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_3, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_3), + ) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=256, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_4: + for o_m_1 in range(64, 128, 64): + for o_n_1 in range(0, 16, 16): + for o_k_1 in range(0, 16, 16): + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 16]), + layout=((8 @ threads_4.d1, 2, 4 @ threads_4.d0, 2, 4 @ threads_4.d2, 2), (1, 8, 16, 64, 128, 512)), + shape=(64, 16), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 3][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 3][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 16]), + layout=ShardLayout( + layout=Layout(((2, 8), (2, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_4, + ), + shape=(16, 16), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=16, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_4), + ) + with scope_3[1:] as scope_2: + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_5: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 16]), + layout=((2 @ threads_5.d0, 8 @ threads_5.d2, 2, 4 @ threads_5.d1, 2, 4 @ threads_5.d3, 2), (1024, 1, 8, 16, 64, 128, 512)), + shape=(128, 16), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 384),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((2, 4, 8, 4), (128, 32, 4, 1)), +), names=("d0", "d1", "d2", "d3") + ) as threads_6: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 128, 0:0 + 16]), + layout=((2 @ threads_6.d0, 8 @ threads_6.d2, 2, 4 @ threads_6.d1, 2, 4 @ threads_6.d3, 2), (1024, 1, 8, 16, 64, 128, 512)), + shape=(128, 16), + ) + T.copy(value_view_1, out) diff --git a/tests/fixtures/schedule/tir/wgmma_two_schedules.py b/tests/fixtures/schedule/tir/wgmma_two_schedules.py new file mode 100644 index 00000000..724496ff --- /dev/null +++ b/tests/fixtures/schedule/tir/wgmma_two_schedules.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +from tilefoundry import prim_func +from tilefoundry.dsl import T, Tensor +from tilefoundry.ir.types import B, ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.storage import StorageKind +from tilefoundry.target import CudaTarget + + +@prim_func(target=CudaTarget("nvidia.h200_sxm")) +def gemm( + a: Tensor[(64, 32), "bf16"], b: Tensor[(32, 32), "bf16"], c: Tensor[(64, 32), "bf16"], d: Tensor[(32, 32), "bf16"], out: Tensor[(64, 32), "bf16"] +): + with Mesh((Topology("cta", 1),), Layout((1,), (1,)), names=("d0",)) as cta: + acc = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "f32", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + value = T.alloc_tensor( + tensor_type=Tensor[ + (64, 32), "bf16", Layout((8, 2, 4, 2, 4, 4), (1, 8, 16, 64, 128, 512)), "rmem" + ] + ) + with Mesh( + (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") + ) as scope_5: + lhs_2_stages = (T.tensor_view(4096, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(5120, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_2_stages = (T.tensor_view(3072, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(3584, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) + rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads: + T.fill(acc, 0.0) + for k in range(0, 32, 16): + with scope_5[:1, :32] as scope: + tile = T.tensor_view( + T.ptr_of(c[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_1: + T.copy_async_tensor(tile, lhs_2_stages[(k // 16) % 2]) + tile_1 = T.tensor_view( + T.ptr_of(d[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_2: + T.copy_async_tensor(tile_1, rhs_2_stages[(k // 16) % 2]) + tile_2 = T.tensor_view( + T.ptr_of(a[0:0 + 64, k:k + 16]), + layout=Layout((64, 16), (32, 1)), + shape=(64, 16), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_3: + T.copy_async_tensor(tile_2, lhs_stages[(k // 16) % 2]) + tile_3 = T.tensor_view( + T.ptr_of(b[k:k + 16, 0:0 + 32]), + layout=Layout((16, 32), (32, 1)), + shape=(16, 32), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=0, + outer=Layout((32,), (1,)), +), names=("d0",) + ) as threads_4: + T.copy_async_tensor(tile_3, rhs_stages[(k // 16) % 2]) + with scope_5[1:] as scope_2: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_5: + for o_m in range(0, 64, 64): + for o_n in range(0, 32, 32): + for o_k in range(0, 16, 16): + acc_view = T.tensor_view( + T.ptr_of(acc[o_m:o_m + 64, o_n:o_n + 32]), + layout=((8 @ threads_5.d1, 2, 4 @ threads_5.d0, 2, 4 @ threads_5.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view = T.tensor_view( + T.ptr_of(lhs_stages[(k // 16) % 2][o_m:o_m + 64, o_k:o_k + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_5, + ), + shape=(64, 16), + ) + rhs_view = T.tensor_view( + T.ptr_of(rhs_stages[(k // 16) % 2][o_k:o_k + 16, o_n:o_n + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_5, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view, + lhs_view, + rhs_view, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_5), + ) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_6: + for o_m_1 in range(0, 64, 64): + for o_n_1 in range(0, 32, 32): + for o_k_1 in range(0, 16, 16): + acc_view_1 = T.tensor_view( + T.ptr_of(acc[o_m_1:o_m_1 + 64, o_n_1:o_n_1 + 32]), + layout=((8 @ threads_6.d1, 2, 4 @ threads_6.d0, 2, 4 @ threads_6.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + lhs_view_1 = T.tensor_view( + T.ptr_of(lhs_2_stages[(k // 16) % 2][o_m_1:o_m_1 + 64, o_k_1:o_k_1 + 16]), + layout=ShardLayout( + layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), + attrs=(B(), B(), B()), + mesh=threads_6, + ), + shape=(64, 16), + ) + rhs_view_1 = T.tensor_view( + T.ptr_of(rhs_2_stages[(k // 16) % 2][o_k_1:o_k_1 + 16, o_n_1:o_n_1 + 32]), + layout=ShardLayout( + layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), + attrs=(B(), B(), B()), + mesh=threads_6, + ), + shape=(16, 32), + ) + T.tiled_mma( + acc_view_1, + lhs_view_1, + rhs_view_1, + atom=T.cuda.sm90.Wgmma(n=32, form=T.cuda.sm90.Form.SS, a_major=T.cuda.sm90.Major.K, mesh=threads_6), + ) + with scope_5[1:] as scope_4: + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_7: + value_view = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_7.d1, 2, 4 @ threads_7.d0, 2, 4 @ threads_7.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.cast(acc, value_view) + with Mesh( + (Topology("thread", 256),), ComposedLayout( + inner=None, + offset=128, + outer=Layout((4, 8, 4), (32, 4, 1)), +), names=("d0", "d1", "d2") + ) as threads_8: + value_view_1 = T.tensor_view( + T.ptr_of(value[0:0 + 64, 0:0 + 32]), + layout=((8 @ threads_8.d1, 2, 4 @ threads_8.d0, 2, 4 @ threads_8.d2, 4), (1, 8, 16, 64, 128, 512)), + shape=(64, 32), + ) + T.copy(value_view_1, out) diff --git a/tests/ir/types/test_mesh.py b/tests/ir/types/test_mesh.py index 51b4b55e..0fd7680e 100644 --- a/tests/ir/types/test_mesh.py +++ b/tests/ir/types/test_mesh.py @@ -84,9 +84,31 @@ def test_make_mesh_appends_a_sliced_scope_with_its_offset() -> None: assert not covered_by_scope(make_mesh(CTA, THR[0:128]), make_mesh(CTA, THR[128:256])) -def test_make_mesh_refuses_a_sliced_suffix_replacement() -> None: - with pytest.raises(ValueError): - make_mesh(CT, THR[128:256]) +def test_make_mesh_replaces_a_sliced_suffix_by_component() -> None: + cta = Mesh( + (Topology("cta", 4352),), + Layout((64, 68), (68, 1)), + ) + current = make_mesh(cta[2:3, :], THR) + inner = Mesh( + (Topology("thread", 384),), + ComposedLayout(None, 128, Layout((4, 8, 4), (32, 4, 1))), + ) + + replaced = make_mesh(current, inner) + + assert replaced.layout == ComposedLayout( + None, + 52352, + Layout(((1, 68), (4, 8, 4)), ((68, 1), (32, 4, 1))), + ) + assert separate(replaced) == ( + Mesh( + (Topology("cta", 4352),), + ComposedLayout(None, 136, Layout((1, 68), (68, 1))), + ), + inner, + ) def test_mesh_with_several_levels_slices_in_device_numbering() -> None: diff --git a/tests/schedule/test_fixtures.py b/tests/schedule/test_fixtures.py index 441242c6..11e36610 100644 --- a/tests/schedule/test_fixtures.py +++ b/tests/schedule/test_fixtures.py @@ -10,11 +10,16 @@ from __future__ import annotations import importlib +import importlib.util +from pathlib import Path import pytest from tilefoundry.analysis.api import analyze from tilefoundry.analysis.check import check_program +from tilefoundry.inspection import as_script +from tilefoundry.ir.tir import PrimFunction +from tilefoundry.ir.tir.verify import verify_prim_function PLAIN = ( "gemm_8192x17408x5120_cta_grid", @@ -22,6 +27,15 @@ "gemm_relu_gemm_tiled", "gemm_relu_gemm_untiled", ) +TIR = tuple(sorted((Path(__file__).parents[1] / "fixtures" / "schedule" / "tir").glob("*.py"))) + + +def _prim_in(path: Path) -> PrimFunction: + spec = importlib.util.spec_from_file_location(path.stem, path) + assert spec is not None and spec.loader is not None + loaded = importlib.util.module_from_spec(spec) + spec.loader.exec_module(loaded) + return next(value for value in vars(loaded).values() if isinstance(value, PrimFunction)) @pytest.mark.parametrize("name", PLAIN) @@ -34,3 +48,10 @@ def test_plain_program_is_analyzable(name: str) -> None: check_program(program, entry) result = analyze(program, entry, analysis=("memory", "performance")) assert result.metadata_types + + +@pytest.mark.parametrize("path", TIR, ids=lambda path: path.stem) +def test_tir_program_is_verified_and_canonical(path: Path) -> None: + function = _prim_in(path) + verify_prim_function(function) + assert as_script(function) == path.read_text() From bcfc16b7d3bda5b308b858c5332ca5695cc9611b Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 12:28:33 +0800 Subject: [PATCH 06/19] refactor(layout): isolate swizzle algebra --- docs/spec/code-organization.md | 2 +- docs/spec/shard.md | 13 ++ .../codegen/cuda/tir/memory/tensor_view.py | 8 +- src/tilefoundry/ir/types/layout_algebra.py | 131 ++-------------- src/tilefoundry/ir/types/swizzle_layout.py | 144 ++++++++++++++++++ 5 files changed, 173 insertions(+), 125 deletions(-) create mode 100644 src/tilefoundry/ir/types/swizzle_layout.py diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 7b5d0786..09abf73c 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -20,7 +20,7 @@ truth for the directory's structure and invariants. | `ir/core/` | [core-ir](./core-ir.md) | Shared node algebra: `Module` / `Expr` / `Var` / `Constant` / `Tuple` / `Op` / `Call` / `Stmt` (base class) / `OpSchema` / `ParamDef` / call-graph and ownership queries / typed metadata attach-detach and diagnostics / `@register_op` / `@register_alias` / `op_registry` / `errors`. | | `ir/pattern/` | [core-ir](./core-ir.md) | Operation-declaration predicates: composable pattern values in `pattern.py`, match/binding and rendering mechanics in `match.py`, cross-operand relations in `constraint.py`, and pattern construction/specialization helpers in `utils.py`. | | `ir/types/` | [types](./types.md) | Type-system root: `Type` / `TensorType` / `TupleType` / `UnitType` / `CallableType` / `DType` / `StorageKind` / `resolve_storage` / local projections (`local_type_of`) / tensor-leaf, byte-by-storage, and topology-extent queries / `dim.*` (with their typeinfer). | -| `ir/types/{int_tuple,stride,layout,layout_algebra,shard_layout,mesh}.py` | [shard](./shard.md) | `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`), filed as CuTe files them: int tuples (`flatten` / `unflatten` / `repeat_like` / `product`), strides (`compact_major` / `idx2crd` / `crd2idx`), layouts and the algebra over them each in their own module; mesh construction, separation, and topology-bound checking stay with `Mesh`. | +| `ir/types/{int_tuple,stride,layout,layout_algebra,swizzle_layout,shard_layout,mesh}.py` | [shard](./shard.md) | `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`), filed as CuTe files them: int tuples (`flatten` / `unflatten` / `repeat_like` / `product`), strides (`compact_major` / `idx2crd` / `crd2idx`), layouts, general layout algebra, and swizzle-specialized layout algebra each in their own module; mesh construction, separation, and topology-bound checking stay with `Mesh`. | | `ir/mesh_scope.py` | [shard](./shard.md) | Which scope a statement stands inside and what it admits: `device_layout`, `covered_by_scope`. Neither a type nor a visitor, so it sits beside `ir/isl_interop.py` rather than in either. | | `ir/clause/` | [parser](./parser.md) | Authored `where(layout=..., mesh=..., storage=...)` constraint records: the shared base plus layout, mesh, and storage constraints, attached by the parser and read back by the Python printer. | | `ir/visitor.py` | [visitor-mutator](./visitor-mutator.md) | `ExprFunctor` / `ExprVisitor` / `ExprWalker` / `ExprCollector` / `ExprCloner` / `BindingSubstitutionCloner` / `StmtVisitor` / `StmtMutator` / `StmtExprMutator`, plus `collect_exprs`, value-operand/function-value queries, and the canonical `PrimFunction` walk and rewrite entries. | diff --git a/docs/spec/shard.md b/docs/spec/shard.md index 9eec878d..6756bdd5 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -231,6 +231,14 @@ class Swizzle: shift: int def __call__(self, offset: int) -> int: ... + + +def get_swizzle_portion(layout: Layout | ComposedLayout) -> Swizzle | None: + """Return the final swizzle of a composed layout, when present.""" + + +def make_swizzle(active_y: int, active_z: int) -> Swizzle: + """Build the representable CuTe swizzle that maps Y bits onto Z bits.""" ``` **Terms.** The *Y bits* are the ones read out of the index @@ -246,6 +254,11 @@ the ones they are XORed onto - Because the two ranges do not overlap, a `Swizzle` is an involution: `swizzle(swizzle(offset)) == offset`, so it is its own left and right inverse. `bits == 0` is the identity. + - Swizzle-specialized `composition`, `left_inverse`, and `right_inverse` + MUST follow the corresponding CuTe `swizzle_layout.hpp` overloads. + `get_swizzle_portion` returns `None` for a layout with no final swizzle; + `make_swizzle` MUST reject mask pairs that are not equal-width contiguous + runs rather than inventing a representation. - A `Swizzle` is a mapping on an index, not a layout. It is not a `LayoutBase`, it states no `shape`, and it MUST NOT be a `TensorType.layout` or a `ShardLayout.layout` on its own. It reaches a diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py index f92816b2..4990fa87 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -27,7 +27,6 @@ from tilefoundry.ir.types import ComposedLayout from tilefoundry.ir.types.dim import DimAdd, DimMul, DimSub, DimVar from tilefoundry.ir.types.layout import Layout, LayoutBase, flatten -from tilefoundry.ir.types.layout_algebra import swizzle_of from tilefoundry.ir.types.shard_layout import ( Broadcast, Dynamic, @@ -38,6 +37,7 @@ from tilefoundry.ir.types.shard_layout import ShardLayout as SL from tilefoundry.ir.types.storage import StorageKind from tilefoundry.ir.types.stride import compact_row_major +from tilefoundry.ir.types.swizzle_layout import get_swizzle_portion from tilefoundry.ir.types.utils import shape_numel_upper_bound, upper_bound from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.target import CudaTarget @@ -57,7 +57,7 @@ def _render_layout_type(layout: LayoutBase) -> str: underneath. The XOR mapping is not expressible as strides, so there is no affine form to fall back to. """ - swizzle = swizzle_of(layout) + swizzle = get_swizzle_portion(layout) if swizzle is not None: return ( f"cute::ComposedLayout, " @@ -76,7 +76,7 @@ def _render_layout_value(layout: LayoutBase, dim, stride) -> str: *dim* and *stride* render one shape entry and one stride entry, which is where a runtime-provided extent reaches the emitted layout. """ - swizzle = swizzle_of(layout) + swizzle = get_swizzle_portion(layout) if swizzle is not None: return ( f"cute::make_composed_layout(" @@ -204,7 +204,7 @@ def render_shard_layout_value(var_name: str, sl: SL, dynamic_extents=None, stora """ sll = sl.layout if storage is StorageKind.RMEM: - if swizzle_of(sll) is not None: + if get_swizzle_portion(sll) is not None: raise NotImplementedError( "render_shard_layout_value: a Swizzle states how a shared-memory " "bank pattern is arranged; a register engine has no such addresses, " diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index bac418da..1e1ced05 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -15,8 +15,10 @@ from tilefoundry.ir.types.layout import flatten +from . import swizzle_layout from .layout import ComposedLayout, Layout, Swizzle, size from .stride import compact_col_major, idx2crd +from .swizzle_layout import get_swizzle_portion class NotProjectable(ValueError): @@ -46,17 +48,6 @@ def _stride(layout: Layout) -> tuple[int, ...]: return compact_col_major(_shape(layout)) -def swizzle_of(layout: object) -> Optional[Swizzle]: - """The ``Swizzle`` a composed layout applies last, or ``None``. - - CuTe ``get_swizzle_portion``, answering ``None`` rather than the identity - ``Swizzle<0,4,3>`` so a caller can branch on "is this swizzled at all". - """ - if isinstance(layout, ComposedLayout) and isinstance(layout.inner, Swizzle): - return layout.inner - return None - - def apply(layout: Union[Layout, ComposedLayout], coord: int) -> int: """``crd2idx`` of a 1-D domain coord: decompose by shape, dot with strides. @@ -76,29 +67,6 @@ def apply(layout: Union[Layout, ComposedLayout], coord: int) -> int: return idx -def _crd2idx_unbounded(layout: Layout, coord: int) -> int: - """``apply`` as CuTe's ``crd2idx`` has it: the last mode is not wrapped. - - ``apply`` wraps every mode, which is the same answer for an in-domain - coord and the one this module's callers want. The swizzle composition - rules feed a *bit mask* through a layout instead, which is routinely - larger than the domain, and CuTe leaves the final mode unwrapped so those - high bits keep contributing. Only that port reads this. - """ - shape = _shape(layout) - stride = _stride(layout) - idx = 0 - rem = coord - last = len(shape) - 1 - for position, (s, d) in enumerate(zip(shape, stride)): - if position == last: - idx += rem * d - else: - idx += (rem % s) * d - rem //= s - return idx - - def cosize(layout: Union[Layout, ComposedLayout]) -> int: """The codomain extent. @@ -106,7 +74,7 @@ def cosize(layout: Union[Layout, ComposedLayout]) -> int: so it adds nothing to it: CuTe ``cosize`` of a swizzled composed layout is ``cosize`` of the layout underneath (``swizzle_layout.hpp:172``). """ - if swizzle_of(layout) is not None: + if get_swizzle_portion(layout) is not None: return cosize(layout.outer) return apply(layout, size(layout) - 1) + 1 @@ -117,7 +85,7 @@ def coalesce(layout: Union[Layout, ComposedLayout]): Coalescing renames the domain and leaves the index mapping alone, so a swizzled composed layout coalesces underneath its swizzle. """ - if swizzle_of(layout) is not None: + if get_swizzle_portion(layout) is not None: return ComposedLayout( inner=layout.inner, offset=layout.offset, outer=coalesce(layout.outer) ) @@ -275,87 +243,15 @@ def _check_admissible(scope: ComposedLayout) -> None: raise NotProjectable("outer layout is not inverse-projectable (injective + compact)") -def _make_swizzle(active_y: int, active_z: int) -> Swizzle: - """CuTe ``make_swizzle()``: the swizzle that XORs *Y* onto *Z*. - - The two masks must hold the same number of bits; their trailing-zero - counts give ``base`` and the signed ``shift``, and the reconstructed - ``swizzle_code`` must give the masks back, which is how CuTe checks that - the pair is a swizzle it can represent at all. - """ - bits_y, bits_z = active_y.bit_count(), active_z.bit_count() - if bits_y != bits_z: - raise NotImplementedError( - f"composition: the Y mask {active_y:#x} holds {bits_y} bits and the Z " - f"mask {active_z:#x} holds {bits_z}; only an equal-width pair is a " - f"Swizzle" - ) - if bits_y == 0: - return Swizzle(0, 0, 0) - trailing_y = (active_y & -active_y).bit_length() - 1 - trailing_z = (active_z & -active_z).bit_length() - 1 - swizzle = Swizzle(bits_y, min(trailing_y, trailing_z), trailing_y - trailing_z) - if swizzle.swizzle_code != (active_y | active_z): - raise NotImplementedError( - f"composition: the mask pair ({active_y:#x}, {active_z:#x}) is not a " - f"Swizzle; its bits are not two contiguous equal-width runs" - ) - return swizzle - - def composition(left, right, offset: int = 0): - """CuTe ``composition`` for the swizzle cases (``swizzle_layout.hpp:302``). - - ``composition(Swizzle, Layout)`` builds a swizzled layout, which states - ``Swizzle(offset + Layout(coord))``. - - ``composition(Layout, Swizzle)`` would otherwise want the ``Swizzle`` in - ``outer``, which has no domain to be a domain-side component of. CuTe - instead reads which of the swizzle's bits the layout leaves active, - rebuilds a ``Swizzle`` over those, and puts it back on the inner side. - """ - if isinstance(left, Swizzle) and isinstance(right, Layout): - if left.bits == 0 and offset == 0: - return right - return ComposedLayout(inner=left, offset=offset, outer=right) - if isinstance(left, Layout) and isinstance(right, Swizzle): - if offset: - raise NotImplementedError( - f"composition: a non-zero offset ({offset}) between a Layout and a " - f"Swizzle has no canonical ComposedLayout form" - ) - active_y = _crd2idx_unbounded(left, right.yyy_mask) - active_z = _crd2idx_unbounded(left, right.zzz_mask) - return composition(_make_swizzle(active_y, active_z), left) + """CuTe ``composition``, dispatched to the supported overloads.""" + if swizzle_layout._supports_composition(left, right): + return swizzle_layout.composition(left, right, offset) raise NotImplementedError( f"composition: no rule for {type(left).__name__} ∘ {type(right).__name__}" ) -def _swizzled_inverse(layout: ComposedLayout, inverse_of_layout): - """CuTe's swizzled ``left_inverse``/``right_inverse`` (``swizzle_layout.hpp:344``). - - ``inverse(Swizzle(offset + outer(c)))`` passes the swizzle back to the - left of the inverted ``outer``, which ``composition(Layout, Swizzle)`` - then canonicalizes back into this IR's one legal shape. CuTe's non-zero - ``offset`` branch composes ``inverse(offset)`` between the two, which - lands a bare ``Swizzle`` in ``outer``; that is not a layout, so this - refuses it by name rather than building something unrepresentable. - """ - if layout.offset != 0: - raise NotImplementedError( - f"inverse: a swizzled composed layout with a non-zero offset " - f"({layout.offset}) inverts to a Swizzle on the domain side, which " - f"ComposedLayout.outer cannot hold" - ) - if not isinstance(layout.outer, Layout): - raise NotImplementedError( - f"inverse: a swizzled composed layout inverts through its outer " - f"Layout; this one states {type(layout.outer).__name__}" - ) - return composition(inverse_of_layout(layout.outer), layout.inner) - - def left_inverse(layout: Union[Layout, ComposedLayout, Swizzle]): """CuTe ``left_inverse``, dispatched. @@ -366,10 +262,8 @@ def left_inverse(layout: Union[Layout, ComposedLayout, Swizzle]): outer=None)`` (``outer=None`` ≡ identity), i.e. ``image⁻¹(t) = outer⁻¹(t − offset)``. """ - if isinstance(layout, Swizzle): - return layout - if swizzle_of(layout) is not None: - return _swizzled_inverse(layout, left_inverse) + if swizzle_layout._supports_inverse(layout): + return swizzle_layout.left_inverse(layout) if isinstance(layout, ComposedLayout): _check_admissible(layout) return ComposedLayout( @@ -386,10 +280,8 @@ def right_inverse(layout: Union[Layout, ComposedLayout, Swizzle]): A ``Swizzle`` is an involution -- its Y and Z bit ranges do not overlap -- so it is its own inverse on both sides (``swizzle_layout.hpp:371``). """ - if isinstance(layout, Swizzle): - return layout - if swizzle_of(layout) is not None: - return _swizzled_inverse(layout, right_inverse) + if swizzle_layout._supports_inverse(layout): + return swizzle_layout.right_inverse(layout) if isinstance(layout, ComposedLayout): _check_admissible(layout) return ComposedLayout( @@ -455,7 +347,6 @@ def contains(scope: ComposedLayout, t: int) -> bool: "ASYNC_WIDTHS", "NotProjectable", "Run", - "swizzle_of", "composition", "cosize", "apply", diff --git a/src/tilefoundry/ir/types/swizzle_layout.py b/src/tilefoundry/ir/types/swizzle_layout.py new file mode 100644 index 00000000..0ac3be95 --- /dev/null +++ b/src/tilefoundry/ir/types/swizzle_layout.py @@ -0,0 +1,144 @@ +"""Provide CuTe layout algebra specializations involving a ``Swizzle``. + +The operations here mirror ``cute/swizzle_layout.hpp``. The ``Swizzle`` +value type itself remains in :mod:`tilefoundry.ir.types.layout`, alongside +the other layout value types. +""" + +from __future__ import annotations + +from typing import Optional + +from .int_tuple import flatten +from .layout import ComposedLayout, Layout, Swizzle +from .stride import compact_col_major + + +def _shape(layout: Layout) -> tuple[int, ...]: + return flatten(layout.shape) + + +def _stride(layout: Layout) -> tuple[int, ...]: + if layout.strides is not None: + return layout.strides + return compact_col_major(_shape(layout)) + + +def _crd2idx_unbounded(layout: Layout, coord: int) -> int: + """Apply CuTe ``crd2idx`` while leaving the final mode unwrapped.""" + shape = _shape(layout) + stride = _stride(layout) + idx = 0 + rem = coord + last = len(shape) - 1 + for position, (extent, step) in enumerate(zip(shape, stride)): + if position == last: + idx += rem * step + else: + idx += (rem % extent) * step + rem //= extent + return idx + + +def get_swizzle_portion(layout: object) -> Optional[Swizzle]: + """Return the final ``Swizzle`` of a composed layout, or ``None``. + + This mirrors CuTe ``get_swizzle_portion`` while using ``None`` rather + than an identity ``Swizzle<0,4,3>`` for a non-swizzled layout. + """ + if isinstance(layout, ComposedLayout) and isinstance(layout.inner, Swizzle): + return layout.inner + return None + + +def make_swizzle(active_y: int, active_z: int) -> Swizzle: + """CuTe ``make_swizzle()``: build the swizzle XORing *Y* onto *Z*.""" + bits_y, bits_z = active_y.bit_count(), active_z.bit_count() + if bits_y != bits_z: + raise NotImplementedError( + f"composition: the Y mask {active_y:#x} holds {bits_y} bits and the Z " + f"mask {active_z:#x} holds {bits_z}; only an equal-width pair is a " + f"Swizzle" + ) + if bits_y == 0: + return Swizzle(0, 0, 0) + trailing_y = (active_y & -active_y).bit_length() - 1 + trailing_z = (active_z & -active_z).bit_length() - 1 + swizzle = Swizzle(bits_y, min(trailing_y, trailing_z), trailing_y - trailing_z) + if swizzle.swizzle_code != (active_y | active_z): + raise NotImplementedError( + f"composition: the mask pair ({active_y:#x}, {active_z:#x}) is not a " + f"Swizzle; its bits are not two contiguous equal-width runs" + ) + return swizzle + + +def _supports_composition(left: object, right: object) -> bool: + return (isinstance(left, Swizzle) and isinstance(right, Layout)) or ( + isinstance(left, Layout) and isinstance(right, Swizzle) + ) + + +def composition(left, right, offset: int = 0): + """CuTe ``composition`` overloads involving a ``Swizzle``.""" + if isinstance(left, Swizzle) and isinstance(right, Layout): + if left.bits == 0 and offset == 0: + return right + return ComposedLayout(inner=left, offset=offset, outer=right) + if isinstance(left, Layout) and isinstance(right, Swizzle): + if offset: + raise NotImplementedError( + f"composition: a non-zero offset ({offset}) between a Layout and a " + f"Swizzle has no canonical ComposedLayout form" + ) + active_y = _crd2idx_unbounded(left, right.yyy_mask) + active_z = _crd2idx_unbounded(left, right.zzz_mask) + return composition(make_swizzle(active_y, active_z), left) + raise NotImplementedError( + f"composition: no swizzle rule for {type(left).__name__} ∘ {type(right).__name__}" + ) + + +def _supports_inverse(layout: object) -> bool: + return isinstance(layout, Swizzle) or get_swizzle_portion(layout) is not None + + +def _inverse(layout, inverse_of_layout): + """CuTe's composed swizzle inverse, shared by both inverse overloads.""" + if isinstance(layout, Swizzle): + return layout + if layout.offset != 0: + raise NotImplementedError( + f"inverse: a swizzled composed layout with a non-zero offset " + f"({layout.offset}) inverts to a Swizzle on the domain side, which " + f"ComposedLayout.outer cannot hold" + ) + if not isinstance(layout.outer, Layout): + raise NotImplementedError( + f"inverse: a swizzled composed layout inverts through its outer " + f"Layout; this one states {type(layout.outer).__name__}" + ) + return composition(inverse_of_layout(layout.outer), layout.inner) + + +def left_inverse(layout): + """CuTe's swizzled ``left_inverse`` overload.""" + from .layout_algebra import left_inverse as inverse_of_layout # noqa: PLC0415 + + return _inverse(layout, inverse_of_layout) + + +def right_inverse(layout): + """CuTe's swizzled ``right_inverse`` overload.""" + from .layout_algebra import right_inverse as inverse_of_layout # noqa: PLC0415 + + return _inverse(layout, inverse_of_layout) + + +__all__ = [ + "composition", + "get_swizzle_portion", + "left_inverse", + "make_swizzle", + "right_inverse", +] From ac235bae33b65552d6f95bc3b7e489265afa037a Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 12:34:05 +0800 Subject: [PATCH 07/19] refactor(layout): keep swizzle dependencies one-way --- src/tilefoundry/ir/types/layout.py | 38 +++++++++++ src/tilefoundry/ir/types/layout_algebra.py | 44 +++++------- src/tilefoundry/ir/types/swizzle_layout.py | 79 ++++++---------------- 3 files changed, 76 insertions(+), 85 deletions(-) diff --git a/src/tilefoundry/ir/types/layout.py b/src/tilefoundry/ir/types/layout.py index 89504d1b..96157776 100644 --- a/src/tilefoundry/ir/types/layout.py +++ b/src/tilefoundry/ir/types/layout.py @@ -5,6 +5,7 @@ from .int_tuple import flatten as _flat from .int_tuple import product +from .stride import compact_col_major class LayoutBase: @@ -84,6 +85,28 @@ def __call__(self, offset: int) -> int: return offset ^ moved +def make_swizzle(active_y: int, active_z: int) -> Swizzle: + """CuTe ``make_swizzle()``: build the swizzle XORing *Y* onto *Z*.""" + bits_y, bits_z = active_y.bit_count(), active_z.bit_count() + if bits_y != bits_z: + raise NotImplementedError( + f"composition: the Y mask {active_y:#x} holds {bits_y} bits and the Z " + f"mask {active_z:#x} holds {bits_z}; only an equal-width pair is a " + f"Swizzle" + ) + if bits_y == 0: + return Swizzle(0, 0, 0) + trailing_y = (active_y & -active_y).bit_length() - 1 + trailing_z = (active_z & -active_z).bit_length() - 1 + swizzle = Swizzle(bits_y, min(trailing_y, trailing_z), trailing_y - trailing_z) + if swizzle.swizzle_code != (active_y | active_z): + raise NotImplementedError( + f"composition: the mask pair ({active_y:#x}, {active_z:#x}) is not a " + f"Swizzle; its bits are not two contiguous equal-width runs" + ) + return swizzle + + @dataclass(frozen=True) class ComposedLayout(LayoutBase): """Represent ``image(c) = inner(offset + outer(c))``. @@ -112,6 +135,18 @@ def shape(self) -> tuple: EMPTY_LAYOUT = Layout(shape=(), strides=()) +def flat_shape(layout: Layout) -> tuple[int, ...]: + """Return CuTe ``flatten(layout.shape())`` as a flat tuple.""" + return _flat(layout.shape) + + +def flat_stride(layout: Layout) -> tuple[int, ...]: + """Flatten stated strides, or synthesize the compact column-major default.""" + if layout.strides is not None: + return _flat(layout.strides) + return compact_col_major(flat_shape(layout)) + + def size(layout: Layout) -> int: return product(layout.shape) @@ -178,7 +213,10 @@ def take(layout: LayoutBase, begin: int, end: int) -> "Layout": "Swizzle", "ComposedLayout", "EMPTY_LAYOUT", + "flat_shape", + "flat_stride", "get", + "make_swizzle", "rank", "take", ] diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index 1e1ced05..d7905447 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -16,7 +16,7 @@ from tilefoundry.ir.types.layout import flatten from . import swizzle_layout -from .layout import ComposedLayout, Layout, Swizzle, size +from .layout import ComposedLayout, Layout, Swizzle, flat_shape, flat_stride, size from .stride import compact_col_major, idx2crd from .swizzle_layout import get_swizzle_portion @@ -38,16 +38,6 @@ class Run: mode: int -def _shape(layout: Layout) -> tuple[int, ...]: - return flatten(layout.shape) - - -def _stride(layout: Layout) -> tuple[int, ...]: - if layout.strides is not None: - return layout.strides - return compact_col_major(_shape(layout)) - - def apply(layout: Union[Layout, ComposedLayout], coord: int) -> int: """``crd2idx`` of a 1-D domain coord: decompose by shape, dot with strides. @@ -57,8 +47,8 @@ def apply(layout: Union[Layout, ComposedLayout], coord: int) -> int: """ if isinstance(layout, ComposedLayout): return _apply_any(layout, coord) - shape = _shape(layout) - stride = _stride(layout) + shape = flat_shape(layout) + stride = flat_stride(layout) idx = 0 rem = coord for s, d in zip(shape, stride): @@ -91,7 +81,7 @@ def coalesce(layout: Union[Layout, ComposedLayout]): ) result_shape: list[int] = [1] result_stride: list[int] = [0] - for shape, stride in zip(_shape(layout), _stride(layout)): + for shape, stride in zip(flat_shape(layout), flat_stride(layout)): if shape == 1: continue if result_shape[-1] == 1: @@ -150,7 +140,7 @@ def complement(layout: Layout, max_idx: int = 1) -> Layout: result_shape: list[int] = [] result_stride: list[int] = [] current_idx = 1 - for stride, shape in sorted(zip(_stride(layout), _shape(layout))): + for stride, shape in sorted(zip(flat_stride(layout), flat_shape(layout))): if stride == 0 or shape == 1: continue if current_idx > shape * stride: @@ -165,7 +155,7 @@ def complement(layout: Layout, max_idx: int = 1) -> Layout: def _make_flat(a: Layout, b: Layout) -> Layout: """Concatenate two flat layouts into one (CuTe ``make_layout`` after flatten).""" - return Layout(shape=_shape(a) + _shape(b), strides=_stride(a) + _stride(b)) + return Layout(shape=flat_shape(a) + flat_shape(b), strides=flat_stride(a) + flat_stride(b)) def is_inverse_projectable(layout: Layout) -> bool: @@ -182,7 +172,9 @@ def is_inverse_projectable(layout: Layout) -> bool: """ current = 1 modes = sorted( - (stride, shape) for shape, stride in zip(_shape(layout), _stride(layout)) if shape != 1 + (stride, shape) + for shape, stride in zip(flat_shape(layout), flat_stride(layout)) + if shape != 1 ) for stride, shape in modes: if stride == 0 or stride % current != 0: @@ -196,8 +188,8 @@ def _right_inverse_layout(layout: Layout) -> Layout: result_shape: list[int] = [] result_stride: list[int] = [] current_idx = 1 - shape = _shape(layout) - stride = _stride(layout) + shape = flat_shape(layout) + stride = flat_stride(layout) triples = sorted(zip(stride, shape, compact_col_major(shape))) for st, sh, rstride in triples: if sh == 1: @@ -224,7 +216,7 @@ def _is_identity_inner(inner: object) -> bool: if inner is None: return True if isinstance(inner, Layout): - return _stride(inner) == compact_col_major(_shape(inner)) + return flat_stride(inner) == compact_col_major(flat_shape(inner)) return False @@ -245,7 +237,7 @@ def _check_admissible(scope: ComposedLayout) -> None: def composition(left, right, offset: int = 0): """CuTe ``composition``, dispatched to the supported overloads.""" - if swizzle_layout._supports_composition(left, right): + if swizzle_layout.supports_composition(left, right): return swizzle_layout.composition(left, right, offset) raise NotImplementedError( f"composition: no rule for {type(left).__name__} ∘ {type(right).__name__}" @@ -262,8 +254,8 @@ def left_inverse(layout: Union[Layout, ComposedLayout, Swizzle]): outer=None)`` (``outer=None`` ≡ identity), i.e. ``image⁻¹(t) = outer⁻¹(t − offset)``. """ - if swizzle_layout._supports_inverse(layout): - return swizzle_layout.left_inverse(layout) + if swizzle_layout.supports_inverse(layout): + return swizzle_layout.inverse(layout, left_inverse) if isinstance(layout, ComposedLayout): _check_admissible(layout) return ComposedLayout( @@ -280,8 +272,8 @@ def right_inverse(layout: Union[Layout, ComposedLayout, Swizzle]): A ``Swizzle`` is an involution -- its Y and Z bit ranges do not overlap -- so it is its own inverse on both sides (``swizzle_layout.hpp:371``). """ - if swizzle_layout._supports_inverse(layout): - return swizzle_layout.right_inverse(layout) + if swizzle_layout.supports_inverse(layout): + return swizzle_layout.inverse(layout, right_inverse) if isinstance(layout, ComposedLayout): _check_admissible(layout) return ComposedLayout( @@ -334,7 +326,7 @@ def project(scope: ComposedLayout, t: int) -> Optional[tuple[int, ...]]: if image(scope, coord_1d) != t: return None - shape = _shape(outer) + shape = flat_shape(outer) return idx2crd(coord_1d, shape, compact_col_major(shape)) diff --git a/src/tilefoundry/ir/types/swizzle_layout.py b/src/tilefoundry/ir/types/swizzle_layout.py index 0ac3be95..e3e13c3f 100644 --- a/src/tilefoundry/ir/types/swizzle_layout.py +++ b/src/tilefoundry/ir/types/swizzle_layout.py @@ -7,27 +7,22 @@ from __future__ import annotations -from typing import Optional +from typing import Callable, Optional -from .int_tuple import flatten -from .layout import ComposedLayout, Layout, Swizzle -from .stride import compact_col_major - - -def _shape(layout: Layout) -> tuple[int, ...]: - return flatten(layout.shape) - - -def _stride(layout: Layout) -> tuple[int, ...]: - if layout.strides is not None: - return layout.strides - return compact_col_major(_shape(layout)) +from .layout import ( + ComposedLayout, + Layout, + Swizzle, + flat_shape, + flat_stride, + make_swizzle, +) def _crd2idx_unbounded(layout: Layout, coord: int) -> int: """Apply CuTe ``crd2idx`` while leaving the final mode unwrapped.""" - shape = _shape(layout) - stride = _stride(layout) + shape = flat_shape(layout) + stride = flat_stride(layout) idx = 0 rem = coord last = len(shape) - 1 @@ -51,29 +46,8 @@ def get_swizzle_portion(layout: object) -> Optional[Swizzle]: return None -def make_swizzle(active_y: int, active_z: int) -> Swizzle: - """CuTe ``make_swizzle()``: build the swizzle XORing *Y* onto *Z*.""" - bits_y, bits_z = active_y.bit_count(), active_z.bit_count() - if bits_y != bits_z: - raise NotImplementedError( - f"composition: the Y mask {active_y:#x} holds {bits_y} bits and the Z " - f"mask {active_z:#x} holds {bits_z}; only an equal-width pair is a " - f"Swizzle" - ) - if bits_y == 0: - return Swizzle(0, 0, 0) - trailing_y = (active_y & -active_y).bit_length() - 1 - trailing_z = (active_z & -active_z).bit_length() - 1 - swizzle = Swizzle(bits_y, min(trailing_y, trailing_z), trailing_y - trailing_z) - if swizzle.swizzle_code != (active_y | active_z): - raise NotImplementedError( - f"composition: the mask pair ({active_y:#x}, {active_z:#x}) is not a " - f"Swizzle; its bits are not two contiguous equal-width runs" - ) - return swizzle - - -def _supports_composition(left: object, right: object) -> bool: +def supports_composition(left: object, right: object) -> bool: + """Return whether the operands select a CuTe swizzle composition overload.""" return (isinstance(left, Swizzle) and isinstance(right, Layout)) or ( isinstance(left, Layout) and isinstance(right, Swizzle) ) @@ -99,12 +73,13 @@ def composition(left, right, offset: int = 0): ) -def _supports_inverse(layout: object) -> bool: +def supports_inverse(layout: object) -> bool: + """Return whether a CuTe swizzle inverse overload accepts *layout*.""" return isinstance(layout, Swizzle) or get_swizzle_portion(layout) is not None -def _inverse(layout, inverse_of_layout): - """CuTe's composed swizzle inverse, shared by both inverse overloads.""" +def inverse(layout, inverse_of_layout: Callable): + """Apply the CuTe swizzle inverse using the injected general inverse.""" if isinstance(layout, Swizzle): return layout if layout.offset != 0: @@ -121,24 +96,10 @@ def _inverse(layout, inverse_of_layout): return composition(inverse_of_layout(layout.outer), layout.inner) -def left_inverse(layout): - """CuTe's swizzled ``left_inverse`` overload.""" - from .layout_algebra import left_inverse as inverse_of_layout # noqa: PLC0415 - - return _inverse(layout, inverse_of_layout) - - -def right_inverse(layout): - """CuTe's swizzled ``right_inverse`` overload.""" - from .layout_algebra import right_inverse as inverse_of_layout # noqa: PLC0415 - - return _inverse(layout, inverse_of_layout) - - __all__ = [ "composition", "get_swizzle_portion", - "left_inverse", - "make_swizzle", - "right_inverse", + "inverse", + "supports_composition", + "supports_inverse", ] From bdfce0dadfa541788bad77d1dfd3c55f2d02441d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 12:42:06 +0800 Subject: [PATCH 08/19] refactor(pattern): align shard layout matching --- docs/spec/core-ir.md | 8 +++- src/tilefoundry/ir/pattern/pattern.py | 50 ++++++---------------- src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py | 25 +++++++++-- src/tilefoundry/ir/tir/cuda/nn/wgmma.py | 17 +++++++- tests/ir/core/test_overload.py | 8 +++- tests/ir/core/test_pattern.py | 10 ++++- 6 files changed, 71 insertions(+), 47 deletions(-) diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 1d0fd80e..40194efc 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -634,14 +634,20 @@ arrangement by default. A sliced layout must be stated explicitly with With `per_mode=True`, `LayoutPattern` checks each top-level mode independently; `MeshPattern` requires this explicit form because each mesh level uses its own numbering space. `MeshPattern` never changes the supplied pattern implicitly. +`ShardLayoutPattern` names the same `layout`, `attrs`, and `mesh` fields as +`ShardLayout`: `layout` and `mesh` are nested patterns, while `attrs` remains +an exact structural value. Its mesh pattern may state bare and sliced forms +explicitly; the matcher does not normalize one into the other. Two consumer surfaces: - **Parser dispatch** — `ParamDef.pattern` ([§2.3](#23-op)) is matched against an argument's `Expr.type` during overload resolution. Subclasses used: - `ScalarPattern` (rank-0), `TensorPattern(rank?, dtype?)` (non-scalar), and + `ScalarPattern` (rank-0), `TensorPattern(shape?, dtype?)` (non-scalar), and `AndPattern(parts)` (conjunction). Two singletons are exported as convenience: `Scalar = ScalarPattern()` and `Tensor = TensorPattern()`. + A tensor rank is stated by giving `shape` that many positions; wildcard + positions constrain only the sequence length. - **Specialization dispatch** — patterns appearing in `hir.Function.specializations` ([hir.md §1.1](./hir.md#11-function)) and `tir.PrimFunction.specializations` describe which runtime diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index e2f2e859..d45727e1 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -6,7 +6,6 @@ from typing import Any from tilefoundry.ir.types import ( - Broadcast, ComposedLayout, Layout, Mesh, @@ -20,7 +19,6 @@ from tilefoundry.ir.types.layout_algebra import ( ASYNC_WIDTHS, box_runs, - frame_of, is_inverse_projectable, ) from tilefoundry.ir.types.mesh import separate @@ -686,7 +684,6 @@ def describe(self, name: str = UNNAMED_PLACE) -> str: class TensorPattern(Pattern): """Match a non-scalar ``TensorType`` field by field.""" - rank: int | None = None dtype: Any = None shape: tuple | None = None storage: Any = None @@ -696,7 +693,6 @@ def match(self, subject, captures=None): if not isinstance(subject, TensorType) or subject.shape == (): return None wanted = ( - (self.rank, len(subject.shape)), (self.dtype, subject.dtype), (self.storage, subject.storage), ) @@ -720,7 +716,6 @@ def refusal(self, subject, captures=None) -> str | None: None if self.shape is None else SequencePattern(*self.shape), tuple(subject.shape), ), - ("rank", self.rank, len(subject.shape)), ("dtype", self.dtype, subject.dtype), ("storage", self.storage, subject.storage), ("layout", self.layout, subject.layout), @@ -743,8 +738,6 @@ def describe(self, name: str = UNNAMED_PLACE, arrangements=None) -> str: stated = [] if self.shape is not None: stated.append("shape=" + written_tuple(tuple(written_place(x) for x in self.shape))) - if self.rank is not None: - stated.append(f"rank={self.rank}") if self.dtype is not None: stated.append(f"dtype={_named(self.dtype)}") if self.storage is not None: @@ -759,7 +752,6 @@ def describe(self, name: str = UNNAMED_PLACE, arrangements=None) -> str: def relations(self) -> tuple[str, ...]: return relations_of( ( - self.rank, self.dtype, self.storage, *(self.shape or ()), @@ -770,51 +762,37 @@ def relations(self) -> tuple[str, ...]: @dataclass(frozen=True) class ShardLayoutPattern(Pattern): - """Match a sharded layout's arrangement, shard attrs, and mesh frame.""" + """Match a sharded layout's layout, shard attrs, and mesh frame.""" - arrangement: object + layout: object attrs: tuple - mesh: Mesh + mesh: MeshPattern def match(self, subject, captures=None): if not isinstance(subject, ShardLayout): return None - extra = len(subject.attrs) - len(self.attrs) - if ( - extra < 0 - or subject.attrs[extra:] != self.attrs - or any(not isinstance(attr, Broadcast) for attr in subject.attrs[:extra]) + held = Match(dict(captures or {})) + for pattern, value in ( + (self.layout, subject.layout), + (self.attrs, subject.attrs), + (self.mesh, subject.mesh), ): - return None - framed = frame_of(subject.mesh.layout) - if framed is None: - return None - frame = framed[1] - if extra: - if not isinstance(frame, Layout): + held = matched(pattern, value, held.captures) + if held is None: return None - frame = Layout(frame.shape[extra:], frame.strides[extra:]) - subject_names = tuple( - getattr(topology, "name", topology) for topology in subject.mesh.topologies - ) - wanted_names = tuple( - getattr(topology, "name", topology) for topology in self.mesh.topologies - ) - if frame != self.mesh.layout or subject_names != wanted_names: - return None - return self.reads(subject.layout, captures) + return held def reads(self, layout, captures=None): - return matched(self.arrangement, layout, captures) + return matched(self.layout, layout, captures) def accepts_layout(self, layout) -> bool: return self.reads(layout) is not None def alternatives(self, bindings=()) -> tuple: - return alternatives_of(self.arrangement, bindings) + return alternatives_of(self.layout, bindings) def relations(self) -> tuple[str, ...]: - return relations_of((self.arrangement,)) + return relations_of((self.layout,)) def rules(self, arrangements=None) -> tuple[str, ...]: items = self.alternatives() if arrangements is None else tuple(arrangements) diff --git a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py index 9ea50dd7..e483cd39 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py @@ -2,7 +2,15 @@ from __future__ import annotations -from tilefoundry.ir.pattern import ShardLayoutPattern, TensorPattern, arrangement_pattern +from tilefoundry.ir.pattern import ( + ComposedLayoutPattern, + MeshPattern, + OrPattern, + ShardLayoutPattern, + TensorPattern, + WildcardPattern, + arrangement_pattern, +) from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology from tilefoundry.ir.types.storage import StorageKind as S @@ -30,6 +38,15 @@ mesh=WARP, ) +_WARP_LAYOUT = arrangement_pattern(WARP.layout, per_mode=True) +_WARP_PATTERN = MeshPattern( + ("thread",), + OrPattern( + ComposedLayoutPattern(offset=WildcardPattern(), outer=_WARP_LAYOUT), + _WARP_LAYOUT, + ), +) + def _fragment(shape: tuple, dtype, held: ShardLayout) -> TensorPattern: return TensorPattern( @@ -37,9 +54,9 @@ def _fragment(shape: tuple, dtype, held: ShardLayout) -> TensorPattern: dtype=dtype, storage=S.RMEM, layout=ShardLayoutPattern( - arrangement_pattern(held.layout), - held.attrs, - WARP, + layout=arrangement_pattern(held.layout), + attrs=held.attrs, + mesh=_WARP_PATTERN, ), ) diff --git a/src/tilefoundry/ir/tir/cuda/nn/wgmma.py b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py index 8544dce5..38ef3a13 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/wgmma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py @@ -11,13 +11,17 @@ ComposedLayoutPattern, GuardPattern, LayoutPattern, + MeshPattern, MultipleOfPattern, OneOfPattern, + OrPattern, RangePattern, ShardLayoutPattern, SwitchPattern, SwizzlePattern, TensorPattern, + WildcardPattern, + arrangement_pattern, ) from tilefoundry.ir.pattern.match import is_symbolic from tilefoundry.ir.types import Broadcast, DType, Layout, Mesh, Split, Topology @@ -145,13 +149,22 @@ def Fragment(rows: int, cols) -> LayoutPattern: SHARED_BY_ALL = (Broadcast(), Broadcast(), Broadcast()) HELD_PER_THREAD = (Split(2), Split(0), Split(4)) +_WARPGROUP_LAYOUT = arrangement_pattern(WARPGROUP.layout, per_mode=True) +_WARPGROUP_PATTERN = MeshPattern( + ("thread",), + OrPattern( + ComposedLayoutPattern(offset=WildcardPattern(), outer=_WARPGROUP_LAYOUT), + _WARPGROUP_LAYOUT, + ), +) + def shared(arrangement) -> ShardLayoutPattern: - return ShardLayoutPattern(arrangement, SHARED_BY_ALL, WARPGROUP) + return ShardLayoutPattern(arrangement, SHARED_BY_ALL, _WARPGROUP_PATTERN) def held(arrangement) -> ShardLayoutPattern: - return ShardLayoutPattern(arrangement, HELD_PER_THREAD, WARPGROUP) + return ShardLayoutPattern(arrangement, HELD_PER_THREAD, _WARPGROUP_PATTERN) N = DimVar("n", 8, 257) diff --git a/tests/ir/core/test_overload.py b/tests/ir/core/test_overload.py index 8ee31be8..ea42b190 100644 --- a/tests/ir/core/test_overload.py +++ b/tests/ir/core/test_overload.py @@ -9,7 +9,7 @@ from tilefoundry.ir.core.op_schema import OpSchema from tilefoundry.ir.core.overload import OverloadError, filter_candidates, resolve from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.pattern import Scalar, Tensor, TensorPattern +from tilefoundry.ir.pattern import Scalar, Tensor, TensorPattern, WildcardPattern from tilefoundry.ir.types import TensorType _S = TensorType.umat_scalar() @@ -42,7 +42,11 @@ def __init__(self, **kw: Any) -> None: ... def test_resolve_picks_first_matching_candidate() -> None: """Arity + pattern filter; first-match wins; raises when no match.""" - rank2 = _schema("matmul", TensorPattern(rank=2), TensorPattern(rank=2)) + rank2 = _schema( + "matmul", + TensorPattern(shape=(WildcardPattern(),) * 2), + TensorPattern(shape=(WildcardPattern(),) * 2), + ) any_t = _schema("matmul", Tensor, Tensor) assert resolve([rank2, any_t], [_T2, _T2]) is rank2 diff --git a/tests/ir/core/test_pattern.py b/tests/ir/core/test_pattern.py index 0366f655..e698582e 100644 --- a/tests/ir/core/test_pattern.py +++ b/tests/ir/core/test_pattern.py @@ -8,6 +8,7 @@ Scalar, Tensor, TensorPattern, + WildcardPattern, ) from tilefoundry.ir.types import DType, TensorType @@ -24,12 +25,17 @@ def test_pattern_match_contract() -> None: assert not Tensor.match(TensorType.umat_scalar()) assert not Tensor.match(type("FakeTy", (), {"shape": (3, 4)})()) - rank2_bf16 = TensorPattern(rank=2, dtype=DType.bf16) + rank2_bf16 = TensorPattern(shape=(WildcardPattern(),) * 2, dtype=DType.bf16) assert rank2_bf16.match(_tensor((3, 4), DType.bf16)) assert not rank2_bf16.match(_tensor((3,), DType.bf16)) assert not rank2_bf16.match(_tensor((3, 4), DType.f32)) - combined = AndPattern(parts=(TensorPattern(rank=2), TensorPattern(dtype=DType.f16))) + combined = AndPattern( + parts=( + TensorPattern(shape=(WildcardPattern(),) * 2), + TensorPattern(dtype=DType.f16), + ) + ) assert combined.match(_tensor((3, 4), DType.f16)) assert not combined.match(_tensor((3,), DType.f16)) assert AndPattern(parts=()).match(TensorType.umat_scalar()) From fdd974fb35d91483b0d0b547dbb9f450d4c39a98 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 12:47:53 +0800 Subject: [PATCH 09/19] fix(pattern): preserve shard frame topology matching --- src/tilefoundry/ir/pattern/pattern.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index d45727e1..e2b19436 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -771,6 +771,11 @@ class ShardLayoutPattern(Pattern): def match(self, subject, captures=None): if not isinstance(subject, ShardLayout): return None + subject_names = tuple( + getattr(topology, "name", topology) for topology in subject.mesh.topologies + ) + if subject_names != self.mesh.topologies: + return None held = Match(dict(captures or {})) for pattern, value in ( (self.layout, subject.layout), From 7c67756cc3be8e49b9041ecc942d726fdcf61687 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 12:59:56 +0800 Subject: [PATCH 10/19] refactor(layout): use profiled coalescing for vectors --- docs/spec/shard.md | 20 ++++++ src/tilefoundry/ir/pattern/pattern.py | 31 +++++++-- src/tilefoundry/ir/types/layout_algebra.py | 80 +++++++++++++++++++--- 3 files changed, 118 insertions(+), 13 deletions(-) diff --git a/docs/spec/shard.md b/docs/spec/shard.md index 6756bdd5..360080ca 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -875,6 +875,19 @@ def shard_layout_local_shape( ... +def coalesce(layout: Layout | ComposedLayout, trg_profile=...): + """Merge contiguous modes, optionally within profile-selected groups. + + Args: + layout: Layout to simplify under CuTe's mode-zero-fast convention. + trg_profile: Optional nesting whose terminals select groups to merge. + + Returns: + The equivalent layout with contiguous modes merged. + """ + ... + + def is_inverse_projectable(layout: Layout) -> bool: """Return whether a layout admits the supported inverse projection. @@ -962,6 +975,13 @@ def contains(scope: ComposedLayout, t: int) -> bool: only when `require_static=False`; strict mode MUST reject it. - `try_c_order_strides` MUST return `None` unless every shape entry is a non-boolean integer. + - `coalesce(layout)` MUST flatten and merge contiguous modes under CuTe's + mode-zero-fast convention. With `trg_profile`, it MUST apply that rule at + each profile terminal, preserve unmatched trailing modes, and reject a + profile that asks for more modes at any nesting level with that level in + the diagnostic. A row-major consumer MUST reverse modes within each of its + groups before calling this CuTe operation; `coalesce` itself does not + reinterpret storage order. - Mesh-scope projection MUST accept only an identity inner mapping and an inverse-projectable primitive outer layout; other layouts MUST raise `NotProjectable` rather than guess a projection. diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index e2b19436..2db4fb54 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -18,7 +18,7 @@ from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.layout_algebra import ( ASYNC_WIDTHS, - box_runs, + coalesce, is_inverse_projectable, ) from tilefoundry.ir.types.mesh import separate @@ -86,6 +86,20 @@ def describe(self, name: str = UNNAMED_PLACE) -> str: ) +def _reverse_group(group): + if isinstance(group, tuple): + return tuple(_reverse_group(mode) for mode in reversed(group)) + return group + + +def _row_major_groups_for_cute(layout: Layout) -> Layout: + """Reverse modes within each tile axis for CuTe's mode-0-fast algebra.""" + return Layout( + shape=tuple(_reverse_group(group) for group in layout.shape), + strides=tuple(_reverse_group(group) for group in layout.strides), + ) + + def vector_widths(layout, element_bits: int) -> tuple[int, ...]: """Every cp.async width that divides every run in an arrangement.""" widest = ASYNC_WIDTHS[-1] @@ -101,11 +115,20 @@ def vector_widths(layout, element_bits: int) -> tuple[int, ...]: for value in flatten(group) ): return () - runs = box_runs(held, element_bits, None, limit=None) - unit = [run.extent for run in runs if run.step == 1] + grouped = coalesce( + _row_major_groups_for_cute(held), + (0,) * len(held.shape), + ) + runs = tuple( + sorted( + zip(flatten(grouped.shape), flatten(grouped.strides)), + key=lambda run: run[1], + ) + ) + unit = [extent for extent, step in runs if step == 1] if len(unit) != 1: return () - counted = (unit[0], *(run.step for run in runs if run.step != 1)) + counted = (unit[0], *(step for _, step in runs if step != 1)) return tuple( width for width in ASYNC_WIDTHS diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index d7905447..e92b65ba 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -69,16 +69,11 @@ def cosize(layout: Union[Layout, ComposedLayout]) -> int: return apply(layout, size(layout) - 1) + 1 -def coalesce(layout: Union[Layout, ComposedLayout]): - """Flatten + merge contiguous modes, drop shape-1 modes (CuTe ``coalesce``). +_NO_PROFILE = object() - Coalescing renames the domain and leaves the index mapping alone, so a - swizzled composed layout coalesces underneath its swizzle. - """ - if get_swizzle_portion(layout) is not None: - return ComposedLayout( - inner=layout.inner, offset=layout.offset, outer=coalesce(layout.outer) - ) + +def _coalesce_flat(layout: Layout) -> Layout: + """Apply the flat CuTe ``coalesce`` rule to one layout.""" result_shape: list[int] = [1] result_stride: list[int] = [0] for shape, stride in zip(flat_shape(layout), flat_stride(layout)): @@ -95,6 +90,73 @@ def coalesce(layout: Union[Layout, ComposedLayout]): return Layout(shape=tuple(result_shape), strides=tuple(result_stride)) +def _profile_place(path: tuple[int, ...]) -> str: + return "profile" + "".join(f"[{index}]" for index in path) + + +def _coalesce_profile(layout: Layout, profile, path: tuple[int, ...]) -> Layout: + """Apply flat coalescing at the terminals selected by one profile.""" + if not isinstance(profile, tuple): + return _coalesce_flat(layout) + + strides = layout.strides + if strides is None: + strides = compact_col_major(layout.shape) + if not isinstance(strides, tuple) or len(layout.shape) != len(strides): + raise ValueError( + f"coalesce: layout shape has {len(layout.shape)} modes at " + f"{_profile_place(path)}, but its strides do not" + ) + if len(profile) > len(layout.shape): + raise ValueError( + f"coalesce: {_profile_place(path)} has {len(profile)} modes, but the " + f"layout there has {len(layout.shape)}" + ) + + result_shape: list = [] + result_stride: list = [] + for index, (shape, stride) in enumerate(zip(layout.shape, strides)): + if index >= len(profile): + result_shape.append(shape) + result_stride.append(stride) + continue + + nested = isinstance(shape, tuple) + child = Layout( + shape=shape if nested else (shape,), + strides=stride if isinstance(stride, tuple) else (stride,), + ) + child_profile = profile[index] + child = _coalesce_profile(child, child_profile, (*path, index)) + if nested or isinstance(child_profile, tuple): + result_shape.append(child.shape) + result_stride.append(child.strides) + else: + result_shape.append(child.shape[0]) + result_stride.append(child.strides[0]) + return Layout(shape=tuple(result_shape), strides=tuple(result_stride)) + + +def coalesce(layout: Union[Layout, ComposedLayout], trg_profile=_NO_PROFILE): + """CuTe ``coalesce``, optionally applied at ``trg_profile`` terminals. + + Coalescing renames the domain and leaves the index mapping alone, so a + swizzled composed layout coalesces underneath its swizzle. A tuple profile + transforms the corresponding top-level modes and retains any modes beyond + its length, while a non-tuple terminal applies flat coalescing. + """ + if get_swizzle_portion(layout) is not None: + outer = ( + coalesce(layout.outer) + if trg_profile is _NO_PROFILE + else coalesce(layout.outer, trg_profile) + ) + return ComposedLayout(inner=layout.inner, offset=layout.offset, outer=outer) + if trg_profile is _NO_PROFILE: + return _coalesce_flat(layout) + return _coalesce_profile(layout, trg_profile, ()) + + def frame_of(layout: Union[Layout, ComposedLayout]) -> tuple[int, Layout] | None: """Read a bare layout or an affine composition as offset + outer layout.""" if isinstance(layout, Layout): From 6fe5e6e5125b14d5c92f5a9571633b5ae21f251f Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 13:08:05 +0800 Subject: [PATCH 11/19] refactor(tir): localize instruction pattern facts --- docs/spec/tir.md | 5 +- src/tilefoundry/ir/pattern/__init__.py | 8 +-- src/tilefoundry/ir/pattern/pattern.py | 31 +++++----- src/tilefoundry/ir/pattern/utils.py | 16 ++--- src/tilefoundry/ir/tir/async_copy.py | 18 +++--- .../ir/tir/cuda/memory/copy_async_tensor.py | 59 ++++++++++++++++--- .../ir/tir/cuda/memory/ldmatrix.py | 6 +- src/tilefoundry/ir/tir/memory/copy.py | 8 +-- src/tilefoundry/ir/tir/memory/fill.py | 6 +- src/tilefoundry/ir/types/layout_algebra.py | 50 ---------------- 10 files changed, 100 insertions(+), 107 deletions(-) diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 826ffdc2..38761cf2 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -953,6 +953,9 @@ producer issues copies, groups them, and a consumer waits on the group queue. ##### CopyAsync ```python +ASYNC_WIDTHS = (4, 8, 16) + + class CopyAsync(Op): """Effect form; async gmem→smem copy, non-blocking. @@ -969,7 +972,7 @@ class CopyAsync(Op): - constraints: - Lowers to `tilefoundry::ops::copy_async(src, dst)`. - `src` is gmem and `dst` is smem, with the same dtype. Each layout MUST - admit the same 4-, 8-, or 16-byte vector width and MUST walk the same tile + admit the same width from `ASYNC_WIDTHS` and MUST walk the same tile mode at step 1. For a `ShardLayout`, vector width is read from the whole tile arrangement; its other strides ensure every participant's start is aligned. Two split layouts compare their tile modes only when their mesh diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py index 0db17007..67294a7e 100644 --- a/src/tilefoundry/ir/pattern/__init__.py +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -57,9 +57,9 @@ arrangement_pattern, dtype_place, locate_dim_var, - moved_tile, + operand_tile, storage_place, - vector, + whole_vectors, ) __all__ = [ @@ -108,12 +108,12 @@ "is_symbolic", "locate_dim_var", "matched", - "moved_tile", + "operand_tile", "refusals_between", "relations_of", "resolved", "storage_place", "VectorPattern", - "vector", + "whole_vectors", "vector_widths", ] diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index 2db4fb54..bc5be1b7 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -16,11 +16,7 @@ ) from tilefoundry.ir.types.int_tuple import congruent from tilefoundry.ir.types.layout import flatten -from tilefoundry.ir.types.layout_algebra import ( - ASYNC_WIDTHS, - coalesce, - is_inverse_projectable, -) +from tilefoundry.ir.types.layout_algebra import coalesce, is_inverse_projectable from tilefoundry.ir.types.mesh import separate from .constraint import affine_part @@ -100,9 +96,11 @@ def _row_major_groups_for_cute(layout: Layout) -> Layout: ) -def vector_widths(layout, element_bits: int) -> tuple[int, ...]: - """Every cp.async width that divides every run in an arrangement.""" - widest = ASYNC_WIDTHS[-1] +def vector_widths(layout, element_bits: int, widths: tuple[int, ...]) -> tuple[int, ...]: + """Every requested byte width that divides every run in an arrangement.""" + if not widths: + return () + widest = widths[-1] if isinstance(layout, ShardLayout): layout = layout.layout inner = getattr(layout, "inner", None) @@ -131,7 +129,7 @@ def vector_widths(layout, element_bits: int) -> tuple[int, ...]: counted = (unit[0], *(step for _, step in runs if step != 1)) return tuple( width - for width in ASYNC_WIDTHS + for width in widths if width <= widest and all(value * element_bits % (width * 8) == 0 for value in counted) ) @@ -139,18 +137,19 @@ def vector_widths(layout, element_bits: int) -> tuple[int, ...]: @dataclass(frozen=True) class VectorPattern(Pattern): - """An arrangement that moves whole 4-, 8-, or 16-byte vectors.""" + """An arrangement that moves whole vectors at the requested byte widths.""" width: CapturePattern dtype: str + widths: tuple[int, ...] - def widths(self, subject, captures) -> tuple[int, ...]: + def available_widths(self, subject, captures) -> tuple[int, ...]: bits = getattr(dict(captures or {}).get(self.dtype), "bit_width", None) - return () if type(bits) is not int else vector_widths(subject, bits) + return () if type(bits) is not int else vector_widths(subject, bits, self.widths) def match(self, subject, captures=None): held = dict(captures or {}) - widths = self.widths(subject, held) + widths = self.available_widths(subject, held) if not widths: return None if self.width.name in held: @@ -160,11 +159,13 @@ def match(self, subject, captures=None): def refusal(self, subject, captures=None) -> str | None: if self.match(subject, captures) is not None: return None - sizes = ", ".join(map(str, ASYNC_WIDTHS[:-1])) + f" or {ASYNC_WIDTHS[-1]}" + sizes = " or ".join(map(str, self.widths)) + if len(self.widths) > 2: + sizes = ", ".join(map(str, self.widths[:-1])) + f" or {self.widths[-1]}" return ( f"{subject!r} moves no whole vector of {sizes} bytes -- its run at step 1 " "and every other step are no whole number of one -- so the two ends share " - "no run wide enough for cp.async" + "no run wide enough for the requested vector widths" ) def describe(self, name: str = UNNAMED_PLACE) -> str: diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py index 5aa179ca..216cf9e4 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -4,7 +4,6 @@ from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.types import ComposedLayout, Mesh, StorageKind, Swizzle -from tilefoundry.ir.types.layout_algebra import ASYNC_WIDTHS from .pattern import ( AttrPattern, @@ -27,8 +26,8 @@ WHOLE_BYTES = AttrPattern("bit_width", MultipleOfPattern(8)) -def moved_tile(index: int, storage=None, layout=None) -> TensorPattern: - """A moved tensor tile, with dtype and storage captures named by end.""" +def operand_tile(index: int, storage=None, layout=None) -> TensorPattern: + """A tensor tile whose dtype and storage captures are named by operand slot.""" storages = MOVED_STORAGES if storage is None else storage return TensorPattern( dtype=CapturePattern(dtype_place(index), WHOLE_BYTES), @@ -51,11 +50,12 @@ def storage_place(index: int) -> str: return f"storage{index}" -def vector(index: int) -> VectorPattern: - """End *index* of a cp.async, counted in that end's element dtype.""" +def whole_vectors(index: int, widths: tuple[int, ...]) -> VectorPattern: + """Whole vectors at *widths*, counted in operand *index*'s element dtype.""" return VectorPattern( - CapturePattern("width", OneOfPattern(ASYNC_WIDTHS)), + CapturePattern("width", OneOfPattern(widths)), dtype_place(index), + widths, ) @@ -147,7 +147,7 @@ def _mangle_variant_name(name: str, specializations: tuple[Pattern, ...]) -> str "arrangement_pattern", "dtype_place", "locate_dim_var", - "moved_tile", + "operand_tile", "storage_place", - "vector", + "whole_vectors", ] diff --git a/src/tilefoundry/ir/tir/async_copy.py b/src/tilefoundry/ir/tir/async_copy.py index b956c965..5b3f712e 100644 --- a/src/tilefoundry/ir/tir/async_copy.py +++ b/src/tilefoundry/ir/tir/async_copy.py @@ -5,18 +5,14 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import ( - DistinctConstraint, - SameModesConstraint, - any_threads, - moved_tile, - vector, -) +from tilefoundry.ir.pattern import DistinctConstraint, SameModesConstraint, utils from tilefoundry.ir.tir.verify import verify_between, verify_operands from tilefoundry.ir.types import Layout, UnitType from tilefoundry.ir.types.storage import StorageKind as S from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt +ASYNC_WIDTHS = (4, 8, 16) + @register_op(dialect="T", category="async", name="copy_async") class CopyAsync(Op): @@ -25,12 +21,12 @@ class CopyAsync(Op): src = ParamDef( kind="input", effect=MemoryEffect.READ, - pattern=moved_tile(0, S.GMEM, vector(0)), + pattern=utils.operand_tile(0, S.GMEM, utils.whole_vectors(0, ASYNC_WIDTHS)), ) dst = ParamDef( kind="input", effect=MemoryEffect.WRITE, - pattern=moved_tile(1, S.SMEM, vector(1)), + pattern=utils.operand_tile(1, S.SMEM, utils.whole_vectors(1, ASYNC_WIDTHS)), ) between = ( DistinctConstraint("storage", "src", "dst"), @@ -42,7 +38,7 @@ class CopyAsync(Op): optional=True, default=None, ) - scope = any_threads() + scope = utils.any_threads() @register_typeinfer(CopyAsync) @@ -99,4 +95,4 @@ def _(call: "Call", ctx: "VerifyContext") -> None: ctx.error(call, f"CpAsyncWait.n must be a non-negative int, got {n!r}") -__all__ = ["CopyAsync", "CpAsyncCommit", "CpAsyncWait"] +__all__ = ["ASYNC_WIDTHS", "CopyAsync", "CpAsyncCommit", "CpAsyncWait"] diff --git a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py index 8958d8cb..dfe7e218 100644 --- a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py +++ b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py @@ -2,7 +2,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace from enum import Enum from tilefoundry.ir.core import Op @@ -26,17 +26,15 @@ SwitchPattern, SwizzlePattern, affine_part, - dtype_place, matched, - moved_tile, relations_of, - storage_place, + utils, ) from tilefoundry.ir.pattern.match import written_place, written_tuple from tilefoundry.ir.tir.verify import verify_between, verify_operands from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, UnitType from tilefoundry.ir.types.int_tuple import flatten -from tilefoundry.ir.types.layout_algebra import box_runs, frame_of +from tilefoundry.ir.types.layout_algebra import frame_of from tilefoundry.ir.types.storage import StorageKind as S from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -54,6 +52,47 @@ TMA_STORAGES = (S.GMEM, S.SMEM) +@dataclass(frozen=True) +class Run: + """One contiguous run of modes from one logical tile axis.""" + + extent: int + step: int + axis: int + mode: int + + +def box_runs( + layout: Layout, + element_bits: int, + span: int | None, + limit: int | None = BOX_EXTENT, +) -> tuple[Run, ...]: + """Read TMA box runs by tile axis, ordered by increasing step.""" + runs: list[Run] = [] + for axis, (extents, steps) in enumerate(zip(layout.shape, layout.strides)): + modes = tuple(enumerate(zip(flatten(extents), flatten(steps)))) + for mode, (extent, step) in reversed(modes): + if extent == 1: + continue + last = runs[-1] if runs and runs[-1].axis == axis else None + joined = None if last is None else last.extent * extent + if ( + last is not None + and step == last.step * last.extent + and (limit is None or joined <= limit) + and not ( + span is not None + and last.step == 1 + and joined * element_bits > span * 8 + ) + ): + runs[-1] = replace(last, extent=joined) + else: + runs.append(Run(extent, step, axis, mode)) + return tuple(sorted(runs, key=lambda run: run.step)) + + class TmaSwizzle(Enum): """The shared-memory swizzle selected when the tensor map is encoded.""" @@ -324,19 +363,19 @@ class CopyAsyncTensor(Op): src = ParamDef( kind="input", effect=MemoryEffect.READ, - pattern=moved_tile( + pattern=utils.operand_tile( 0, TMA_STORAGES, - TmaOperandPattern(storage_place(0), dtype_place(0)), + TmaOperandPattern(utils.storage_place(0), utils.dtype_place(0)), ), ) dst = ParamDef( kind="input", effect=MemoryEffect.WRITE, - pattern=moved_tile( + pattern=utils.operand_tile( 1, TMA_STORAGES, - TmaOperandPattern(storage_place(1), dtype_place(1)), + TmaOperandPattern(utils.storage_place(1), utils.dtype_place(1)), ), ) between = ( @@ -387,6 +426,7 @@ def verify_copy_async_tensor(call: "Call", ctx: "VerifyContext") -> None: "BoxFamily", "BoxPattern", "CopyAsyncTensor", + "Run", "SWIZZLE_PLACE", "TENSORMAP_READING", "TMA_RANK", @@ -397,5 +437,6 @@ def verify_copy_async_tensor(call: "Call", ctx: "VerifyContext") -> None: "TmaGlobalPattern", "TmaOperandPattern", "TmaSwizzle", + "box_runs", "unframed", ] diff --git a/src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py b/src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py index f7f1609e..057007ad 100644 --- a/src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py +++ b/src/tilefoundry/ir/tir/cuda/memory/ldmatrix.py @@ -5,7 +5,7 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import moved_tile +from tilefoundry.ir.pattern import utils from tilefoundry.ir.tir.cuda.nn.sm80_mma import Mma from tilefoundry.ir.tir.memory.copy import Copy from tilefoundry.ir.types import Mesh @@ -20,7 +20,9 @@ class LdMatrix(Op): capability = "tensor_core" - src = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=moved_tile(0, S.SMEM)) + src = ParamDef( + kind="input", effect=MemoryEffect.READ, pattern=utils.operand_tile(0, S.SMEM) + ) dst = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=Mma.A) scope = ParamDef( kind="attribute", diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index 5a50f7af..568f14b9 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -11,7 +11,7 @@ from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import any_threads, moved_tile +from tilefoundry.ir.pattern import utils from tilefoundry.ir.types import LayoutBase, UnitType from tilefoundry.ir.types.shard_layout import ShardLayout from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -21,12 +21,12 @@ class Copy(Op): """Copies ``src`` into ``dst`` (in-place memory write).""" - src = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=moved_tile(0)) - dst = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=moved_tile(1)) + src = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=utils.operand_tile(0)) + dst = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=utils.operand_tile(1)) rmem_layout = ParamDef(kind="attribute", annotation=LayoutBase, optional=True, default=None) smem_layout = ParamDef(kind="attribute", annotation=LayoutBase, optional=True, default=None) - scope = any_threads() + scope = utils.any_threads() @register_typeinfer(Copy) diff --git a/src/tilefoundry/ir/tir/memory/fill.py b/src/tilefoundry/ir/tir/memory/fill.py index 3f3cb61a..348e292a 100644 --- a/src/tilefoundry/ir/tir/memory/fill.py +++ b/src/tilefoundry/ir/tir/memory/fill.py @@ -9,7 +9,7 @@ from tilefoundry.ir.core import Constant, Op from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op -from tilefoundry.ir.pattern import Scalar, any_threads, moved_tile +from tilefoundry.ir.pattern import Scalar, utils from tilefoundry.ir.types import UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -18,9 +18,9 @@ class Fill(Op): """Fills ``tensor`` element-wise with ``value`` (rank-0 scalar).""" - tensor = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=moved_tile(0)) + tensor = ParamDef(kind="input", effect=MemoryEffect.WRITE, pattern=utils.operand_tile(0)) value = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=Scalar) - scope = any_threads() + scope = utils.any_threads() @register_typeinfer(Fill) diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index e92b65ba..8bc58031 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -10,11 +10,8 @@ from __future__ import annotations -from dataclasses import dataclass, replace from typing import Optional, Union -from tilefoundry.ir.types.layout import flatten - from . import swizzle_layout from .layout import ComposedLayout, Layout, Swizzle, flat_shape, flat_stride, size from .stride import compact_col_major, idx2crd @@ -25,19 +22,6 @@ class NotProjectable(ValueError): """A layout cannot serve as a mesh execution scope (not inverse-projectable).""" -ASYNC_WIDTHS = (4, 8, 16) - - -@dataclass(frozen=True) -class Run: - """One contiguous run of modes from one logical tile axis.""" - - extent: int - step: int - axis: int - mode: int - - def apply(layout: Union[Layout, ComposedLayout], coord: int) -> int: """``crd2idx`` of a 1-D domain coord: decompose by shape, dot with strides. @@ -166,37 +150,6 @@ def frame_of(layout: Union[Layout, ComposedLayout]) -> tuple[int, Layout] | None return layout.offset, layout.outer -def box_runs( - layout: Layout, - element_bits: int, - span: int | None, - limit: int | None = 256, -) -> tuple[Run, ...]: - """Read contiguous runs by tile axis, ordered by increasing step.""" - runs: list[Run] = [] - for axis, (extents, steps) in enumerate(zip(layout.shape, layout.strides)): - modes = tuple(enumerate(zip(flatten(extents), flatten(steps)))) - for mode, (extent, step) in reversed(modes): - if extent == 1: - continue - last = runs[-1] if runs and runs[-1].axis == axis else None - joined = None if last is None else last.extent * extent - if ( - last is not None - and step == last.step * last.extent - and (limit is None or joined <= limit) - and not ( - span is not None - and last.step == 1 - and joined * element_bits > span * 8 - ) - ): - runs[-1] = replace(last, extent=joined) - else: - runs.append(Run(extent, step, axis, mode)) - return tuple(sorted(runs, key=lambda run: run.step)) - - def complement(layout: Layout, max_idx: int = 1) -> Layout: """CuTe ``complement``: the modes that fill the gaps below ``max_idx``.""" result_shape: list[int] = [] @@ -398,15 +351,12 @@ def contains(scope: ComposedLayout, t: int) -> bool: __all__ = [ - "ASYNC_WIDTHS", "NotProjectable", - "Run", "composition", "cosize", "apply", "coalesce", "frame_of", - "box_runs", "complement", "is_inverse_projectable", "right_inverse", From 4796caf6f2ebfe2ea5a931b411bff1f5c66ca818 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 13:30:29 +0800 Subject: [PATCH 12/19] refactor(layout): align helpers with owning modules --- docs/spec/code-organization.md | 9 ++ docs/spec/shard.md | 77 ++++-------- src/tilefoundry/ir/mesh_scope.py | 8 +- src/tilefoundry/ir/pattern/constraint.py | 9 +- .../ir/tir/cuda/memory/copy_async_tensor.py | 47 +++---- src/tilefoundry/ir/tir/cuda/nn/mma_atom.py | 17 +-- src/tilefoundry/ir/tir/sync.py | 2 +- src/tilefoundry/ir/types/layout.py | 38 ++++++ src/tilefoundry/ir/types/layout_algebra.py | 115 +++--------------- src/tilefoundry/ir/types/mesh.py | 34 ++++-- tests/ir/types/test_mma_fragment_layouts.py | 5 +- 11 files changed, 155 insertions(+), 206 deletions(-) diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 09abf73c..bdb7a975 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -102,6 +102,15 @@ physical directory layout reflects that boundary directly. contracts are distinct even though both are consumed across the codegen boundary. +`ir/types/layout_algebra.py` has a mechanically checkable public surface: +every name in its `__all__` MUST have the same name in CuTe, with +`is_inverse_projectable` as its sole exception. Python's missing overload +dispatch requires the separately named `supports_composition`, +`supports_inverse`, and `inverse` adapters in `swizzle_layout.py`; +`NotProjectable` in `layout.py` is the named Python diagnostic needed by +composed application and inversion. These are explicit exceptions, not a +license for unrelated helpers in the algebra modules. + `ir/pattern/`, `ir/clause/`, `visitor_registry/`, and `dump/` are cross-cutting packages; their stable responsibilities are owned by [core-ir](./core-ir.md), [parser](./parser.md), [visitor-registry](./visitor-registry.md), and [inspection](./inspection.md), diff --git a/docs/spec/shard.md b/docs/spec/shard.md index 360080ca..030d904e 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -274,7 +274,7 @@ the ones they are XORed onto carries the same `Swizzle` through: a window's start is a constant shift of the index, which `offset` already states. - A swizzled `ComposedLayout` MUST NOT serve as a `Mesh` execution scope: - [§9](#9-layout-construction-and-mesh-scope-projection) admits an identity + [§9](#9-layout-construction-and-algebra) admits an identity `inner` only. --- @@ -307,6 +307,16 @@ class Mesh: names: tuple[str, ...] = () def __getitem__(self, key) -> Mesh: ... + + +def levels(mesh: Mesh) -> tuple[Layout, ...]: + """Return each topology level's arrangement in its own numbering.""" + ... + + +def starts(mesh: Mesh) -> tuple[int, ...]: + """Decode each topology level's start from device numbering.""" + ... ``` - constraints: @@ -406,10 +416,12 @@ value is consumed by a region. guessed level. - A reader that asks for an exact participant set at a selected topology level MUST take that level's projected layout as above, and MUST refuse a Mesh - that names no such level. For a plain `Layout`, the set is - `{apply(layout, c) | 0 <= c < size(layout)}`. For a sliced `ComposedLayout`, - it is `{image(layout, c) | 0 <= c < size(layout.outer)}`. - - That image MUST be static, positive, inverse-projectable, duplicate-free, + that names no such level. For both a plain `Layout` and a sliced + `ComposedLayout`, the set is + `{apply(layout, c) | 0 <= c < size(layout)}`; composition applies its + retained offset and inner mapping. + - That participant set MUST be static, positive, inverse-projectable, + duplicate-free, and contained in `[0, selected_topology.size)`. A plain Mesh MUST cover the complete selected domain. A strict subdomain MUST use a sliced Mesh so its offset is retained rather than collapsed to an extent. @@ -796,7 +808,7 @@ copy that displacement to a materialized consumer. --- -## 9. Layout construction and mesh-scope projection +## 9. Layout construction and algebra Making the steps of a compact arrangement and reading an index back into the coordinate that reaches it are not operations on layouts, so they are filed @@ -807,6 +819,11 @@ class NotProjectable(ValueError): """Report that a layout cannot serve as a mesh execution scope.""" +def apply(layout: Layout | ComposedLayout, coord: int) -> int: + """Apply a layout to one linear domain coordinate.""" + ... + + def prefix_product(shape: tuple[int, ...]) -> tuple[int, ...]: """Return exclusive prefix-product strides. @@ -889,7 +906,7 @@ def coalesce(layout: Layout | ComposedLayout, trg_profile=...): def is_inverse_projectable(layout: Layout) -> bool: - """Return whether a layout admits the supported inverse projection. + """Return whether a layout admits the supported inverse construction. Args: layout: Primitive layout to inspect. @@ -924,43 +941,6 @@ def right_inverse(layout: Layout | ComposedLayout): ... -def image(scope: ComposedLayout, coord: int) -> int: - """Map a domain coordinate into a mesh execution scope. - - Args: - scope: Admissible composed mesh scope. - coord: Flat domain coordinate. - - Returns: - Runtime thread index. - """ - ... - - -def project(scope: ComposedLayout, t: int) -> tuple[int, ...] | None: - """Project a runtime thread into a mesh execution scope. - - Args: - scope: Admissible composed mesh scope. - t: Runtime thread index. - - Returns: - The multidimensional coordinate, or None when outside the scope. - """ - ... - - -def contains(scope: ComposedLayout, t: int) -> bool: - """Return whether a runtime thread participates in a mesh scope. - - Args: - scope: Admissible composed mesh scope. - t: Runtime thread index. - - Returns: - Whether the thread participates. - """ - ... ``` - constraints: @@ -982,9 +962,6 @@ def contains(scope: ComposedLayout, t: int) -> bool: the diagnostic. A row-major consumer MUST reverse modes within each of its groups before calling this CuTe operation; `coalesce` itself does not reinterpret storage order. - - Mesh-scope projection MUST accept only an identity inner mapping and an - inverse-projectable primitive outer layout; other layouts MUST raise - `NotProjectable` rather than guess a projection. - - `project` MUST return `None` for an index outside the scope or outside the - outer layout's round-tripping image; `contains` MUST report the same test as - a boolean. + - Inverting a composed mesh layout MUST accept only an identity inner mapping + and an inverse-projectable primitive outer layout; other layouts MUST raise + `NotProjectable` rather than guess an inverse. diff --git a/src/tilefoundry/ir/mesh_scope.py b/src/tilefoundry/ir/mesh_scope.py index f674867f..f7aabbd0 100644 --- a/src/tilefoundry/ir/mesh_scope.py +++ b/src/tilefoundry/ir/mesh_scope.py @@ -12,7 +12,7 @@ from tilefoundry.ir.types.int_tuple import flatten, product from tilefoundry.ir.types.layout import ComposedLayout, Layout, size from tilefoundry.ir.types.layout_algebra import is_inverse_projectable -from tilefoundry.ir.types.mesh import Mesh, _levels, _starts, check_topology +from tilefoundry.ir.types.mesh import Mesh, check_topology, levels, starts from tilefoundry.ir.types.storage import StorageKind, resolve_storage from tilefoundry.ir.types.stride import compact_major @@ -31,7 +31,7 @@ def device_layout(mesh: Mesh) -> Layout: units = compact_major(sizes, major="row") if all( isinstance(one, int) for one in sizes ) else (1,) * len(sizes) - for arrangement, unit in zip(_levels(mesh), units): + for arrangement, unit in zip(levels(mesh), units): stated = arrangement.strides shape.extend(flatten(arrangement.shape)) strides.extend( @@ -85,7 +85,7 @@ def covered_by_scope(mesh: Mesh, current: Mesh) -> bool: scope = { getattr(topology, "name", topology): _selected(arrangement, start) for topology, arrangement, start in zip( - current.topologies, _levels(current), _starts(current) + current.topologies, levels(current), starts(current) ) } return all( @@ -93,7 +93,7 @@ def covered_by_scope(mesh: Mesh, current: Mesh) -> bool: and _selected(arrangement, start) == scope[getattr(topology, "name", topology)] for topology, arrangement, start in zip( - mesh.topologies, _levels(mesh), _starts(mesh) + mesh.topologies, levels(mesh), starts(mesh) ) ) diff --git a/src/tilefoundry/ir/pattern/constraint.py b/src/tilefoundry/ir/pattern/constraint.py index 732cba5f..397d25c0 100644 --- a/src/tilefoundry/ir/pattern/constraint.py +++ b/src/tilefoundry/ir/pattern/constraint.py @@ -64,13 +64,18 @@ def written(self) -> str: return f"{self.left}.{self.field} = {self.right}.{self.field}" -def affine_part(layout): - """Return the strided affine part beneath shard frames and swizzles.""" +def affine_part(layout, *, plain: bool = False): + """Return the strided affine part beneath shard frames and swizzles. + + When *plain* is true, refuse a composition with a transform or offset. + """ if isinstance(layout, ShardLayout): if not all(isinstance(attr, Broadcast) for attr in layout.attrs): return None layout = layout.layout if isinstance(layout, ComposedLayout): + if plain and (layout.inner is not None or layout.offset != 0): + return None if layout.inner is not None and not isinstance(layout.inner, Swizzle): return None layout = layout.outer diff --git a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py index dfe7e218..446df7f1 100644 --- a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py +++ b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py @@ -34,7 +34,6 @@ from tilefoundry.ir.tir.verify import verify_between, verify_operands from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, UnitType from tilefoundry.ir.types.int_tuple import flatten -from tilefoundry.ir.types.layout_algebra import frame_of from tilefoundry.ir.types.storage import StorageKind as S from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -106,27 +105,6 @@ def swizzle(self) -> Swizzle | None: return None if self is TmaSwizzle.NONE else Swizzle(self.value.bit_length() - 5, 4, 3) -def unframed(layout): - """Drop a shard frame only when every participant sees the whole tile.""" - if isinstance(layout, ShardLayout) and affine_part(layout) is not None: - return layout.layout - return layout - - -def _affine(subject): - framed = frame_of(unframed(subject)) - layout = framed[1] if framed is not None and framed[0] == 0 else None - if not isinstance(layout, Layout) or layout.strides is None: - return None - if any( - type(value) is not int - for group in (layout.shape, layout.strides) - for value in flatten(group) - ): - return None - return layout - - def _missed_place(places, values, captures, first: int = 0) -> tuple | None: held = captures for index, (place, value) in enumerate(zip(places, values), first): @@ -146,8 +124,12 @@ class BoxPattern(Pattern): span: int | None = None def reading(self, subject, captures) -> tuple[tuple | None, str | None]: - layout = _affine(subject) - if layout is None: + layout = affine_part(subject, plain=True) + if layout is None or any( + type(value) is not int + for group in (layout.shape, layout.strides) + for value in flatten(group) + ): return None, f"{subject!r} is no static strided arrangement" width = getattr(captures.get(self.dtype), "bit_width", None) if type(width) is not int: @@ -216,10 +198,14 @@ class BoxFamily(SwitchPattern): """Every unswizzled or swizzled shared-memory box.""" def match(self, subject, captures=None): - return super().match(unframed(subject), captures) + if isinstance(subject, ShardLayout) and affine_part(subject) is not None: + subject = subject.layout + return super().match(subject, captures) def refusal(self, subject, captures=None) -> str | None: - layout = unframed(subject) + layout = subject + if isinstance(layout, ShardLayout) and affine_part(layout) is not None: + layout = layout.layout transform = layout.inner if isinstance(layout, ComposedLayout) else None if transform is not None and layout.offset != 0: return f"it is reached through {transform!r} at offset {layout.offset}, not 0" @@ -244,8 +230,12 @@ class TensorMapPattern(Pattern): dtype: str def reading(self, subject) -> tuple[tuple | None, str | None]: - layout = _affine(subject) - if layout is None: + layout = affine_part(subject, plain=True) + if layout is None or any( + type(value) is not int + for group in (layout.shape, layout.strides) + for value in flatten(group) + ): return None, f"{subject!r} is no static strided tensor a tensormap describes" modes = [ (extent, step) @@ -438,5 +428,4 @@ def verify_copy_async_tensor(call: "Call", ctx: "VerifyContext") -> None: "TmaOperandPattern", "TmaSwizzle", "box_runs", - "unframed", ] diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py index 35e0a3cb..e2132bf5 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py @@ -22,9 +22,10 @@ resolved, ) from tilefoundry.ir.pattern.match import written_binding, written_bindings, written_place -from tilefoundry.ir.types import Layout, Mesh +from tilefoundry.ir.types import ComposedLayout, Layout, Mesh from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.layout_algebra import coalesce, frame_of +from tilefoundry.ir.types.layout_algebra import coalesce +from tilefoundry.ir.types.mesh import levels, starts class MmaAtom: @@ -261,13 +262,15 @@ def physical_frames_match(left: Mesh, right: Mesh) -> bool: return False def frame(mesh): - framed = frame_of(mesh.layout) - if framed is None: + if isinstance(mesh.layout, ComposedLayout) and mesh.layout.inner is not None: return None - offset, layout = framed - if layout.strides is None: + try: + arranged = levels(mesh) + except (IndexError, TypeError, ValueError): return None - return offset, coalesce(_reverse(layout)) + if any(layout.strides is None for layout in arranged): + return None + return starts(mesh), tuple(coalesce(_reverse(layout)) for layout in arranged) a, b = frame(left), frame(right) return a is not None and b is not None and a == b diff --git a/src/tilefoundry/ir/tir/sync.py b/src/tilefoundry/ir/tir/sync.py index 6337f01c..90d4e903 100644 --- a/src/tilefoundry/ir/tir/sync.py +++ b/src/tilefoundry/ir/tir/sync.py @@ -12,7 +12,7 @@ from tilefoundry.ir.types import UnitType from tilefoundry.ir.types.int_tuple import flatten, product from tilefoundry.ir.types.layout import ComposedLayout, Layout, get -from tilefoundry.ir.types.layout_algebra import apply as _apply +from tilefoundry.ir.types.layout import apply as _apply from tilefoundry.ir.types.layout_algebra import size as _size from tilefoundry.ir.types.mesh import Mesh from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt diff --git a/src/tilefoundry/ir/types/layout.py b/src/tilefoundry/ir/types/layout.py index 96157776..55501fa5 100644 --- a/src/tilefoundry/ir/types/layout.py +++ b/src/tilefoundry/ir/types/layout.py @@ -16,6 +16,10 @@ def domain_rank(self) -> int: return len(flatten(self.shape)) +class NotProjectable(ValueError): + """A layout cannot serve as a mesh execution scope (not inverse-projectable).""" + + @dataclass(frozen=True) class Layout(LayoutBase): """Cute-style layout: shape + per-axis cute strides.""" @@ -147,6 +151,38 @@ def flat_stride(layout: Layout) -> tuple[int, ...]: return compact_col_major(flat_shape(layout)) +def apply(layout: Layout | ComposedLayout, coord: int) -> int: + """``crd2idx`` of a 1-D domain coord: decompose by shape, dot with strides. + + A ``ComposedLayout`` applies its components in order, so this is + ``inner(offset + outer(coord))`` — the swizzle case included, since a + ``Swizzle`` is exactly a mapping on that index. + """ + if isinstance(layout, ComposedLayout): + return _apply_any(layout, coord) + shape = flat_shape(layout) + stride = flat_stride(layout) + idx = 0 + rem = coord + for extent, step in zip(shape, stride): + idx += (rem % extent) * step + rem //= extent + return idx + + +def _apply_any(layout, x: int) -> int: + """Apply a ``Layout`` / ``ComposedLayout`` (``None`` ≡ identity) to ``x``.""" + if layout is None: + return x + if isinstance(layout, Swizzle): + return layout(x) + if isinstance(layout, Layout): + return apply(layout, x) + if isinstance(layout, ComposedLayout): + return _apply_any(layout.inner, layout.offset + _apply_any(layout.outer, x)) + raise NotProjectable(f"cannot apply layout of type {type(layout).__name__}") + + def size(layout: Layout) -> int: return product(layout.shape) @@ -206,6 +242,8 @@ def take(layout: LayoutBase, begin: int, end: int) -> "Layout": __all__ = [ "LayoutBase", + "NotProjectable", + "apply", "flatten", "size", "unflatten", diff --git a/src/tilefoundry/ir/types/layout_algebra.py b/src/tilefoundry/ir/types/layout_algebra.py index 8bc58031..502c9f4e 100644 --- a/src/tilefoundry/ir/types/layout_algebra.py +++ b/src/tilefoundry/ir/types/layout_algebra.py @@ -1,46 +1,31 @@ -"""Provide flat CuTe layout algebra for mesh execution scopes. +"""Provide the supported CuTe algebra for ``Layout`` and ``ComposedLayout``. -The restricted port supports coordinate application, inverses, containment, -and projection for ``Layout`` and ``ComposedLayout``, and the CuTe -``swizzle_layout.hpp`` specializations for a ``Swizzle`` in ``inner``. -Execution scopes must be injective and inverse-projectable. +The restricted port covers coalescing, composition, complements, and inverses, +including the CuTe ``swizzle_layout.hpp`` specializations for a ``Swizzle`` in +``inner``. Layout application lives with the layout types themselves. -See [shard §9](docs/spec/shard.md#9-layout-construction-and-mesh-scope-projection). +See [shard §9](docs/spec/shard.md#9-layout-construction-and-algebra). """ from __future__ import annotations -from typing import Optional, Union +from typing import Union from . import swizzle_layout -from .layout import ComposedLayout, Layout, Swizzle, flat_shape, flat_stride, size -from .stride import compact_col_major, idx2crd +from .layout import ( + ComposedLayout, + Layout, + NotProjectable, + Swizzle, + apply, + flat_shape, + flat_stride, + size, +) +from .stride import compact_col_major from .swizzle_layout import get_swizzle_portion -class NotProjectable(ValueError): - """A layout cannot serve as a mesh execution scope (not inverse-projectable).""" - - -def apply(layout: Union[Layout, ComposedLayout], coord: int) -> int: - """``crd2idx`` of a 1-D domain coord: decompose by shape, dot with strides. - - A ``ComposedLayout`` applies its components in order, so this is - ``inner(offset + outer(coord))`` — the swizzle case included, since a - ``Swizzle`` is exactly a mapping on that index. - """ - if isinstance(layout, ComposedLayout): - return _apply_any(layout, coord) - shape = flat_shape(layout) - stride = flat_stride(layout) - idx = 0 - rem = coord - for s, d in zip(shape, stride): - idx += (rem % s) * d - rem //= s - return idx - - def cosize(layout: Union[Layout, ComposedLayout]) -> int: """The codomain extent. @@ -141,15 +126,6 @@ def coalesce(layout: Union[Layout, ComposedLayout], trg_profile=_NO_PROFILE): return _coalesce_profile(layout, trg_profile, ()) -def frame_of(layout: Union[Layout, ComposedLayout]) -> tuple[int, Layout] | None: - """Read a bare layout or an affine composition as offset + outer layout.""" - if isinstance(layout, Layout): - return 0, layout - if layout.inner is not None or not isinstance(layout.outer, Layout): - return None - return layout.offset, layout.outer - - def complement(layout: Layout, max_idx: int = 1) -> Layout: """CuTe ``complement``: the modes that fill the gaps below ``max_idx``.""" result_shape: list[int] = [] @@ -299,69 +275,12 @@ def right_inverse(layout: Union[Layout, ComposedLayout, Swizzle]): return _right_inverse_layout(layout) -def _apply_any(layout, x: int) -> int: - """Apply a ``Layout`` / ``ComposedLayout`` (``None`` ≡ identity) to ``x``.""" - if layout is None: - return x - if isinstance(layout, Swizzle): - return layout(x) - if isinstance(layout, Layout): - return apply(layout, x) - if isinstance(layout, ComposedLayout): - return _apply_any(layout.inner, layout.offset + _apply_any(layout.outer, x)) - raise NotProjectable(f"cannot apply layout of type {type(layout).__name__}") - - -def image(scope: ComposedLayout, coord: int) -> int: - """``inner(offset + outer(coord))`` for a 1-D domain coord.""" - _check_admissible(scope) - return scope.offset + apply(scope.outer, coord) - - -def project(scope: ComposedLayout, t: int) -> Optional[tuple[int, ...]]: - """Project. - - Recover the multi-dim domain coord ``(warp, lane, …)`` of thread ``t``, - or ``None`` if ``t`` is not in this scope. Raises ``NotProjectable`` if the - scope itself is inadmissible (non-identity inner / non-injective outer). - - Built on the composed ``left_inverse``: ``coord_1d = left_inverse(scope)(t)`` - (``= left_inverse(outer)(t − offset)``); the multi-dim coord is ``idx2crd`` - of that over ``outer``'s shape. Returns ``None`` unless the coord is - in-domain *and* round-trips (``image(coord) == t``). - """ - _check_admissible(scope) - outer = scope.outer - if t - scope.offset < 0: - return None - coord_1d = _apply_any(left_inverse(scope), t) - if not (0 <= coord_1d < size(outer)): - return None - - if image(scope, coord_1d) != t: - return None - - shape = flat_shape(outer) - return idx2crd(coord_1d, shape, compact_col_major(shape)) - - -def contains(scope: ComposedLayout, t: int) -> bool: - """Does thread ``t`` execute this mesh scope's body.""" - return project(scope, t) is not None - - __all__ = [ - "NotProjectable", "composition", "cosize", - "apply", "coalesce", - "frame_of", "complement", "is_inverse_projectable", "right_inverse", "left_inverse", - "image", - "project", - "contains", ] diff --git a/src/tilefoundry/ir/types/mesh.py b/src/tilefoundry/ir/types/mesh.py index 1e38cb8b..24b50064 100644 --- a/src/tilefoundry/ir/types/mesh.py +++ b/src/tilefoundry/ir/types/mesh.py @@ -73,8 +73,8 @@ def __getitem__(self, key) -> "Mesh": """ if isinstance(self.layout, ComposedLayout): raise ValueError("cannot slice an already-sliced mesh (nested slice unsupported)") - levels = _levels(self) - rank = sum(len(flatten(level.shape)) for level in levels) + held_levels = levels(self) + rank = sum(len(flatten(level.shape)) for level in held_levels) keys = key if isinstance(key, tuple) else (key,) if len(keys) > rank: raise ValueError(f"mesh slice has {len(keys)} indices but the mesh has {rank} axes") @@ -88,7 +88,7 @@ def __getitem__(self, key) -> "Mesh": if len(self.topologies) == 1 else compact_major(tuple(topology.size for topology in self.topologies)) ) - for level, unit in zip(levels, units): + for level, unit in zip(held_levels, units): level_shape = tuple(flatten(level.shape)) stated = level.strides level_strides = ( @@ -231,7 +231,7 @@ def _nested(layout, topologies: tuple) -> "Layout | ComposedLayout": return Layout(shape=tuple(shape), strides=tuple(strides)) -def _levels(mesh: Mesh) -> tuple[Layout, ...]: +def levels(mesh: Mesh) -> tuple[Layout, ...]: """Each level's arrangement, in that level's own numbering.""" stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout if stated is None: @@ -242,7 +242,7 @@ def _levels(mesh: Mesh) -> tuple[Layout, ...]: return tuple(get(stated, index) for index in range(_rank(stated))) -def _starts(mesh: Mesh) -> tuple[int, ...]: +def starts(mesh: Mesh) -> tuple[int, ...]: """Where each level's run starts, decoded from the device-numbered offset.""" offset = mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0 if not isinstance(offset, int): @@ -263,7 +263,7 @@ def check_topology(mesh: Mesh) -> None: """ if isinstance(mesh.layout, ComposedLayout): return - for topology, arrangement in zip(mesh.topologies, _levels(mesh)): + for topology, arrangement in zip(mesh.topologies, levels(mesh)): declared = getattr(topology, "size", None) if not isinstance(declared, int) or isinstance(declared, bool): continue @@ -315,8 +315,8 @@ def make_mesh(*meshes: Mesh) -> Mesh: if set(here).isdisjoint(there): result = _joined( (*result.topologies, *inner.topologies), - (*_levels(result), *_levels(inner)), - (*_starts(result), *_starts(inner)), + (*levels(result), *levels(inner)), + (*starts(result), *starts(inner)), (*result.names, *inner.names), sliced=isinstance(result.layout, ComposedLayout) or isinstance(inner.layout, ComposedLayout), @@ -325,12 +325,12 @@ def make_mesh(*meshes: Mesh) -> Mesh: result = inner elif len(there) < len(here) and here[-len(there) :] == there: kept = len(here) - len(there) - above = _levels(result)[:kept] + above = levels(result)[:kept] named = sum(len(flatten(level.shape)) for level in above) result = _joined( (*result.topologies[:kept], *inner.topologies), - (*above, *_levels(inner)), - (*_starts(result)[:kept], *_starts(inner)), + (*above, *levels(inner)), + (*starts(result)[:kept], *starts(inner)), (*result.names[:named], *inner.names), sliced=isinstance(result.layout, ComposedLayout) or isinstance(inner.layout, ComposedLayout), @@ -351,7 +351,7 @@ def separate(mesh: Mesh) -> tuple[Mesh, ...]: sliced = isinstance(mesh.layout, ComposedLayout) names_at = 0 separated: list[Mesh] = [] - for topology, level, start in zip(mesh.topologies, _levels(mesh), _starts(mesh)): + for topology, level, start in zip(mesh.topologies, levels(mesh), starts(mesh)): axis_count = len(flatten(level.shape)) names = mesh.names[names_at : names_at + axis_count] layout: Layout | ComposedLayout = level @@ -362,4 +362,12 @@ def separate(mesh: Mesh) -> tuple[Mesh, ...]: return tuple(separated) -__all__ = ["Mesh", "Topology", "check_topology", "make_mesh", "separate"] +__all__ = [ + "Mesh", + "Topology", + "check_topology", + "levels", + "make_mesh", + "separate", + "starts", +] diff --git a/tests/ir/types/test_mma_fragment_layouts.py b/tests/ir/types/test_mma_fragment_layouts.py index 1b107169..fd63aa08 100644 --- a/tests/ir/types/test_mma_fragment_layouts.py +++ b/tests/ir/types/test_mma_fragment_layouts.py @@ -12,6 +12,7 @@ from tilefoundry.dsl import T from tilefoundry.ir.core import Call, Var from tilefoundry.ir.hir.sharding.reshard import Reshard +from tilefoundry.ir.tir.cuda.nn.sm80_mma import WARP from tilefoundry.ir.types import DType, ShardLayout, Split, TensorType from tilefoundry.ir.types.int_tuple import flatten, product from tilefoundry.ir.types.storage import StorageKind @@ -23,9 +24,9 @@ def _realized_fragment(role: str) -> ShardLayout: declared = _ATOM.role(role).layout return ShardLayout( - layout=declared.arrangement.fixed(), + layout=declared.layout.fixed(), attrs=declared.attrs, - mesh=declared.mesh, + mesh=WARP, ) From 2a59374a398caaab9a77c6cc4bef39e84b15e581 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 13:37:36 +0800 Subject: [PATCH 13/19] fix(mesh): bound suffix scope replacements --- docs/spec/shard.md | 15 ++++++- src/tilefoundry/ir/mesh_scope.py | 34 ++------------ src/tilefoundry/ir/types/mesh.py | 77 ++++++++++++++++++++++++++++++++ 3 files changed, 94 insertions(+), 32 deletions(-) diff --git a/docs/spec/shard.md b/docs/spec/shard.md index 030d904e..f693f641 100644 --- a/docs/spec/shard.md +++ b/docs/spec/shard.md @@ -317,6 +317,16 @@ def levels(mesh: Mesh) -> tuple[Layout, ...]: def starts(mesh: Mesh) -> tuple[int, ...]: """Decode each topology level's start from device numbering.""" ... + + +def selected_run(arrangement: Layout, start: int) -> tuple[tuple, tuple, int]: + """Reduce one level's selected positions to its joined modes and start.""" + ... + + +def within_scope(mesh: Mesh, current: Mesh) -> bool: + """Return whether each continuous run selected by mesh is within current.""" + ... ``` - constraints: @@ -362,7 +372,10 @@ Mesh composition uses the following rules: level's start and take each replaced level's start and arrangement from the inner mesh. The combined `ComposedLayout.offset` MUST then be re-encoded in device numbering from those per-level starts. Replacing an unsliced suffix - and replacing the whole mesh retain their existing behavior. + and replacing the whole mesh retain their existing behavior. Every replaced + level MUST reduce to one continuous run contained in the enclosing level's + continuous run. A replacement or enclosing selection that does not reduce to + one continuous run MUST be rejected rather than approximated as an interval. - `make_mesh(*meshes)` invokes `check_topology` on its result. For each named level with a concrete declared extent, its position count MUST NOT exceed that extent; symbolic extents are deferred until dimensions are bound. A diff --git a/src/tilefoundry/ir/mesh_scope.py b/src/tilefoundry/ir/mesh_scope.py index f7aabbd0..ae40d524 100644 --- a/src/tilefoundry/ir/mesh_scope.py +++ b/src/tilefoundry/ir/mesh_scope.py @@ -12,7 +12,7 @@ from tilefoundry.ir.types.int_tuple import flatten, product from tilefoundry.ir.types.layout import ComposedLayout, Layout, size from tilefoundry.ir.types.layout_algebra import is_inverse_projectable -from tilefoundry.ir.types.mesh import Mesh, check_topology, levels, starts +from tilefoundry.ir.types.mesh import Mesh, check_topology, levels, selected_run, starts from tilefoundry.ir.types.storage import StorageKind, resolve_storage from tilefoundry.ir.types.stride import compact_major @@ -45,34 +45,6 @@ def device_layout(mesh: Mesh) -> Layout: return Layout(shape=tuple(shape), strides=tuple(strides)) -def _selected(arrangement: Layout, start: int) -> tuple[tuple, tuple, int]: - """One level's positions as a set of them reads: its modes, and where it starts. - - An axis of one position names no instance, and modes written in another - order state the same positions, so the modes come back sorted by step with - the adjacent ones joined and the ones of a single position left out. - """ - strides = arrangement.strides - if strides is None: - return tuple(flatten(arrangement.shape)), (), start - modes = [ - (extent, stride) - for extent, stride in zip(flatten(arrangement.shape), flatten(strides)) - if extent != 1 - ] - joined: list[list] = [] - for extent, stride in sorted(modes, key=lambda mode: (mode[1], mode[0])): - if joined and joined[-1][0] * joined[-1][1] == stride: - joined[-1][0] *= extent - else: - joined.append([extent, stride]) - return ( - tuple(extent for extent, _ in joined), - tuple(stride for _, stride in joined), - start, - ) - - def covered_by_scope(mesh: Mesh, current: Mesh) -> bool: """Whether *mesh* selects exactly the positions the enclosing scope does. @@ -83,14 +55,14 @@ def covered_by_scope(mesh: Mesh, current: Mesh) -> bool: however either of them wrote the axes down. """ scope = { - getattr(topology, "name", topology): _selected(arrangement, start) + getattr(topology, "name", topology): selected_run(arrangement, start) for topology, arrangement, start in zip( current.topologies, levels(current), starts(current) ) } return all( getattr(topology, "name", topology) in scope - and _selected(arrangement, start) + and selected_run(arrangement, start) == scope[getattr(topology, "name", topology)] for topology, arrangement, start in zip( mesh.topologies, levels(mesh), starts(mesh) diff --git a/src/tilefoundry/ir/types/mesh.py b/src/tilefoundry/ir/types/mesh.py index 24b50064..bdfc7ec0 100644 --- a/src/tilefoundry/ir/types/mesh.py +++ b/src/tilefoundry/ir/types/mesh.py @@ -255,6 +255,61 @@ def starts(mesh: Mesh) -> tuple[int, ...]: return tuple(idx2crd(offset, sizes, compact_major(sizes))) +def selected_run(arrangement: Layout, start: int) -> tuple[tuple, tuple, int]: + """Reduce one level's selected positions to its joined modes and start.""" + strides = arrangement.strides + if strides is None: + return tuple(flatten(arrangement.shape)), (), start + modes = [ + (extent, stride) + for extent, stride in zip(flatten(arrangement.shape), flatten(strides)) + if extent != 1 + ] + joined: list[list] = [] + for extent, stride in sorted(modes, key=lambda mode: (mode[1], mode[0])): + if joined and joined[-1][0] * joined[-1][1] == stride: + joined[-1][0] *= extent + else: + joined.append([extent, stride]) + return ( + tuple(extent for extent, _ in joined), + tuple(stride for _, stride in joined), + start, + ) + + +def _continuous_interval(run: tuple[tuple, tuple, int]) -> tuple[int, int] | None: + extents, strides, start = run + if not isinstance(start, int): + return None + if not extents: + return start, start + 1 + if len(extents) != 1 or strides != (1,) or not isinstance(extents[0], int): + return None + return start, start + extents[0] + + +def within_scope(mesh: Mesh, current: Mesh) -> bool: + """Whether each continuous run selected by *mesh* is within *current*.""" + scope = { + getattr(topology, "name", topology): selected_run(arrangement, start) + for topology, arrangement, start in zip( + current.topologies, levels(current), starts(current) + ) + } + for topology, arrangement, start in zip( + mesh.topologies, levels(mesh), starts(mesh) + ): + name = getattr(topology, "name", topology) + inner = _continuous_interval(selected_run(arrangement, start)) + outer = _continuous_interval(scope[name]) if name in scope else None + if inner is None or outer is None: + return False + if not (outer[0] <= inner[0] and inner[1] <= outer[1]): + return False + return True + + def check_topology(mesh: Mesh) -> None: """Reject static mesh positions beyond their declared topology extents. @@ -324,6 +379,7 @@ def make_mesh(*meshes: Mesh) -> Mesh: elif set(here) <= set(there): result = inner elif len(there) < len(here) and here[-len(there) :] == there: + current = result kept = len(here) - len(there) above = levels(result)[:kept] named = sum(len(flatten(level.shape)) for level in above) @@ -335,6 +391,25 @@ def make_mesh(*meshes: Mesh) -> Mesh: sliced=isinstance(result.layout, ComposedLayout) or isinstance(inner.layout, ComposedLayout), ) + if not within_scope(result, current): + parent_runs = { + name: selected_run(arrangement, start) + for name, arrangement, start in zip( + here, levels(current), starts(current) + ) + if name in there + } + inner_runs = { + name: selected_run(arrangement, start) + for name, arrangement, start in zip( + there, levels(inner), starts(inner) + ) + } + raise ValueError( + f"replacement scope selects runs {inner_runs}, outside parent " + f"scope runs {parent_runs}; both must be continuous and each " + "replacement run must be contained in its parent run" + ) else: shared = sorted(set(here) & set(there)) unnamed = sorted(set(here) - set(there)) @@ -368,6 +443,8 @@ def separate(mesh: Mesh) -> tuple[Mesh, ...]: "check_topology", "levels", "make_mesh", + "selected_run", "separate", "starts", + "within_scope", ] From df02412b35b4b59dd5b9818de966966a89685395 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 15:44:50 +0800 Subject: [PATCH 14/19] refactor(pattern): construct patterns from layouts --- docs/spec/core-ir.md | 5 ++-- src/tilefoundry/ir/pattern/__init__.py | 2 -- src/tilefoundry/ir/pattern/pattern.py | 32 ++++++++++++++++++++ src/tilefoundry/ir/pattern/utils.py | 35 +--------------------- src/tilefoundry/ir/tir/cuda/nn/mma_atom.py | 6 ++-- src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py | 6 ++-- src/tilefoundry/ir/tir/cuda/nn/wgmma.py | 3 +- 7 files changed, 43 insertions(+), 46 deletions(-) diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 40194efc..9aa663ed 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -624,13 +624,14 @@ The implementation is split by responsibility under `ir/pattern/`: description helpers. An unstated (`None`) pattern field admits any value. - `constraint.py` owns cross-operand `Constraint`, `DistinctConstraint`, `SameConstraint`, and `SameModesConstraint` values. -- `utils.py` owns exact-layout construction plus specialization naming and - dimension lookup. +- `utils.py` owns specialization naming and dimension lookup. `LayoutPattern` matches only a bare `Layout`, preserves its nested mode structure, and checks `forward` and `injective` over the whole flattened arrangement by default. A sliced layout must be stated explicitly with `ComposedLayoutPattern`; callers that accept both forms use `OrPattern`. +`LayoutPattern.from_layout(layout, ...)` constructs the exact bare or composed +pattern for an authored arrangement, preserving its nested structure. With `per_mode=True`, `LayoutPattern` checks each top-level mode independently; `MeshPattern` requires this explicit form because each mesh level uses its own numbering space. `MeshPattern` never changes the supplied pattern implicitly. diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py index 67294a7e..fb358cef 100644 --- a/src/tilefoundry/ir/pattern/__init__.py +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -54,7 +54,6 @@ WHOLE_BYTES, _mangle_variant_name, any_threads, - arrangement_pattern, dtype_place, locate_dim_var, operand_tile, @@ -101,7 +100,6 @@ "affine_part", "alternatives_of", "any_threads", - "arrangement_pattern", "between_rules", "evaluated", "dtype_place", diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index bc5be1b7..30264512 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -520,6 +520,38 @@ class LayoutPattern(Pattern): injective: bool = True per_mode: bool = False + @classmethod + def from_layout( + cls, + layout, + *, + forward: bool = True, + injective: bool = True, + per_mode: bool = False, + ): + """Build the exact pattern for one authored arrangement.""" + if isinstance(layout, ComposedLayout): + inner = layout.inner + return ComposedLayoutPattern( + SwizzlePattern(inner.bits, inner.base, inner.shift) + if isinstance(inner, Swizzle) + else inner, + layout.offset, + cls.from_layout( + layout.outer, + forward=forward, + injective=injective, + per_mode=per_mode, + ), + ) + return cls( + tuple(layout.shape), + tuple(layout.strides), + forward=forward, + injective=injective, + per_mode=per_mode, + ) + def positions(self) -> tuple: return (*flatten(self.shape), *flatten(self.strides)) diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py index 216cf9e4..7033ed44 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -3,7 +3,7 @@ from __future__ import annotations from tilefoundry.ir.core.param_def import ParamDef -from tilefoundry.ir.types import ComposedLayout, Mesh, StorageKind, Swizzle +from tilefoundry.ir.types import Mesh, StorageKind from .pattern import ( AttrPattern, @@ -16,7 +16,6 @@ OrPattern, Pattern, RangePattern, - SwizzlePattern, TensorPattern, VectorPattern, WildcardPattern, @@ -87,37 +86,6 @@ def any_threads() -> ParamDef: ) -def arrangement_pattern( - layout, - *, - forward: bool = True, - injective: bool = True, - per_mode: bool = False, -): - """Build the exact pattern for one authored arrangement.""" - if isinstance(layout, ComposedLayout): - inner = layout.inner - return ComposedLayoutPattern( - SwizzlePattern(inner.bits, inner.base, inner.shift) - if isinstance(inner, Swizzle) - else inner, - layout.offset, - arrangement_pattern( - layout.outer, - forward=forward, - injective=injective, - per_mode=per_mode, - ), - ) - return LayoutPattern( - tuple(layout.shape), - tuple(layout.strides), - forward=forward, - injective=injective, - per_mode=per_mode, - ) - - def locate_dim_var(params: tuple, name: str) -> tuple[int, int] | None: """Return the first parameter/axis carrying a ``DimVar`` named *name*.""" for index, param in enumerate(params): @@ -144,7 +112,6 @@ def _mangle_variant_name(name: str, specializations: tuple[Pattern, ...]) -> str "WHOLE_BYTES", "_mangle_variant_name", "any_threads", - "arrangement_pattern", "dtype_place", "locate_dim_var", "operand_tile", diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py index e2132bf5..59fc31e4 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py @@ -10,6 +10,7 @@ ABSENT, CapturePattern, ComposedLayoutPattern, + LayoutPattern, Match, MeshPattern, MultipleOfPattern, @@ -17,7 +18,6 @@ OrPattern, Pattern, alternatives_of, - arrangement_pattern, matched, resolved, ) @@ -145,10 +145,10 @@ def required_scope(self) -> Mesh: def scope_pattern(cls) -> MeshPattern: topology, = cls.scope.topologies size = topology.size - bare = arrangement_pattern(cls.scope.layout, per_mode=True) + bare = LayoutPattern.from_layout(cls.scope.layout, per_mode=True) sliced = ComposedLayoutPattern( offset=CapturePattern("p0", MultipleOfPattern(size)), - outer=arrangement_pattern(cls.scope.layout, per_mode=True), + outer=LayoutPattern.from_layout(cls.scope.layout, per_mode=True), ) return MeshPattern((topology.name,), OrPattern(sliced, bare)) diff --git a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py index e483cd39..e16ebe72 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py @@ -4,12 +4,12 @@ from tilefoundry.ir.pattern import ( ComposedLayoutPattern, + LayoutPattern, MeshPattern, OrPattern, ShardLayoutPattern, TensorPattern, WildcardPattern, - arrangement_pattern, ) from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology from tilefoundry.ir.types.storage import StorageKind as S @@ -38,7 +38,7 @@ mesh=WARP, ) -_WARP_LAYOUT = arrangement_pattern(WARP.layout, per_mode=True) +_WARP_LAYOUT = LayoutPattern.from_layout(WARP.layout, per_mode=True) _WARP_PATTERN = MeshPattern( ("thread",), OrPattern( @@ -54,7 +54,7 @@ def _fragment(shape: tuple, dtype, held: ShardLayout) -> TensorPattern: dtype=dtype, storage=S.RMEM, layout=ShardLayoutPattern( - layout=arrangement_pattern(held.layout), + layout=LayoutPattern.from_layout(held.layout), attrs=held.attrs, mesh=_WARP_PATTERN, ), diff --git a/src/tilefoundry/ir/tir/cuda/nn/wgmma.py b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py index 38ef3a13..27a2f4a2 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/wgmma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py @@ -21,7 +21,6 @@ SwizzlePattern, TensorPattern, WildcardPattern, - arrangement_pattern, ) from tilefoundry.ir.pattern.match import is_symbolic from tilefoundry.ir.types import Broadcast, DType, Layout, Mesh, Split, Topology @@ -149,7 +148,7 @@ def Fragment(rows: int, cols) -> LayoutPattern: SHARED_BY_ALL = (Broadcast(), Broadcast(), Broadcast()) HELD_PER_THREAD = (Split(2), Split(0), Split(4)) -_WARPGROUP_LAYOUT = arrangement_pattern(WARPGROUP.layout, per_mode=True) +_WARPGROUP_LAYOUT = LayoutPattern.from_layout(WARPGROUP.layout, per_mode=True) _WARPGROUP_PATTERN = MeshPattern( ("thread",), OrPattern( From 1186a5b96e53079df26b0014f522d218eeee8a39 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 16:14:11 +0800 Subject: [PATCH 15/19] refactor(pattern): name computed layout predicates --- docs/spec/core-ir.md | 25 +- src/tilefoundry/ir/pattern/__init__.py | 10 +- src/tilefoundry/ir/pattern/constraint.py | 30 +- src/tilefoundry/ir/pattern/pattern.py | 259 +++++------- src/tilefoundry/ir/pattern/predicates.py | 392 ++++++++++++++++++ src/tilefoundry/ir/pattern/utils.py | 20 +- .../ir/tir/cuda/memory/copy_async_tensor.py | 270 ++---------- src/tilefoundry/ir/tir/cuda/nn/mma.py | 5 +- src/tilefoundry/ir/tir/cuda/nn/mma_atom.py | 8 +- src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py | 8 +- src/tilefoundry/ir/tir/cuda/nn/wgmma.py | 8 +- tests/ir/pattern/test_mesh_pattern.py | 40 +- 12 files changed, 634 insertions(+), 441 deletions(-) create mode 100644 src/tilefoundry/ir/pattern/predicates.py diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 9aa663ed..7a5cd910 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -613,28 +613,35 @@ class Pattern: The implementation is split by responsibility under `ir/pattern/`: -- `pattern.py` defines `Pattern` and the composable classes +- `pattern.py` defines `Pattern`, the computed-condition base `Predicate`, and + the composable classes `OrPattern`, `AndPattern`, `SequencePattern`, `CapturePattern`, `ConstraintPattern`, `GuardPattern`, `SwitchPattern`, `RangePattern`, `MultipleOfPattern`, `OneOfPattern`, `AttrPattern`, `BitsPattern`, `LayoutPattern`, `SwizzlePattern`, `ComposedLayoutPattern`, `MeshPattern`, `ShardLayoutPattern`, `ScalarPattern`, `TensorPattern`, and `WildcardPattern`. It also owns the `Scalar` and `Tensor` singletons. +- `predicates.py` defines named arrangement predicates: `Forward`, + `Injective`, `WholeVectors`, `PlainArrangement`, `BoxDims`, and `TensorMap`. - `match.py` owns matches, captures, symbolic resolution, and the shared description helpers. An unstated (`None`) pattern field admits any value. - `constraint.py` owns cross-operand `Constraint`, `DistinctConstraint`, `SameConstraint`, and `SameModesConstraint` values. - `utils.py` owns specialization naming and dimension lookup. -`LayoutPattern` matches only a bare `Layout`, preserves its nested mode -structure, and checks `forward` and `injective` over the whole flattened -arrangement by default. A sliced layout must be stated explicitly with -`ComposedLayoutPattern`; callers that accept both forms use `OrPattern`. +`LayoutPattern` optionally matches a bare `Layout`'s nested `shape` and +`strides`, then applies its table of named predicates. Omitting both structural +fields leaves the structure unconstrained and lets predicates read through +supported composed or sharded forms. `Forward()` and `Injective()` express the +corresponding computed properties; they are not implicit. `LayoutPattern.from_layout(layout, ...)` constructs the exact bare or composed -pattern for an authored arrangement, preserving its nested structure. -With `per_mode=True`, `LayoutPattern` checks each top-level mode independently; -`MeshPattern` requires this explicit form because each mesh level uses its own -numbering space. `MeshPattern` never changes the supplied pattern implicitly. +pattern for an authored arrangement, preserving its nested structure and the +explicitly supplied predicate table. `Forward(per_mode=True)` and +`Injective(per_mode=True)` check each top-level mode independently. +`MeshPattern` rejects any supplied arrangement predicate that exposes +`per_mode=False`, because each mesh level uses its own numbering space; an +empty predicate table is allowed. It never changes the supplied pattern +implicitly. `ShardLayoutPattern` names the same `layout`, `attrs`, and `mesh` fields as `ShardLayout`: `layout` and `mesh` are nested patterns, while `attrs` remains an exact structural value. Its mesh pattern may state bare and sliced forms diff --git a/src/tilefoundry/ir/pattern/__init__.py b/src/tilefoundry/ir/pattern/__init__.py index fb358cef..ba2db48b 100644 --- a/src/tilefoundry/ir/pattern/__init__.py +++ b/src/tilefoundry/ir/pattern/__init__.py @@ -1,11 +1,11 @@ """Public operation-declaration pattern language.""" +from . import predicates from .constraint import ( Constraint, DistinctConstraint, SameConstraint, SameModesConstraint, - affine_part, ) from .match import ( ABSENT, @@ -36,6 +36,7 @@ OneOfPattern, OrPattern, Pattern, + Predicate, RangePattern, Scalar, ScalarPattern, @@ -45,9 +46,7 @@ SwizzlePattern, Tensor, TensorPattern, - VectorPattern, WildcardPattern, - vector_widths, ) from .utils import ( MOVED_STORAGES, @@ -82,6 +81,7 @@ "OneOfPattern", "OrPattern", "Pattern", + "Predicate", "RangePattern", "SameConstraint", "SameModesConstraint", @@ -97,7 +97,6 @@ "WildcardPattern", "WHOLE_BYTES", "_mangle_variant_name", - "affine_part", "alternatives_of", "any_threads", "between_rules", @@ -111,7 +110,6 @@ "relations_of", "resolved", "storage_place", - "VectorPattern", + "predicates", "whole_vectors", - "vector_widths", ] diff --git a/src/tilefoundry/ir/pattern/constraint.py b/src/tilefoundry/ir/pattern/constraint.py index 397d25c0..62ca3317 100644 --- a/src/tilefoundry/ir/pattern/constraint.py +++ b/src/tilefoundry/ir/pattern/constraint.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from math import prod -from tilefoundry.ir.types import Broadcast, ComposedLayout, Layout, ShardLayout, Swizzle +from tilefoundry.ir.types import ComposedLayout, Layout, ShardLayout, Swizzle from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.tensor_type import TensorType @@ -64,24 +64,6 @@ def written(self) -> str: return f"{self.left}.{self.field} = {self.right}.{self.field}" -def affine_part(layout, *, plain: bool = False): - """Return the strided affine part beneath shard frames and swizzles. - - When *plain* is true, refuse a composition with a transform or offset. - """ - if isinstance(layout, ShardLayout): - if not all(isinstance(attr, Broadcast) for attr in layout.attrs): - return None - layout = layout.layout - if isinstance(layout, ComposedLayout): - if plain and (layout.inner is not None or layout.offset != 0): - return None - if layout.inner is not None and not isinstance(layout.inner, Swizzle): - return None - layout = layout.outer - return layout if isinstance(layout, Layout) and layout.strides is not None else None - - @dataclass(frozen=True, init=False) class SameModesConstraint(Constraint): """Require two tensor layouts to run along the same final mode of one axis.""" @@ -97,8 +79,13 @@ def pair(self, operands: dict): @staticmethod def reading(tensor: TensorType, arrangement=None): - layout = affine_part(tensor.layout if arrangement is None else arrangement) - if layout is None: + layout = tensor.layout if arrangement is None else arrangement + if isinstance(layout, ComposedLayout): + if layout.inner is not None and not isinstance(layout.inner, Swizzle): + layout = None + else: + layout = layout.outer + if not isinstance(layout, Layout) or layout.strides is None: return None, f"{tensor.layout!r} is no strided arrangement" shape, strides = tuple(layout.shape), tuple(layout.strides) extents = tuple(tensor.shape) @@ -190,5 +177,4 @@ def refused(self, operands: dict) -> str: "DistinctConstraint", "SameConstraint", "SameModesConstraint", - "affine_part", ] diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index 30264512..75c71af2 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -16,13 +16,10 @@ ) from tilefoundry.ir.types.int_tuple import congruent from tilefoundry.ir.types.layout import flatten -from tilefoundry.ir.types.layout_algebra import coalesce, is_inverse_projectable from tilefoundry.ir.types.mesh import separate -from .constraint import affine_part from .match import ( ABSENT, - ARRANGEMENT, UNNAMED_PLACE, Match, _named, @@ -66,113 +63,59 @@ def resolve(self, bindings): @dataclass(frozen=True) -class WildcardPattern(Pattern): - """Match any value without binding it.""" - - def match(self, subject, captures=None) -> Match: - return Match(dict(captures or {})) - - def describe(self, name: str = UNNAMED_PLACE) -> str: - return name - - -VECTOR_READING = ( - "every run: each tile axis's modes walked fastest first, contiguous ones joined; " - "the run at step 1 and every other step a whole number of vectors" -) - - -def _reverse_group(group): - if isinstance(group, tuple): - return tuple(_reverse_group(mode) for mode in reversed(group)) - return group - - -def _row_major_groups_for_cute(layout: Layout) -> Layout: - """Reverse modes within each tile axis for CuTe's mode-0-fast algebra.""" - return Layout( - shape=tuple(_reverse_group(group) for group in layout.shape), - strides=tuple(_reverse_group(group) for group in layout.strides), - ) - - -def vector_widths(layout, element_bits: int, widths: tuple[int, ...]) -> tuple[int, ...]: - """Every requested byte width that divides every run in an arrangement.""" - if not widths: - return () - widest = widths[-1] - if isinstance(layout, ShardLayout): - layout = layout.layout - inner = getattr(layout, "inner", None) - if inner is not None and hasattr(inner, "base"): - widest = min(widest, 1 << inner.base) - held = affine_part(layout) - if held is None or any( - type(value) is not int - for group in (held.shape, held.strides) - for value in flatten(group) - ): - return () - grouped = coalesce( - _row_major_groups_for_cute(held), - (0,) * len(held.shape), - ) - runs = tuple( - sorted( - zip(flatten(grouped.shape), flatten(grouped.strides)), - key=lambda run: run[1], - ) - ) - unit = [extent for extent, step in runs if step == 1] - if len(unit) != 1: - return () - counted = (unit[0], *(step for _, step in runs if step != 1)) - return tuple( - width - for width in widths - if width <= widest - and all(value * element_bits % (width * 8) == 0 for value in counted) - ) - - -@dataclass(frozen=True) -class VectorPattern(Pattern): - """An arrangement that moves whole vectors at the requested byte widths.""" - - width: CapturePattern - dtype: str - widths: tuple[int, ...] +class Predicate(Pattern): + """A named computed condition over one authored arrangement.""" + + @staticmethod + def arrangement(subject) -> Layout | None: + """Read the static strided layout beneath shard and composition wrappers.""" + if isinstance(subject, ShardLayout): + subject = subject.layout + if isinstance(subject, ComposedLayout): + if subject.inner is not None and not isinstance(subject.inner, Swizzle): + return None + subject = subject.outer + if not isinstance(subject, Layout) or subject.strides is None: + return None + extents = tuple(flatten(subject.shape)) + strides = tuple(flatten(subject.strides)) + if any(type(number) is not int for number in (*extents, *strides)): + return None + if any(extent <= 0 for extent in extents): + return None + return subject - def available_widths(self, subject, captures) -> tuple[int, ...]: - bits = getattr(dict(captures or {}).get(self.dtype), "bit_width", None) - return () if type(bits) is not int else vector_widths(subject, bits, self.widths) + def holds(self, arrangement: Layout, captures: dict) -> bool: + raise NotImplementedError - def match(self, subject, captures=None): + def match(self, subject, captures=None) -> Match | None: held = dict(captures or {}) - widths = self.available_widths(subject, held) - if not widths: - return None - if self.width.name in held: - return Match(held) if held[self.width.name] in widths else None - return matched(self.width, widths[-1], held) + arrangement = self.arrangement(subject) + return Match(held) if arrangement is not None and self.holds(arrangement, held) else None def refusal(self, subject, captures=None) -> str | None: - if self.match(subject, captures) is not None: - return None - sizes = " or ".join(map(str, self.widths)) - if len(self.widths) > 2: - sizes = ", ".join(map(str, self.widths[:-1])) + f" or {self.widths[-1]}" return ( - f"{subject!r} moves no whole vector of {sizes} bytes -- its run at step 1 " - "and every other step are no whole number of one -- so the two ends share " - "no run wide enough for the requested vector widths" + None + if self.match(subject, captures) is not None + else f"{subject!r} does not satisfy {self.describe()}" ) def describe(self, name: str = UNNAMED_PLACE) -> str: - return f"vectors of {self.width.name} bytes" + raise NotImplementedError def relations(self) -> tuple[str, ...]: - return (VECTOR_READING, *relations_of((self.width,))) + raise NotImplementedError + + +@dataclass(frozen=True) +class WildcardPattern(Pattern): + """Match any value without binding it.""" + + def match(self, subject, captures=None) -> Match: + return Match(dict(captures or {})) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return name @dataclass(frozen=True, init=False) @@ -512,24 +455,21 @@ def fixed(self): @dataclass(frozen=True) class LayoutPattern(Pattern): - """An affine layout, preserving grouping and optional whole-layout rules.""" + """An optional layout structure with named computed predicates.""" - shape: tuple - strides: tuple - forward: bool = True - injective: bool = True - per_mode: bool = False + shape: tuple | None = None + strides: tuple | None = None + predicates: tuple[Predicate, ...] = field(default_factory=tuple) @classmethod def from_layout( cls, layout, *, - forward: bool = True, - injective: bool = True, - per_mode: bool = False, + predicates: tuple[Predicate, ...] = (), ): """Build the exact pattern for one authored arrangement.""" + held = tuple(predicates) if isinstance(layout, ComposedLayout): inner = layout.inner return ComposedLayoutPattern( @@ -539,71 +479,83 @@ def from_layout( layout.offset, cls.from_layout( layout.outer, - forward=forward, - injective=injective, - per_mode=per_mode, + predicates=held, ), ) return cls( tuple(layout.shape), tuple(layout.strides), - forward=forward, - injective=injective, - per_mode=per_mode, + predicates=held, ) def positions(self) -> tuple: - return (*flatten(self.shape), *flatten(self.strides)) + return ( + *(flatten(self.shape) if self.shape is not None else ()), + *(flatten(self.strides) if self.strides is not None else ()), + ) - def match(self, subject, captures=None): - layout = subject - if not isinstance(layout, Layout) or layout.strides is None: + def _match_structure(self, subject, captures=None) -> Match | None: + held = Match(dict(captures or {})) + if self.shape is None and self.strides is None: + return held + if not isinstance(subject, Layout) or subject.strides is None: return None - if not ( - congruent(layout.shape, self.shape) - and congruent(layout.strides, self.strides) - ): + if self.shape is not None and not congruent(subject.shape, self.shape): + return None + if self.strides is not None and not congruent(subject.strides, self.strides): return None - extents = tuple(flatten(layout.shape)) - strides = tuple(flatten(layout.strides)) + extents = tuple(flatten(subject.shape)) + strides = tuple(flatten(subject.strides)) if any(type(number) is not int for number in (*extents, *strides)): return None if any(number <= 0 for number in extents): return None - held = matched(SequencePattern(*self.positions()), (*extents, *strides), captures) + values = ( + *(extents if self.shape is not None else ()), + *(strides if self.strides is not None else ()), + ) + return matched(SequencePattern(*self.positions()), values, held.captures) + + def match(self, subject, captures=None): + held = self._match_structure(subject, captures) if held is None: return None - arrangements = ( - tuple( - Layout(tuple(flatten(shape)), tuple(flatten(steps))) - for shape, steps in zip(layout.shape, layout.strides) - ) - if self.per_mode - else (Layout(extents, strides),) - ) - for arrangement in arrangements: - held_strides = tuple(flatten(arrangement.strides)) - if self.forward and any(stride < 0 for stride in held_strides): - return None - if self.injective and not is_inverse_projectable(arrangement): + for predicate in self.predicates: + held = matched(predicate, subject, held.captures) + if held is None: return None return held + def refusal(self, subject, captures=None) -> str | None: + held = self._match_structure(subject, captures) + if held is None: + return f"{subject!r} is not {self.describe()}" + for predicate in self.predicates: + found = matched(predicate, subject, held.captures) + if found is not None: + held = found + continue + explained = getattr(predicate, "refusal", None) + return ( + explained(subject, held.captures) + if explained is not None + else f"{subject!r} does not satisfy {predicate.describe()}" + ) + return None + def describe(self, name: str = UNNAMED_PLACE) -> str: - return ( - f"Layout({written_grouped(tuple(self.shape))}, {written_grouped(tuple(self.strides))})" - ) + if self.shape is None and self.strides is None: + return "layout" + shape = name if self.shape is None else written_grouped(tuple(self.shape)) + strides = name if self.strides is None else written_grouped(tuple(self.strides)) + return f"Layout({shape}, {strides})" def relations(self) -> tuple[str, ...]: - lines = list(relations_of(self.positions())) - subject = "each top-level mode" if self.per_mode else ARRANGEMENT - if self.forward: - lines.append(f"{subject} has no backward step") - if self.injective: - lines.append(f"{subject} reaches each of its own slots exactly once") - return tuple(lines) + return (*relations_of(self.positions()), *relations_of(self.predicates)) def fixed(self): + if self.shape is None or self.strides is None: + return None if any(isinstance(value, Pattern) for value in self.positions()): return None return Layout(tuple(self.shape), tuple(self.strides)) @@ -695,10 +647,12 @@ def require_per_mode(pattern) -> None: return if isinstance(pattern, ComposedLayoutPattern): pattern = pattern.outer - if not isinstance(pattern, LayoutPattern) or not pattern.per_mode: + if not isinstance(pattern, LayoutPattern) or any( + getattr(predicate, "per_mode", True) is False for predicate in pattern.predicates + ): raise ValueError( - "MeshPattern layout must be a LayoutPattern(per_mode=True), " - "or a ComposedLayoutPattern whose outer uses per_mode=True" + "MeshPattern layout predicates must check each top-level mode, " + "or a ComposedLayoutPattern outer must use per-mode predicates" ) require_per_mode(self.layout) @@ -890,6 +844,7 @@ def describe(self, name: str = UNNAMED_PLACE, arrangements=None) -> str: "OneOfPattern", "OrPattern", "Pattern", + "Predicate", "RangePattern", "Scalar", "ScalarPattern", @@ -899,7 +854,5 @@ def describe(self, name: str = UNNAMED_PLACE, arrangements=None) -> str: "SwitchPattern", "Tensor", "TensorPattern", - "VectorPattern", "WildcardPattern", - "vector_widths", ] diff --git a/src/tilefoundry/ir/pattern/predicates.py b/src/tilefoundry/ir/pattern/predicates.py new file mode 100644 index 00000000..bfef4f48 --- /dev/null +++ b/src/tilefoundry/ir/pattern/predicates.py @@ -0,0 +1,392 @@ +"""Named computed predicates for operation-declaration layout patterns.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +from tilefoundry.ir.types import ComposedLayout, Layout, ShardLayout +from tilefoundry.ir.types.int_tuple import flatten +from tilefoundry.ir.types.layout_algebra import coalesce, is_inverse_projectable + +from .match import ( + ARRANGEMENT, + UNNAMED_PLACE, + Match, + matched, + relations_of, + written_place, + written_tuple, +) +from .pattern import CapturePattern, Predicate, SequencePattern + +VECTOR_READING = ( + "every run: each tile axis's modes walked fastest first, contiguous ones joined; " + "the run at step 1 and every other step a whole number of vectors" +) +TENSORMAP_READING = "every tensormap: one dim per mode of the tile, the mode at step 1 first" + + +def _arrangements(layout: Layout, per_mode: bool) -> tuple[Layout, ...]: + if per_mode: + return tuple( + Layout(tuple(flatten(shape)), tuple(flatten(steps))) + for shape, steps in zip(layout.shape, layout.strides) + ) + return (Layout(tuple(flatten(layout.shape)), tuple(flatten(layout.strides))),) + + +@dataclass(frozen=True) +class Forward(Predicate): + """Require nonnegative steps, across the whole layout or per top-level mode.""" + + per_mode: bool = False + + def holds(self, arrangement: Layout, captures: dict) -> bool: + return all( + all(step >= 0 for step in flatten(part.strides)) + for part in _arrangements(arrangement, self.per_mode) + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + subject = "each top-level mode" if self.per_mode else ARRANGEMENT + return f"{subject} with no backward step" + + def relations(self) -> tuple[str, ...]: + subject = "each top-level mode" if self.per_mode else ARRANGEMENT + return (f"{subject} has no backward step",) + + +@dataclass(frozen=True) +class Injective(Predicate): + """Require every slot to be reached once, across the layout or per mode.""" + + per_mode: bool = False + + def holds(self, arrangement: Layout, captures: dict) -> bool: + return all( + is_inverse_projectable(part) for part in _arrangements(arrangement, self.per_mode) + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + subject = "each top-level mode" if self.per_mode else ARRANGEMENT + return f"{subject} reaching each of its own slots exactly once" + + def relations(self) -> tuple[str, ...]: + subject = "each top-level mode" if self.per_mode else ARRANGEMENT + return (f"{subject} reaches each of its own slots exactly once",) + + +def _reverse_group(group): + if isinstance(group, tuple): + return tuple(_reverse_group(mode) for mode in reversed(group)) + return group + + +def _row_major_groups_for_cute(layout: Layout) -> Layout: + """Reverse modes within each tile axis for CuTe's mode-0-fast algebra.""" + return Layout( + shape=tuple(_reverse_group(group) for group in layout.shape), + strides=tuple(_reverse_group(group) for group in layout.strides), + ) + + +def _vector_widths(layout, element_bits: int, widths: tuple[int, ...]) -> tuple[int, ...]: + """Every requested byte width that divides every run in an arrangement.""" + if not widths: + return () + widest = widths[-1] + stated = layout.layout if isinstance(layout, ShardLayout) else layout + inner = getattr(stated, "inner", None) + if inner is not None and hasattr(inner, "base"): + widest = min(widest, 1 << inner.base) + held = Predicate.arrangement(layout) + if held is None: + return () + grouped = coalesce( + _row_major_groups_for_cute(held), + (0,) * len(held.shape), + ) + runs = tuple( + sorted( + zip(flatten(grouped.shape), flatten(grouped.strides)), + key=lambda run: run[1], + ) + ) + unit = [extent for extent, step in runs if step == 1] + if len(unit) != 1: + return () + counted = (unit[0], *(step for _, step in runs if step != 1)) + return tuple( + width + for width in widths + if width <= widest and all(value * element_bits % (width * 8) == 0 for value in counted) + ) + + +@dataclass(frozen=True) +class WholeVectors(Predicate): + """Require whole vectors at one of the requested byte widths.""" + + width: CapturePattern + dtype: str + widths: tuple[int, ...] + + def available_widths(self, subject, captures) -> tuple[int, ...]: + bits = getattr(dict(captures or {}).get(self.dtype), "bit_width", None) + return () if type(bits) is not int else _vector_widths(subject, bits, self.widths) + + def match(self, subject, captures=None): + held = dict(captures or {}) + widths = self.available_widths(subject, held) + if not widths: + return None + if self.width.name in held: + return Match(held) if held[self.width.name] in widths else None + return matched(self.width, widths[-1], held) + + def refusal(self, subject, captures=None) -> str | None: + if self.match(subject, captures) is not None: + return None + sizes = " or ".join(map(str, self.widths)) + if len(self.widths) > 2: + sizes = ", ".join(map(str, self.widths[:-1])) + f" or {self.widths[-1]}" + return ( + f"{subject!r} moves no whole vector of {sizes} bytes -- its run at step 1 " + "and every other step are no whole number of one -- so the two ends share " + "no run wide enough for the requested vector widths" + ) + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return f"vectors of {self.width.name} bytes" + + def relations(self) -> tuple[str, ...]: + return (VECTOR_READING, *relations_of((self.width,))) + + +@dataclass(frozen=True) +class PlainArrangement(Predicate): + """Require an arrangement with no transform and no nonzero offset.""" + + @staticmethod + def _stated(subject): + return subject.layout if isinstance(subject, ShardLayout) else subject + + def holds(self, arrangement: Layout, captures: dict) -> bool: + return True + + def match(self, subject, captures=None): + stated = self._stated(subject) + if isinstance(stated, ComposedLayout) and (stated.inner is not None or stated.offset != 0): + return None + return super().match(subject, captures) + + def refusal(self, subject, captures=None) -> str | None: + if self.match(subject, captures) is not None: + return None + stated = self._stated(subject) + if isinstance(stated, ComposedLayout): + reached = [] + if stated.inner is not None: + reached.append(f"through {stated.inner!r}") + if stated.offset != 0: + reached.append(f"at offset {stated.offset}") + if reached: + return f"it is reached {' '.join(reached)}, not as a plain arrangement" + return f"{subject!r} is no static strided arrangement" + + def describe(self, name: str = UNNAMED_PLACE) -> str: + return "a plain arrangement with no transform or offset" + + def relations(self) -> tuple[str, ...]: + return ("every plain arrangement has no transform or nonzero offset",) + + +@dataclass(frozen=True) +class Run: + """One contiguous run of modes from one logical tile axis.""" + + extent: int + step: int + axis: int + mode: int + + +def box_runs( + layout: Layout, + element_bits: int, + span: int | None, + limit: int, +) -> tuple[Run, ...]: + """Read box runs by tile axis, ordered by increasing step.""" + runs: list[Run] = [] + for axis, (extents, steps) in enumerate(zip(layout.shape, layout.strides)): + modes = tuple(enumerate(zip(flatten(extents), flatten(steps)))) + for mode, (extent, step) in reversed(modes): + if extent == 1: + continue + last = runs[-1] if runs and runs[-1].axis == axis else None + joined = None if last is None else last.extent * extent + if ( + last is not None + and step == last.step * last.extent + and joined <= limit + and not (span is not None and last.step == 1 and joined * element_bits > span * 8) + ): + runs[-1] = replace(last, extent=joined) + else: + runs.append(Run(extent, step, axis, mode)) + return tuple(sorted(runs, key=lambda run: run.step)) + + +def _missed_place(places, values, captures, first: int = 0) -> tuple | None: + held = captures + for index, (place, value) in enumerate(zip(places, values), first): + found = matched(place, value, held) + if found is None: + return index, place, value + held = found.captures + return None + + +@dataclass(frozen=True) +class BoxDims(Predicate): + """Require one box dimension per contiguous run of tile-axis modes.""" + + dims: tuple + dtype: str + limit: int + span: int | None = None + + def reading(self, subject, captures) -> tuple[tuple | None, str | None]: + layout = self.arrangement(subject) + if layout is None: + return None, f"{subject!r} is no static strided arrangement" + width = getattr(captures.get(self.dtype), "bit_width", None) + if type(width) is not int: + return None, f"the element it arranges is not bound as {self.dtype}" + runs = box_runs(layout, width, self.span, self.limit) + if not runs: + return None, "it holds one element, which is no box" + if len(runs) > len(self.dims): + return None, ( + f"it is {len(runs)} runs of modes, and a box has at most {len(self.dims)} dims" + ) + extents = tuple(run.extent for run in runs) + extents += (1,) * (len(self.dims) - len(extents)) + if runs[0].step != 1: + return extents, ( + f"its smallest step is {runs[0].step}, and a box lays its dim 0 at step 1" + ) + if self.span is not None and (self.span * 8) % width: + return extents, f"a {self.span}-byte row is no whole number of {self.dtype}" + expected = runs[0].extent if self.span is None else self.span * 8 // width + for index, run in enumerate(runs[1:], 1): + if run.step != expected: + return extents, ( + f"its dim {index} steps {run.step} where a box lays it at " + f"{expected} ({self.laid()})" + ) + expected *= run.extent + return extents, None + + def laid(self) -> str: + rows = "" if self.span is None else f", rows {self.span} B apart" + return f"dim 0 fastest{rows}" + + def match(self, subject, captures=None): + held = dict(captures or {}) + extents, unlaid = self.reading(subject, held) + if extents is None or unlaid is not None: + return None + return matched(SequencePattern(*self.dims), extents, held) + + def refusal(self, subject, captures=None) -> str | None: + held = dict(captures or {}) + extents, why = self.reading(subject, held) + if extents is None: + return why + missed = _missed_place(self.dims, extents, held) + if missed is not None: + index, place, extent = missed + return ( + f"its box dim {index} holds {extent} elements, and a box reads " + f"{written_place(place.pattern, place.name)}" + ) + return why + + def describe(self, name: str = UNNAMED_PLACE) -> str: + dims = written_tuple(tuple(written_place(place) for place in self.dims)) + return f"box {dims}, {self.laid()}" + + def relations(self) -> tuple[str, ...]: + reading = ( + "every box: each tile axis's modes, contiguous ones joined up to " + f"{self.limit} elements, one dim each, in increasing step" + ) + return (reading, *relations_of(self.dims)) + + +@dataclass(frozen=True) +class TensorMap(Predicate): + """Require one tensor-map dimension per nontrivial tile mode.""" + + steps: tuple + dtype: str + + def reading(self, subject) -> tuple[tuple | None, str | None]: + layout = self.arrangement(subject) + if layout is None: + return None, f"{subject!r} is no static strided tensor a tensormap describes" + modes = [ + (extent, step) + for extents, steps in zip(layout.shape, layout.strides) + for extent, step in zip(flatten(extents), flatten(steps)) + if extent > 1 + ] + unit = [mode for mode in modes if mode[1] == 1] + if len(unit) != 1: + return None, ( + f"{len(unit)} of its modes step 1, and a tensormap's dim 0 is its one " + "contiguous mode" + ) + if len(modes) > len(self.steps) + 1: + return None, ( + f"it is {len(modes)} modes, and a tensormap has at most {len(self.steps) + 1} dims" + ) + others = tuple(step for _, step in modes if step != 1) + return others + (0,) * (len(self.steps) - len(others)), None + + def match(self, subject, captures=None): + steps, _ = self.reading(subject) + return None if steps is None else matched(SequencePattern(*self.steps), steps, captures) + + def refusal(self, subject, captures=None) -> str | None: + held = dict(captures or {}) + steps, why = self.reading(subject) + if steps is None: + return why + missed = _missed_place(self.steps, steps, held, first=1) + if missed is not None: + index, place, step = missed + return ( + f"its dim {index} steps {step} elements, and a tensormap reads " + f"{written_place(place.pattern, place.name)}" + ) + return None + + def describe(self, name: str = UNNAMED_PLACE) -> str: + steps = written_tuple(("1", *(written_place(place) for place in self.steps))) + return f"tensormap at {steps}" + + def relations(self) -> tuple[str, ...]: + return (TENSORMAP_READING, *relations_of(self.steps)) + + +__all__ = [ + "BoxDims", + "Forward", + "Injective", + "PlainArrangement", + "TensorMap", + "WholeVectors", +] diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py index 7033ed44..059d1bfa 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -5,6 +5,7 @@ from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.types import Mesh, StorageKind +from . import predicates as P from .pattern import ( AttrPattern, CapturePattern, @@ -17,7 +18,6 @@ Pattern, RangePattern, TensorPattern, - VectorPattern, WildcardPattern, ) @@ -49,12 +49,16 @@ def storage_place(index: int) -> str: return f"storage{index}" -def whole_vectors(index: int, widths: tuple[int, ...]) -> VectorPattern: +def whole_vectors(index: int, widths: tuple[int, ...]) -> LayoutPattern: """Whole vectors at *widths*, counted in operand *index*'s element dtype.""" - return VectorPattern( - CapturePattern("width", OneOfPattern(widths)), - dtype_place(index), - widths, + return LayoutPattern( + predicates=( + P.WholeVectors( + CapturePattern("width", OneOfPattern(widths)), + dtype_place(index), + widths, + ), + ) ) @@ -64,13 +68,13 @@ def whole_vectors(index: int, widths: tuple[int, ...]) -> VectorPattern: outer=LayoutPattern( ((CapturePattern("n", RangePattern(lo=1)),),), ((1,),), - per_mode=True, + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), ), ), LayoutPattern( ((CapturePattern("n", RangePattern(lo=1)),),), ((1,),), - per_mode=True, + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), ), ) diff --git a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py index 446df7f1..39f3d883 100644 --- a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py +++ b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py @@ -2,14 +2,12 @@ from __future__ import annotations -from dataclasses import dataclass, replace from enum import Enum from tilefoundry.ir.core import Op from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op from tilefoundry.ir.pattern import ( - UNNAMED_PLACE, AndPattern, BitsPattern, CapturePattern, @@ -19,21 +17,17 @@ MeshPattern, MultipleOfPattern, OrPattern, - Pattern, RangePattern, SameModesConstraint, - SequencePattern, SwitchPattern, SwizzlePattern, - affine_part, - matched, - relations_of, utils, ) -from tilefoundry.ir.pattern.match import written_place, written_tuple +from tilefoundry.ir.pattern import ( + predicates as P, +) from tilefoundry.ir.tir.verify import verify_between, verify_operands -from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, ShardLayout, Swizzle, UnitType -from tilefoundry.ir.types.int_tuple import flatten +from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, Swizzle, UnitType from tilefoundry.ir.types.storage import StorageKind as S from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -41,57 +35,9 @@ BOX_EXTENT = 256 TMA_UNIT_BITS = 16 * 8 SWIZZLE_PLACE = "swizzle" -BOX_READING = ( - "every box: each tile axis's modes, contiguous ones joined up to " - f"{BOX_EXTENT} elements, one dim each, in increasing step" -) -TENSORMAP_READING = ( - "every tensormap: one dim per mode of the tile, the mode at step 1 first" -) TMA_STORAGES = (S.GMEM, S.SMEM) -@dataclass(frozen=True) -class Run: - """One contiguous run of modes from one logical tile axis.""" - - extent: int - step: int - axis: int - mode: int - - -def box_runs( - layout: Layout, - element_bits: int, - span: int | None, - limit: int | None = BOX_EXTENT, -) -> tuple[Run, ...]: - """Read TMA box runs by tile axis, ordered by increasing step.""" - runs: list[Run] = [] - for axis, (extents, steps) in enumerate(zip(layout.shape, layout.strides)): - modes = tuple(enumerate(zip(flatten(extents), flatten(steps)))) - for mode, (extent, step) in reversed(modes): - if extent == 1: - continue - last = runs[-1] if runs and runs[-1].axis == axis else None - joined = None if last is None else last.extent * extent - if ( - last is not None - and step == last.step * last.extent - and (limit is None or joined <= limit) - and not ( - span is not None - and last.step == 1 - and joined * element_bits > span * 8 - ) - ): - runs[-1] = replace(last, extent=joined) - else: - runs.append(Run(extent, step, axis, mode)) - return tuple(sorted(runs, key=lambda run: run.step)) - - class TmaSwizzle(Enum): """The shared-memory swizzle selected when the tensor map is encoded.""" @@ -105,107 +51,11 @@ def swizzle(self) -> Swizzle | None: return None if self is TmaSwizzle.NONE else Swizzle(self.value.bit_length() - 5, 4, 3) -def _missed_place(places, values, captures, first: int = 0) -> tuple | None: - held = captures - for index, (place, value) in enumerate(zip(places, values), first): - found = matched(place, value, held) - if found is None: - return index, place, value - held = found.captures - return None - - -@dataclass(frozen=True) -class BoxPattern(Pattern): - """One TMA box landed in shared memory.""" - - dims: tuple - dtype: str - span: int | None = None - - def reading(self, subject, captures) -> tuple[tuple | None, str | None]: - layout = affine_part(subject, plain=True) - if layout is None or any( - type(value) is not int - for group in (layout.shape, layout.strides) - for value in flatten(group) - ): - return None, f"{subject!r} is no static strided arrangement" - width = getattr(captures.get(self.dtype), "bit_width", None) - if type(width) is not int: - return None, f"the element it arranges is not bound as {self.dtype}" - runs = box_runs(layout, width, self.span) - if not runs: - return None, "it holds one element, which is no box" - if len(runs) > len(self.dims): - return None, ( - f"it is {len(runs)} runs of modes, and a box has at most " - f"{len(self.dims)} dims" - ) - extents = tuple(run.extent for run in runs) - extents += (1,) * (len(self.dims) - len(extents)) - if runs[0].step != 1: - return extents, ( - f"its smallest step is {runs[0].step}, and a box lays its dim 0 at step 1" - ) - if self.span is not None and (self.span * 8) % width: - return extents, f"a {self.span}-byte row is no whole number of {self.dtype}" - expected = runs[0].extent if self.span is None else self.span * 8 // width - for index, run in enumerate(runs[1:], 1): - if run.step != expected: - return extents, ( - f"its dim {index} steps {run.step} where a box lays it at " - f"{expected} ({self.laid()})" - ) - expected *= run.extent - return extents, None - - def laid(self) -> str: - rows = "" if self.span is None else f", rows {self.span} B apart" - return f"dim 0 fastest{rows}" - - def match(self, subject, captures=None): - held = dict(captures or {}) - extents, unlaid = self.reading(subject, held) - if extents is None or unlaid is not None: - return None - return matched(SequencePattern(*self.dims), extents, held) - - def refusal(self, subject, captures=None) -> str | None: - held = dict(captures or {}) - extents, why = self.reading(subject, held) - if extents is None: - return why - missed = _missed_place(self.dims, extents, held) - if missed is not None: - index, place, extent = missed - return ( - f"its box dim {index} holds {extent} elements, and a box reads " - f"{written_place(place.pattern, place.name)}" - ) - return why - - def describe(self, name: str = UNNAMED_PLACE) -> str: - dims = written_tuple(tuple(written_place(place) for place in self.dims)) - return f"box {dims}, {self.laid()}" - - def relations(self) -> tuple[str, ...]: - return (BOX_READING, *relations_of(self.dims)) - - -@dataclass(frozen=True, init=False) class BoxFamily(SwitchPattern): """Every unswizzled or swizzled shared-memory box.""" - def match(self, subject, captures=None): - if isinstance(subject, ShardLayout) and affine_part(subject) is not None: - subject = subject.layout - return super().match(subject, captures) - def refusal(self, subject, captures=None) -> str | None: layout = subject - if isinstance(layout, ShardLayout) and affine_part(layout) is not None: - layout = layout.layout transform = layout.inner if isinstance(layout, ComposedLayout) else None if transform is not None and layout.offset != 0: return f"it is reached through {transform!r} at offset {layout.offset}, not 0" @@ -222,67 +72,6 @@ def refusal(self, subject, captures=None) -> str | None: return f"it is reached through {transform!r}, and a tensormap swizzles by {written} or not at all" -@dataclass(frozen=True) -class TensorMapPattern(Pattern): - """The global tile described by one tensor map.""" - - steps: tuple - dtype: str - - def reading(self, subject) -> tuple[tuple | None, str | None]: - layout = affine_part(subject, plain=True) - if layout is None or any( - type(value) is not int - for group in (layout.shape, layout.strides) - for value in flatten(group) - ): - return None, f"{subject!r} is no static strided tensor a tensormap describes" - modes = [ - (extent, step) - for extents, steps in zip(layout.shape, layout.strides) - for extent, step in zip(flatten(extents), flatten(steps)) - if extent > 1 - ] - unit = [mode for mode in modes if mode[1] == 1] - if len(unit) != 1: - return None, ( - f"{len(unit)} of its modes step 1, and a tensormap's dim 0 is its one " - "contiguous mode" - ) - if len(modes) > len(self.steps) + 1: - return None, ( - f"it is {len(modes)} modes, and a tensormap has at most " - f"{len(self.steps) + 1} dims" - ) - others = tuple(step for _, step in modes if step != 1) - return others + (0,) * (len(self.steps) - len(others)), None - - def match(self, subject, captures=None): - steps, _ = self.reading(subject) - return None if steps is None else matched(SequencePattern(*self.steps), steps, captures) - - def refusal(self, subject, captures=None) -> str | None: - held = dict(captures or {}) - steps, why = self.reading(subject) - if steps is None: - return why - missed = _missed_place(self.steps, steps, held, first=1) - if missed is not None: - index, place, step = missed - return ( - f"its dim {index} steps {step} elements, and a tensormap reads " - f"{written_place(place.pattern, place.name)}" - ) - return None - - def describe(self, name: str = UNNAMED_PLACE) -> str: - steps = written_tuple(("1", *(written_place(place) for place in self.steps))) - return f"tensormap at {steps}" - - def relations(self) -> tuple[str, ...]: - return (TENSORMAP_READING, *relations_of(self.steps)) - - def _dim(name: str, dtype: str, *, span: int | None = None) -> CapturePattern: parts = [ RangePattern(lo=1, hi=BOX_EXTENT), @@ -298,7 +87,14 @@ def TmaBoxPattern(dtype: str) -> BoxFamily: CapturePattern(f"dim{index}", RangePattern(lo=1, hi=BOX_EXTENT)) for index in range(1, TMA_RANK) ) - boxes = {TmaSwizzle.NONE: BoxPattern((_dim("dim0", dtype), *rest), dtype)} + boxes = { + TmaSwizzle.NONE: LayoutPattern( + predicates=( + P.PlainArrangement(), + P.BoxDims((_dim("dim0", dtype), *rest), dtype, BOX_EXTENT), + ) + ) + } for index, mode in enumerate( (mode for mode in TmaSwizzle if mode.swizzle is not None), TMA_RANK, @@ -307,22 +103,30 @@ def TmaBoxPattern(dtype: str) -> BoxFamily: boxes[mode] = ComposedLayoutPattern( SwizzlePattern(swizzle.bits, swizzle.base, swizzle.shift), 0, - BoxPattern((_dim(f"dim{index}", dtype, span=mode.value), *rest), dtype, mode.value), + LayoutPattern( + predicates=( + P.PlainArrangement(), + P.BoxDims( + (_dim(f"dim{index}", dtype, span=mode.value), *rest), + dtype, + BOX_EXTENT, + span=mode.value, + ), + ) + ), ) return BoxFamily(SWIZZLE_PLACE, boxes) -def TmaGlobalPattern(dtype: str) -> TensorMapPattern: - return TensorMapPattern( - tuple( - CapturePattern( - f"step{index}", - BitsPattern(dtype, MultipleOfPattern(TMA_UNIT_BITS)), - ) - for index in range(1, TMA_RANK) - ), - dtype, +def TmaGlobalPattern(dtype: str) -> LayoutPattern: + steps = tuple( + CapturePattern( + f"step{index}", + BitsPattern(dtype, MultipleOfPattern(TMA_UNIT_BITS)), + ) + for index in range(1, TMA_RANK) ) + return LayoutPattern(predicates=(P.PlainArrangement(), P.TensorMap(steps, dtype))) def TmaOperandPattern(storage: str, dtype: str) -> SwitchPattern: @@ -336,7 +140,11 @@ def TmaOperandPattern(storage: str, dtype: str) -> SwitchPattern: def _warp_scope() -> MeshPattern: - layout = LayoutPattern(((32,),), ((1,),), per_mode=True) + layout = LayoutPattern( + ((32,),), + ((1,),), + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), + ) sliced = ComposedLayoutPattern( offset=CapturePattern("p0", MultipleOfPattern(32)), outer=layout, @@ -412,20 +220,14 @@ def verify_copy_async_tensor(call: "Call", ctx: "VerifyContext") -> None: __all__ = [ "BOX_EXTENT", - "BOX_READING", "BoxFamily", - "BoxPattern", "CopyAsyncTensor", - "Run", "SWIZZLE_PLACE", - "TENSORMAP_READING", "TMA_RANK", "TMA_STORAGES", "TMA_UNIT_BITS", - "TensorMapPattern", "TmaBoxPattern", "TmaGlobalPattern", "TmaOperandPattern", "TmaSwizzle", - "box_runs", ] diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index d7704ccf..866e60c6 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma.py @@ -13,6 +13,9 @@ MultipleOfPattern, OrPattern, ) +from tilefoundry.ir.pattern import ( + predicates as P, +) from tilefoundry.ir.tir.verify import input_params from tilefoundry.ir.types import DType, Mesh, UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -34,7 +37,7 @@ def _warp_layout_pattern() -> LayoutPattern: return LayoutPattern( ((CapturePattern("n", MultipleOfPattern(32)),),), ((1,),), - per_mode=True, + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), ) diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py index 59fc31e4..bd3bba9e 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py @@ -21,6 +21,9 @@ matched, resolved, ) +from tilefoundry.ir.pattern import ( + predicates as P, +) from tilefoundry.ir.pattern.match import written_binding, written_bindings, written_place from tilefoundry.ir.types import ComposedLayout, Layout, Mesh from tilefoundry.ir.types.dim import DimVar @@ -145,10 +148,11 @@ def required_scope(self) -> Mesh: def scope_pattern(cls) -> MeshPattern: topology, = cls.scope.topologies size = topology.size - bare = LayoutPattern.from_layout(cls.scope.layout, per_mode=True) + per_mode = (P.Forward(per_mode=True), P.Injective(per_mode=True)) + bare = LayoutPattern.from_layout(cls.scope.layout, predicates=per_mode) sliced = ComposedLayoutPattern( offset=CapturePattern("p0", MultipleOfPattern(size)), - outer=LayoutPattern.from_layout(cls.scope.layout, per_mode=True), + outer=LayoutPattern.from_layout(cls.scope.layout, predicates=per_mode), ) return MeshPattern((topology.name,), OrPattern(sliced, bare)) diff --git a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py index e16ebe72..03179db4 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py @@ -11,6 +11,9 @@ TensorPattern, WildcardPattern, ) +from tilefoundry.ir.pattern import ( + predicates as P, +) from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology from tilefoundry.ir.types.storage import StorageKind as S @@ -38,7 +41,10 @@ mesh=WARP, ) -_WARP_LAYOUT = LayoutPattern.from_layout(WARP.layout, per_mode=True) +_WARP_LAYOUT = LayoutPattern.from_layout( + WARP.layout, + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), +) _WARP_PATTERN = MeshPattern( ("thread",), OrPattern( diff --git a/src/tilefoundry/ir/tir/cuda/nn/wgmma.py b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py index 27a2f4a2..cf662dcf 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/wgmma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/wgmma.py @@ -22,6 +22,9 @@ TensorPattern, WildcardPattern, ) +from tilefoundry.ir.pattern import ( + predicates as P, +) from tilefoundry.ir.pattern.match import is_symbolic from tilefoundry.ir.types import Broadcast, DType, Layout, Mesh, Split, Topology from tilefoundry.ir.types.dim import DimVar @@ -148,7 +151,10 @@ def Fragment(rows: int, cols) -> LayoutPattern: SHARED_BY_ALL = (Broadcast(), Broadcast(), Broadcast()) HELD_PER_THREAD = (Split(2), Split(0), Split(4)) -_WARPGROUP_LAYOUT = LayoutPattern.from_layout(WARPGROUP.layout, per_mode=True) +_WARPGROUP_LAYOUT = LayoutPattern.from_layout( + WARPGROUP.layout, + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), +) _WARPGROUP_PATTERN = MeshPattern( ("thread",), OrPattern( diff --git a/tests/ir/pattern/test_mesh_pattern.py b/tests/ir/pattern/test_mesh_pattern.py index 23cc95be..2197c390 100644 --- a/tests/ir/pattern/test_mesh_pattern.py +++ b/tests/ir/pattern/test_mesh_pattern.py @@ -7,16 +7,25 @@ LayoutPattern, MeshPattern, MultipleOfPattern, + OneOfPattern, WildcardPattern, ) +from tilefoundry.ir.pattern import ( + predicates as P, +) +from tilefoundry.ir.types import ComposedLayout, DType, Layout def test_mesh_pattern_matches_the_levels_it_names(): - with pytest.raises(ValueError, match="per_mode=True"): + with pytest.raises(ValueError, match="per-mode predicates"): MeshPattern( ("thread",), ComposedLayoutPattern( - outer=LayoutPattern(((128,),), ((1,),)) + outer=LayoutPattern( + ((128,),), + ((1,),), + predicates=(P.Forward(),), + ) ), ) @@ -24,7 +33,11 @@ def test_mesh_pattern_matches_the_levels_it_names(): ("thread",), ComposedLayoutPattern( offset=CapturePattern("p0", MultipleOfPattern(128)), - outer=LayoutPattern(((128,),), ((1,),), per_mode=True), + outer=LayoutPattern( + ((128,),), + ((1,),), + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), + ), ), ) assert warpgroup.match(CT[1:3, 128:256]).captures["p0"] == 128 @@ -33,9 +46,28 @@ def test_mesh_pattern_matches_the_levels_it_names(): ("cta", "thread"), ComposedLayoutPattern( offset=WildcardPattern(), - outer=LayoutPattern(((2,), (128,)), ((1,), (1,)), per_mode=True), + outer=LayoutPattern( + ((2,), (128,)), + ((1,), (1,)), + predicates=(P.Forward(per_mode=True), P.Injective(per_mode=True)), + ), ), ) assert both.match(CT[1:3, 128:256]) is not None assert both.match(CT[0:1, 128:256]) is None assert warpgroup.match(CTA) is None + + +def test_layout_predicates_read_through_composition(): + pattern = LayoutPattern( + predicates=( + P.WholeVectors( + CapturePattern("width", OneOfPattern((4, 8, 16))), + "dtype0", + (4, 8, 16), + ), + ) + ) + subject = ComposedLayout(None, 0, Layout(((128, 4),), ((4, 1),))) + + assert pattern.match(subject, {"dtype0": DType.f32}).captures["width"] == 16 From 361597c78f901bace640fbf74a9756856d6f061f Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 16:28:41 +0800 Subject: [PATCH 16/19] fix(pattern): preserve shard layout predicate rules --- src/tilefoundry/ir/pattern/constraint.py | 7 ++++++- src/tilefoundry/ir/pattern/pattern.py | 3 +++ src/tilefoundry/ir/pattern/predicates.py | 3 ++- 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/tilefoundry/ir/pattern/constraint.py b/src/tilefoundry/ir/pattern/constraint.py index 62ca3317..fba56b96 100644 --- a/src/tilefoundry/ir/pattern/constraint.py +++ b/src/tilefoundry/ir/pattern/constraint.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from math import prod -from tilefoundry.ir.types import ComposedLayout, Layout, ShardLayout, Swizzle +from tilefoundry.ir.types import Broadcast, ComposedLayout, Layout, ShardLayout, Swizzle from tilefoundry.ir.types.layout import flatten from tilefoundry.ir.types.tensor_type import TensorType @@ -80,6 +80,11 @@ def pair(self, operands: dict): @staticmethod def reading(tensor: TensorType, arrangement=None): layout = tensor.layout if arrangement is None else arrangement + if isinstance(layout, ShardLayout): + if not all(isinstance(attr, Broadcast) for attr in layout.attrs): + layout = None + else: + layout = layout.layout if isinstance(layout, ComposedLayout): if layout.inner is not None and not isinstance(layout.inner, Swizzle): layout = None diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index 75c71af2..2c6ed699 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -6,6 +6,7 @@ from typing import Any from tilefoundry.ir.types import ( + Broadcast, ComposedLayout, Layout, Mesh, @@ -70,6 +71,8 @@ class Predicate(Pattern): def arrangement(subject) -> Layout | None: """Read the static strided layout beneath shard and composition wrappers.""" if isinstance(subject, ShardLayout): + if not all(isinstance(attr, Broadcast) for attr in subject.attrs): + return None subject = subject.layout if isinstance(subject, ComposedLayout): if subject.inner is not None and not isinstance(subject.inner, Swizzle): diff --git a/src/tilefoundry/ir/pattern/predicates.py b/src/tilefoundry/ir/pattern/predicates.py index bfef4f48..5945f0da 100644 --- a/src/tilefoundry/ir/pattern/predicates.py +++ b/src/tilefoundry/ir/pattern/predicates.py @@ -133,7 +133,8 @@ class WholeVectors(Predicate): def available_widths(self, subject, captures) -> tuple[int, ...]: bits = getattr(dict(captures or {}).get(self.dtype), "bit_width", None) - return () if type(bits) is not int else _vector_widths(subject, bits, self.widths) + layout = subject.layout if isinstance(subject, ShardLayout) else subject + return () if type(bits) is not int else _vector_widths(layout, bits, self.widths) def match(self, subject, captures=None): held = dict(captures or {}) From f5b13a5a278a1cd77104e92d2df8b867eae64d9f Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 18:11:14 +0800 Subject: [PATCH 17/19] refactor(tir): drive verification from declarations --- docs/spec/tir.md | 16 ++++++---- src/tilefoundry/ir/tir/async_copy.py | 3 -- src/tilefoundry/ir/tir/cast.py | 6 +--- .../ir/tir/cuda/memory/copy_async_tensor.py | 8 ----- src/tilefoundry/ir/tir/cuda/nn/mma.py | 12 +------- src/tilefoundry/ir/tir/verify.py | 19 ++++++++---- .../schedule/tir/sm80_mma_ldmatrix.py | 7 ++++- .../tir/wgmma_cast_between_schedules.py | 30 ++++++++++++++++++- 8 files changed, 60 insertions(+), 41 deletions(-) diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 38761cf2..1924bb8a 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -187,8 +187,10 @@ class PrimFunction(Stmt): `PrimFunction` of that name in the enclosing `Module`, `args` length MUST match the resolved callee's `params`, and the `SymbolRef.type` MUST equal the resolved callee's `CallableType`. When `callable` is - an `Op`, the per-Op verifier registered via - `@register_verify_stmt(Op)` runs. + an `Op`, every input operand MUST match its `ParamDef.pattern` and every + declared `between` relation MUST hold. A per-Op verifier registered via + `@register_verify_stmt(Op)` MAY impose additional rules that the declaration + does not express. ### 1.4 `Evaluate` @@ -206,16 +208,18 @@ The `callable` is one of: - an effect-form `Op` (e.g. `tir.memory.Copy`, `tir.cuda.nn.Mma`, `tir.tensor.Reduce`, `tir.Launch` [§2.3](#23-tir-ops)). `args` are - the Op's operands in `ParamDef` order; the per-Op verifier - registered via `@register_verify_stmt(Op)` runs. + the Op's operands in `ParamDef` order. Verification runs its optional + per-Op verifier, then the declared `between` relations and operand patterns; + a context-dependent operand pattern MAY resolve itself against the callable + before matching. - a `SymbolRef` ([§2.1](#21-symbolref)) — a reference to a callee `PrimFunction` in the enclosing `Module`. `args` follow the callee's parameter order, the final `output_count` positions binding output buffers; the callee is resolved uniquely at module level ([§1.3](#13-primfunction)). -The per-Op verify / codegen handlers are keyed by `Op` type and receive the Op -together with `args`; an `Op` callable carries no result, so its +Per-Op verify and codegen handlers, when present, are keyed by `Op` type and +receive the Op together with `args`; an `Op` callable carries no result, so its `Call` form is unit-typed. The value-producing counterpart is the `Call(Op, args)` Expr diff --git a/src/tilefoundry/ir/tir/async_copy.py b/src/tilefoundry/ir/tir/async_copy.py index 5b3f712e..9d43b12a 100644 --- a/src/tilefoundry/ir/tir/async_copy.py +++ b/src/tilefoundry/ir/tir/async_copy.py @@ -6,7 +6,6 @@ from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op from tilefoundry.ir.pattern import DistinctConstraint, SameModesConstraint, utils -from tilefoundry.ir.tir.verify import verify_between, verify_operands from tilefoundry.ir.types import Layout, UnitType from tilefoundry.ir.types.storage import StorageKind as S from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -57,8 +56,6 @@ def verify_copy_async(call: "Call", ctx: "VerifyContext") -> None: ctx.error(call, f"CopyAsync source must be gmem, got {src.storage}") if src.dtype != dst.dtype: ctx.error(call, f"CopyAsync dtype mismatch: {src.dtype} vs {dst.dtype}") - verify_between(call, ctx) - verify_operands(call, ctx, "copy_async") @register_op(dialect="T", category="async", name="cp_async_commit") diff --git a/src/tilefoundry/ir/tir/cast.py b/src/tilefoundry/ir/tir/cast.py index 622986c1..541bcc27 100644 --- a/src/tilefoundry/ir/tir/cast.py +++ b/src/tilefoundry/ir/tir/cast.py @@ -11,9 +11,8 @@ TensorPattern, any_threads, ) -from tilefoundry.ir.tir.verify import verify_between from tilefoundry.ir.types import StorageKind, UnitType -from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt +from tilefoundry.visitor_registry import register_typeinfer _REGISTER = TensorPattern(storage=StorageKind.RMEM) @@ -36,7 +35,4 @@ def _(call: "Call", ctx: "TypeInferContext") -> UnitType: return UnitType() -register_verify_stmt(Cast)(verify_between) - - __all__ = ["Cast"] diff --git a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py index 39f3d883..e6d073f3 100644 --- a/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py +++ b/src/tilefoundry/ir/tir/cuda/memory/copy_async_tensor.py @@ -26,7 +26,6 @@ from tilefoundry.ir.pattern import ( predicates as P, ) -from tilefoundry.ir.tir.verify import verify_between, verify_operands from tilefoundry.ir.types import ComposedLayout, Layout, Mesh, Swizzle, UnitType from tilefoundry.ir.types.storage import StorageKind as S from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt @@ -209,13 +208,6 @@ def verify_copy_async_tensor(call: "Call", ctx: "VerifyContext") -> None: f"copy_async_tensor moves one tile: src is {tuple(src.shape)} " f"{src.dtype.name} and dst is {tuple(dst.shape)} {dst.dtype.name}", ) - moves = " and ".join(map(str, TMA_STORAGES)) - verify_between( - call, - ctx, - f"copy_async_tensor moves a tile between {moves}, one end each: ", - ) - verify_operands(call, ctx, "copy_async_tensor") __all__ = [ diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index 866e60c6..01e7fe64 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma.py @@ -16,11 +16,10 @@ from tilefoundry.ir.pattern import ( predicates as P, ) -from tilefoundry.ir.tir.verify import input_params from tilefoundry.ir.types import DType, Mesh, UnitType from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt -from .mma_atom import AtomPattern, FromAtom, MmaAtom, physical_frames_match, read_on +from .mma_atom import AtomPattern, FromAtom, MmaAtom, physical_frames_match from .sm80_mma import Mma as _Sm80Mma from .wgmma import Wgmma @@ -89,15 +88,6 @@ def verify_mma(call: "Call", ctx: "VerifyContext") -> None: """Check each operand against its atom and the active physical frame.""" op = call.target atom = op.atom - held = tuple(ctx.type_of(arg) for arg in call.args) - for param, value in zip(input_params(type(op)), held): - pattern = read_on(param.pattern, op) - if pattern.match(value) is None: - ctx.error( - call, - f"MMA {param.name} is not one {atom.reference} reads; " - f"it reads {pattern.describe()}", - ) if ctx.scope is not None and ctx.scope.module is not None: capabilities = ctx.scope.module.target.architecture.instruction_capabilities if op.capability not in capabilities: diff --git a/src/tilefoundry/ir/tir/verify.py b/src/tilefoundry/ir/tir/verify.py index 1b9cc86d..aa825ebf 100644 --- a/src/tilefoundry/ir/tir/verify.py +++ b/src/tilefoundry/ir/tir/verify.py @@ -67,18 +67,24 @@ def verify_between(call, ctx, lead: str = "") -> None: ctx.error(call, lead + rule.refused(operands)) -def verify_operands(call, ctx, label: str) -> None: +def verify_operands(call, ctx) -> None: """Hold each operand to the pattern declared for its parameter.""" + label = type(call.target)._op_schema.name for param, arg in zip(input_params(type(call.target)), call.args): if param.pattern is None: continue value = ctx.type_of(arg) - if param.pattern.match(value) is None: + pattern = ( + param.pattern.read_on(call.target) + if hasattr(param.pattern, "read_on") + else param.pattern + ) + if pattern.match(value) is None: ctx.error( call, f"{label} {param.name} is {tuple(value.shape)} " f"{value.dtype.name} storage={value.storage}: " - f"{param.pattern.refusal(value)}", + f"{pattern.refusal(value)}", ) @@ -212,12 +218,13 @@ def _walk_stmt(stmt, ctx, scope, fn, module_fn_map, bound_var_ids: set[int]): op = stmt.callable op_cls = type(op) fn_verify = verify_stmt_registry.lookup(op_cls) - if fn_verify is None: - raise VerifyError(f"no verify_stmt registered for Op {op_cls.__name__}") ctx.mesh_scope = tuple(scope) call = Call(type=UnitType(), target=op, args=stmt.args) - fn_verify(call, ctx) + if fn_verify is not None: + fn_verify(call, ctx) + verify_between(call, ctx) + verify_operands(call, ctx) for arg in stmt.args: _reject_nested_alloc_tensor(arg, at_letstmt_value=False) diff --git a/tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py b/tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py index 4e181994..04bf835a 100644 --- a/tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py +++ b/tests/fixtures/schedule/tir/sm80_mma_ldmatrix.py @@ -73,7 +73,12 @@ def gemm(a: Tensor[(16, 32), "bf16"], b: Tensor[(32, 8), "bf16"], out: Tensor[(1 outer=Layout((4, 8), (1, 4)), ), names=("d0", "d1") ) as threads_3: - T.ldmatrix(lhs_stages[(k // 16) % 2], ldmatrix) + ldmatrix_view = T.tensor_view( + T.ptr_of(ldmatrix), + layout=((2, 4 @ threads_3.d0, 2, 8 @ threads_3.d1, 2), (1, 2, 8, 16, 128)), + shape=(16, 16), + ) + T.ldmatrix(lhs_stages[(k // 16) % 2], ldmatrix_view) copy = T.alloc_tensor( tensor_type=Tensor[ (16, 8), "bf16", Layout((8, 2, 4, 2), (1, 8, 16, 64)), "rmem" diff --git a/tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py b/tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py index 264aade6..db9499f8 100644 --- a/tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py +++ b/tests/fixtures/schedule/tir/wgmma_cast_between_schedules.py @@ -28,7 +28,35 @@ def gemm( with Mesh( (Topology("thread", 256),), Layout((2, 128), (128, 1)), names=("d0", "d1") ) as scope_4: - T.cast(b_f32, value) + stage_f32 = T.alloc_tensor( + tensor_type=Tensor[ + (32, 32), + "f32", + ((2 @ scope_4.d0, 128 @ scope_4.d1, 4), (512, 4, 1)), + "rmem", + ] + ) + stage_bf16 = T.alloc_tensor( + tensor_type=Tensor[ + (32, 32), + "bf16", + ((2 @ scope_4.d0, 128 @ scope_4.d1, 4), (512, 4, 1)), + "rmem", + ] + ) + b_f32_view = T.tensor_view( + T.ptr_of(b_f32), + layout=((2 @ scope_4.d0, 128 @ scope_4.d1, 4), (512, 4, 1)), + shape=(32, 32), + ) + value_view = T.tensor_view( + T.ptr_of(scratch[0:0 + 1024]), + layout=((2 @ scope_4.d0, 128 @ scope_4.d1, 4), (512, 4, 1)), + shape=(32, 32), + ) + T.copy(b_f32_view, stage_f32) + T.cast(stage_f32, stage_bf16) + T.copy(stage_bf16, value_view) lhs_stages = (T.tensor_view(1024, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16)), T.tensor_view(2048, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((8, 8), (2, 8)), ((128, 8), (64, 1))), shape=(64, 16))) rhs_stages = (T.tensor_view(0, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32)), T.tensor_view(512, dtype='bf16', storage=StorageKind.SMEM, layout=Layout(((2, 8), (4, 8)), ((64, 8), (128, 1))), shape=(16, 32))) with Mesh( From 0eece7d180791a499ebf924f5e0893f4a88ef481 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 18:16:36 +0800 Subject: [PATCH 18/19] fix(tir): retain undeclared op verification guard --- docs/spec/tir.md | 3 ++- src/tilefoundry/ir/tir/verify.py | 11 +++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 1924bb8a..8863f579 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -190,7 +190,8 @@ class PrimFunction(Stmt): an `Op`, every input operand MUST match its `ParamDef.pattern` and every declared `between` relation MUST hold. A per-Op verifier registered via `@register_verify_stmt(Op)` MAY impose additional rules that the declaration - does not express. + does not express. An effect Op with neither an input pattern nor a `between` + relation MUST register such a verifier; an Op stating no contract is rejected. ### 1.4 `Evaluate` diff --git a/src/tilefoundry/ir/tir/verify.py b/src/tilefoundry/ir/tir/verify.py index aa825ebf..f805f7d0 100644 --- a/src/tilefoundry/ir/tir/verify.py +++ b/src/tilefoundry/ir/tir/verify.py @@ -54,7 +54,7 @@ def input_params(op_type: type) -> tuple: return tuple(param for param in op_type._op_schema.signature if param.kind == "input") -def verify_between(call, ctx, lead: str = "") -> None: +def verify_between(call, ctx) -> None: """Hold one call to the relations declared between its operands.""" op_type = type(call.target) rules = between_rules(op_type) @@ -64,7 +64,7 @@ def verify_between(call, ctx, lead: str = "") -> None: operands = dict(zip(names, (ctx.type_of(arg) for arg in call.args))) for rule in rules: if not rule.holds(operands): - ctx.error(call, lead + rule.refused(operands)) + ctx.error(call, rule.refused(operands)) def verify_operands(call, ctx) -> None: @@ -218,6 +218,13 @@ def _walk_stmt(stmt, ctx, scope, fn, module_fn_map, bound_var_ids: set[int]): op = stmt.callable op_cls = type(op) fn_verify = verify_stmt_registry.lookup(op_cls) + if fn_verify is None and not ( + any(param.pattern is not None for param in input_params(op_cls)) + or between_rules(op_cls) + ): + raise VerifyError( + f"Op {op_cls.__name__} states no verifier and no declaration" + ) ctx.mesh_scope = tuple(scope) call = Call(type=UnitType(), target=op, args=stmt.args) From 457de150f7c80618406cd5f17b2f0e9eda5ed68f Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Sat, 26 Sep 2026 19:40:21 +0800 Subject: [PATCH 19/19] refactor(tir): share captures across operands --- docs/spec/core-ir.md | 5 ++++- src/tilefoundry/ir/tir/verify.py | 8 ++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 7a5cd910..00bc2d54 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -655,7 +655,10 @@ Two consumer surfaces: `AndPattern(parts)` (conjunction). Two singletons are exported as convenience: `Scalar = ScalarPattern()` and `Tensor = TensorPattern()`. A tensor rank is stated by giving `shape` that many positions; wildcard - positions constrain only the sequence length. + positions constrain only the sequence length. During effect-Op verification, + input patterns match in `ParamDef` order against one shared capture + environment, so a later operand can require a value captured by an earlier + operand. - **Specialization dispatch** — patterns appearing in `hir.Function.specializations` ([hir.md §1.1](./hir.md#11-function)) and `tir.PrimFunction.specializations` describe which runtime diff --git a/src/tilefoundry/ir/tir/verify.py b/src/tilefoundry/ir/tir/verify.py index f805f7d0..00b8314c 100644 --- a/src/tilefoundry/ir/tir/verify.py +++ b/src/tilefoundry/ir/tir/verify.py @@ -70,6 +70,7 @@ def verify_between(call, ctx) -> None: def verify_operands(call, ctx) -> None: """Hold each operand to the pattern declared for its parameter.""" label = type(call.target)._op_schema.name + held = {} for param, arg in zip(input_params(type(call.target)), call.args): if param.pattern is None: continue @@ -79,13 +80,16 @@ def verify_operands(call, ctx) -> None: if hasattr(param.pattern, "read_on") else param.pattern ) - if pattern.match(value) is None: + found = pattern.match(value, held) + if found is None: ctx.error( call, f"{label} {param.name} is {tuple(value.shape)} " f"{value.dtype.name} storage={value.storage}: " - f"{pattern.refusal(value)}", + f"{pattern.refusal(value, held)}", ) + continue + held = found.captures def verify_prim_function(