feat(cuda): declare CUDA MMA atoms and async copy instructions

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. |
@@ -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),
@@ -135,9 +144,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}/<name>.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

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`, `Mma`, `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

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)).
@@ -613,33 +613,52 @@ 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.
-- `match.py` owns matches, captures, symbolic resolution, layout-frame reading,
-  and the shared description helpers.
+- `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 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.
+- `utils.py` owns specialization naming and dimension lookup.
+
+`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 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
+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. 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

docs/spec/runtime.md

@@ -1194,15 +1194,16 @@ 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 |
+| `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
@@ -1404,32 +1405,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.
@@ -1556,17 +1539,18 @@ __device__ inline void sync(Mesh<TMesh, Topos...> 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`
 
-<!-- generated: ops-tma -->
+<!-- generated: ops-copy-async-bulk -->
 ```cpp
-// include/tilefoundry/runtime/cuda/ops/tma.cuh
+// include/tilefoundry/runtime/cuda/ops/copy_async_bulk.cuh
 template <class Src, class Dst>
-__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);
 ```
 <!-- /generated -->
 
-**`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
@@ -1616,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`
+
+<!-- generated: ops-ldmatrix -->
+```cpp
+// include/tilefoundry/runtime/cuda/ops/ldmatrix.cuh
+template <class Src, class Dst>
+__device__ inline void ldmatrix(Src const &src, Dst &dst);
+```
+<!-- /generated -->
+
+**`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