feat(rocke): IR dtypes in TileMma, a storage type under every Fragment, and MMA issue from the measured backend coverage

dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx942.md

@@ -51,6 +51,14 @@ enforced by `instances/common/gemm_universal.py::is_valid_spec` /
 **f16 / bf16** atoms `16x16x32` and `32x32x16`. Wide-K on CDNA3 exists
 for **fp8 / bf8 only** — there is no dense wide-K f16/bf16 path.
 
+**xf32 (TF32) MFMA: why pre-convert?** The xf32 instructions
+(`16x16x8`, `32x32x4`) take f32 registers and use only their TF32 bits:
+the low 13 mantissa bits of each input are dropped (truncation). An f32
+operand can therefore feed the instruction directly, with no conversion
+(the tiling driver accepts an F32 fragment for a TF32 `TileMma`). Convert the operand to TF32 first only when the inputs must be
+rounded to nearest even instead of truncated; the converted TF32 fragment
+then feeds the instruction unchanged.
+
 **Not present on CDNA3 at all:** block-scale / microscaling (MX) MFMA and
 FP6 / FP4 — these are CDNA4 (gfx950) additions
 ([ROCm Matrix-Cores blog](https://rocm.blogs.amd.com/software-tools-optimization/matrix-cores-cdna/README.html)).

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md

@@ -222,7 +222,7 @@ cooperative map; `classify_epilogue(...)` → the branch. Worked example:
 
 ## Tool catalog (load-bearing knobs; full list in `visualization_api_surface.md`)
 
-- **`MmaTee`** — `from_mma(mma, **overrides)` (canonical refs + dtypes from the atom) OR explicit args. ✗
+- **`MmaTee`** — `from_mma(mma, **overrides)` (canonical refs from the atom + per-operand dtypes from the TileMma: its A, B and result dtypes) OR explicit args. ✗
   `from_mma` with no `a_enc`/`b_enc` renders CANONICAL (it sources canonical from the `canonical_layouts` helper,
   NOT `mma.a_layout`) — for interleaved you MUST pass explicit `a_enc`/`b_enc` and VERIFY the encoding is
   `!= canonical_layouts(mma.traits, mma.subtiles)` + sound (`mma.a_layout` is the CONFIGURED style now, NOT a
@@ -280,7 +280,8 @@ cooperative map; `classify_epilogue(...)` → the branch. Worked example:
 - `mma_operand_layout_sound(layout, canon, role=)` → `ok` / `error`. `mma_operand_repair_hint(layout, canon)`
   → `ok` / `warning` (an in-register reorder makes-it-so — name the permutation) / `error` (cross-lane or
   none). `mma_pair_compatible(a, b, a_canon=, b_canon=)` → both sound AND K-dists match; drives the tee
-  banner. `mma_pair_k_aligned(a, b)` → `ok`/`warning`/`error` (pairwise K-match, per atom).
+  banner. `mma_pair_k_aligned(a, b, a_canon=, b_canon=)` → `ok`/`warning`/`error` (pairwise K-match, per
+  atom, at the positions the canonicals couple).
   `classify_transform(src, tgt)` → `reorder` (cheap, dtype-graded) vs `cross_lane` (expensive/deferred).
 - Tee banner red ERROR / orange WARNING = the pair diagnostic — say what's wrong, which rule, cheapest remedy
   (free symmetry [reposition/source-swap] → dword reorder → sub-dword → cross-lane last).

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/README.md

@@ -6,12 +6,13 @@ tiled GEMM/MMA kernel -- or any custom thread/register distribution -- without h
 bound target so authorship stays target-agnostic.
 
 ```python
+from rocke.core.ir import F16, F32
 from rocke.helpers.tiling import (
     TileMma, make_tensor_desc, make_window, make_fragment,
     fill_fragment, load_fragment, store_fragment,
 )
 
-mma = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a")  # IR-free resolve
+mma = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a")  # IR-free resolve
 acc = make_fragment(mma.c_desc, F32)
 fill_fragment(b, acc, 0)
 for k0 in range(0, K_LEN, 16):

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/geometry.py

@@ -18,6 +18,8 @@
 
 from typing import Any
 
+from ....core.dtypes import dtype_info
+
 _BINARY = {
     "arith.add": lambda a, b: a + b,
     "arith.sub": lambda a, b: a - b,
@@ -112,7 +114,19 @@ def _arch_wave(pipeline: Any) -> tuple[str, int]:
     return pipeline.arch, pipeline.wave_size
 
 
-def _elem_bytes(dtype_name: str) -> int:
-    from ..emit import _BYTE_WIDTH
+def _elem_bits(dtype_name: str) -> int:
+    """Bits of one element of `dtype_name` (``core.dtypes.dtype_info``) -- the width home of the
+    views."""
+    return dtype_info(dtype_name).encoded_bits
+
 
-    return _BYTE_WIDTH[dtype_name]
+def _elem_bytes(dtype_name: str) -> int:
+    """Bytes of one element of `dtype_name`, for the byte-addressed views (the LDS bank model). A
+    sub-byte element has no byte address of its own and is refused."""
+    bits = _elem_bits(dtype_name)
+    if bits % 8:
+        raise NotImplementedError(
+            f"{dtype_name} elements are {bits} bits, not a whole number of bytes -- "
+            "byte-addressed views (LDS bank model) of packed data -- TODO: not implemented yet"
+        )
+    return bits // 8

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/soundness.py

@@ -18,10 +18,11 @@ class MmaSoundnessError(RuntimeError):
 
 def verify_mma_soundness(pipeline: Any) -> int:
     """Gate 2: every recorded MMA's CONSUMED operand encodings must be sound MMA operands against the
-    canonical machine (``mma_operand_layout_sound``) AND share a K-distribution (``mma_pair_k_aligned``) -- bundled
-    as ``mma_pair_compatible``. This is the operand-correctness the addressing round-trip is BLIND to: a
+    canonical machine (``mma_operand_layout_sound``) AND hold the same K at every pair of positions the
+    instruction couples (``mma_pair_k_aligned``; slot for slot on a same-format instruction, through the
+    canonical layouts on a mixed-width one) -- bundled as ``mma_pair_compatible``. This is the operand-correctness the addressing round-trip is BLIND to: a
     scrambled/duplicated operand K passes the address gate but is caught here. ``a_canon``/``b_canon`` are
-    the trusted canonical layouts from the MMA definition; ``a_enc``/``b_enc`` are the kernel's own
+    the trusted canonical layouts from the MMA definition (they also define the A/B coupling); ``a_enc``/``b_enc`` are the kernel's own
     (interleaved) operands. Returns the count verified; raises on the first unsound MMA. Correctness SOT:
     ``docs/mma_is_machinery.md`` (the three-condition sound MAC).
     """

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/api_proposals/PROPOSAL_c_transpose_crossed_ab.md

@@ -17,7 +17,7 @@ hand-plumb three coupled changes and get them all consistent:
 
 ```python
 # today (awkward) — the "ab_swap" crossed path, hand-wired in build_crc_gemm:
-mma = TileMma((warp_n, warp_m, tile_k), a="f16", b="f16", c="f32", target=arch)  # 1. swap M<->N in the shape
+mma = TileMma((warp_n, warp_m, tile_k), a=F16, b=F16, c=F32, target=arch)  # 1. swap M<->N in the shape
 ...
 acc = mma(b_frag, a_frag, acc)          # 2. feed B into the A-slot, A into the B-slot (M<->N relabel)
 ...
@@ -39,7 +39,7 @@ transform (register-identity), not cross-lane.
 
 ```python
 # proposed API
-mma = TileMma((M, N, K), a="f16", b="f16", c="f32", target=arch, c_transpose=True)
+mma = TileMma((M, N, K), a=F16, b=F16, c=F32, target=arch, c_transpose=True)
 acc = mma(a_frag, b_frag, acc)                 # normal call; object handles the crossed feed
 store_fragment(b, c_ptr, make_window(mma.c_desc, (m0, n0)), acc, lane)   # c_desc already (N,M)-oriented
 ```

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/mma_is_machinery.md

@@ -42,7 +42,25 @@ Everything else follows from that one physical fact. Because the atom *physicall
 
 1. **Atom-fixed accumulation.** Which physical slots sum into each output is fixed by the atom — not a choice.
 2. **One $M$ per output on A, one $N$ per output on B.** Every A-slot feeding a given output carries a single fixed $M$; every B-slot a single fixed $N$. If $M$ drifts across those slots, the accumulator holds a sum whose row wanders term to term — not $C_{ij}$ for any $(i,j)$.
-3. **$K$-match, position for position.** $A.K\text{-dist} = B.K\text{-dist}$: $A$ and $B$ agree on which contraction index $K$ sits in each paired slot, and $K$ is well-formed (each index present once).
+3. **$K$-match at coupled positions.** $A$ and $B$ agree on which contraction index $K$ sits at each pair of positions the atom multiplies together, and $K$ is well-formed (each index present once). The atom's canonical A and B layouts define that coupling: an A position and a B position are coupled when they hold the same hardware $K$. For every same-format atom this is slot $s$ with slot $s$, so the rule reads $A.K\text{-dist} = B.K\text{-dist}$. A mixed-width `f8f6f4` atom is the exception: its 8-bit operand holds $K$ in two runs of 16 per lane, its 6/4-bit operand in one run of 32, so the coupled positions differ:
+
+   ```
+   gfx950 16x16x128, fp8 A x fp4 B, lane 0
+     A slots 0-15 = K 0-15, slots 16-31 = K 64-79     (two runs of 16)
+     B slots 0-31 = K 0-31                            (one run of 32)
+     A (lane 0, slot 16) holds K 64  <->  B (lane 32, slot 0) holds K 64
+   ```
+
+   Comparing such a pair slot for slot is wrong both ways: it rejects the canonical fp8 A x fp4 B pair, and it accepts an fp4 B loaded with fp8's map, which multiplies the wrong $K$.
+
+   An atom whose lane holds several $M$ rows of A (`a_m_register_blocks > 1`) multiplies each row with the same B registers, so every $M$ block of A is coupled with B slot for slot:
+
+   ```
+   gfx1250 32x16x128 fp4, lane 0
+     A slots 0-63 = (M 0, K 0-31 and 64-95), slots 64-127 = (M 16, the same K)
+     B slots 0-63 = (N 0, K 0-31 and 64-95)
+     A (lane 0, slot 64) holds K 0  <->  B (lane 0, slot 0) holds K 0
+   ```
 
 > The slots feeding one accumulator are fixed by silicon; a sound MAC labels them with **one $M$, one $N$, and a matched $K$** — nothing else is constrained.
 
@@ -89,18 +107,18 @@ Because A's $M$-placement and B's $N$-placement are free, *they* fix C's native
 Two checks hide under "is this layout OK?", and conflating them is what sends people reaching for canonical:
 
 - **Per-operand soundness** (conditions 1–2) is checkable on *one* fragment: does every slot feeding an output carry a single $M$ (A) or $N$ (B)? It asks nothing about the partner, and nothing about canonical.
-- **Pairwise acceptability** (condition 3) is a property of the **pair**, never of one fragment: do A and B agree on $K$, slot for slot? Neither operand answers it alone — and neither is checked against a canonical reference.
+- **Pairwise acceptability** (condition 3) is a property of the **pair**, never of one fragment: do A and B agree on $K$ at the positions the atom couples (slot for slot on a same-format atom)? Neither operand answers it alone. The canonical layouts enter only to say *which* positions are coupled; neither operand's labels are compared against them.
 
 Consequences:
 
 - "Make this fragment acceptable" always means "reorder it until its $K$-row matches its **partner's**" — never "until it looks canonical."
-- `mma_pair_k_aligned(A, B)` compares the two fragments **to each other**; it never checks either against a canonical reference.
+- `mma_pair_k_aligned(A, B, a_canon=..., b_canon=...)` compares the two fragments **to each other** at coupled positions. The canonicals are required because they define the coupling; neither fragment's labels are checked against them.
 
 ## The traps agents fall into (name them, then don't)
 
 Wrong-turns that come from pattern-matching on training data instead of reasoning from the machinery:
 
-1. **"The MFMA has a fixed layout, so the fragment must be canonical."** — No. The hardware reads slots and multiply-accumulates; the layout is bookkeeping *we* maintain. $K$-pairing is the only **cross-operand** constraint — not canonical resemblance.
+1. **"The MFMA has a fixed layout, so the fragment must be canonical."** — No. The hardware reads slots and multiply-accumulates the positions its canonical coupling pairs (slot $s$ with slot $s$, except on mixed-width `f8f6f4` atoms); the labels are bookkeeping *we* maintain. $K$-pairing at coupled positions is the only **cross-operand** constraint — not canonical resemblance.
 2. **"A lane can only own these $N$ strided elements, so we're capped at dwordx2."** — Reasoning from *one* canonical distribution. $M/N$ placement is free; choose a distribution where the lane owns a contiguous run and the wide store exists. "Canonical caps us" is never a valid conclusion.
 3. **"Just shuffle the canonical C into the store order."** — C's ownership is derived from A/B. The store-friendly C is *produced* by choosing A/B placement, not recovered by post-hoc shuffling a canonical accumulator (often cross-lane and expensive).
 4. **Treating `mma_pair_k_aligned` passing as "canonical-valid."** — It only means A and B agree on $K$. A hand-rolled distribution's correctness is confirmed by **bit-exact**, not by resemblance to canonical.
@@ -110,7 +128,7 @@ Wrong-turns that come from pattern-matching on training data instead of reasonin
 
 For any layout/direction question about an MMA operand or accumulator:
 
-1. **Is this a correctness question?** Correctness = the sound MAC: per-operand $M$/$N$ fixed (1–2) **and** $A.K\text{-dist} = B.K\text{-dist}$ (3). If those hold, the MMA is correct — full stop.
+1. **Is this a correctness question?** Correctness = the sound MAC: per-operand $M$/$N$ fixed (1–2) **and** A and B hold the same $K$ at coupled positions (3). If those hold, the MMA is correct — full stop.
 2. **Place $M/N$ for the goal, not for canon.** Choose A's $M$-placement and B's $N$-placement for the memory behavior you want (coalesced wide load, store-friendly C, feed a downstream MMA).
 3. **Let C be derived.** C emerges at the $M/N$ you chose; shape the *store* around that (or a cheap register reorder *within* the already-derived ownership), not by reshaping a canonical C.
 4. **Verify by bit-exact.** The GPU is the arbiter. "It doesn't look canonical" is not evidence of anything.