Skip to content

Commit 5022edc

Browse files
apstenku123apstenku123
committed
docs(foldttir): TTIR canonicalize/fold MEASURE on gb10 — fold verified (367 lines, native-matching guard counts), routing NO-GO
Fold of the captured §P1 dstates TTIR replicates Triton make_ttir (extsi 102->0, cmpi 109->7, INT_MAX 66->0, andi 54->3) and matches native Triton TTIR counts. Routing the folded TTIR exposed fused-GEMM-accumulate-into-shared-carry which our serial-scalar loop-carry copies cannot lower (mma swizzle vs linear -> MAXDIFF 1.5e3). Fixed lb-name + strict-fragment bugs (tilelang 46ae6c79); the fused case now fails loud (RULE #1). UNFOLD routed stays correct (4.882812e-04, 733 ms vs native 1.16-1.20 ms). Fold ROUTING = honest NO-GO.
1 parent 4a122ef commit 5022edc

1 file changed

Lines changed: 116 additions & 0 deletions

File tree

‎docs/TTIR-CANONICALIZE-FOLD.md‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# TTIR canonicalize/fold — §P1 dstates MEASURE (gb10, sm_121, triton 3.7.0)
2+
3+
Status: **FoldTTIR DONE+VERIFIED. Fold ROUTING = NO-GO (honest, RULE #1).**
4+
5+
## What was proven
6+
7+
The captured TTIR for `_chunk_scan_bwd_dstates_kernel` was snapshotted too early
8+
in Triton's pipeline (`ASTSource.make_ir` returns PRE-optimization TTIR), so it
9+
carried the i32->i64 overflow-guard chain that native Triton folds in `make_ttir`
10+
via the MLIR canonicalizer (int-range analysis proves the guards always-true for
11+
the kernel's i32 index ranges).
12+
13+
`_fold_ttir()` (poc/triton_frontend/_test_harness/jit_to_ttir.py, commit
14+
0a01caac) replicates Triton's nvidia `make_ttir` pass list on the captured module
15+
before serialization. EXECUTED on gb10:
16+
17+
| TTIR | lines | extsi | cmpi | 2147483647 | andi |
18+
|---------------------|------:|------:|-----:|-----------:|-----:|
19+
| UNFOLD (pre-fold) | 895 | 102 | 109 | 66 | 54 |
20+
| **FOLD (from HEAD)**| 367 | **0** |**7** | **0** |**3** |
21+
| native Triton cache | 299*| 0 | 7 | 0 | 3 |
22+
23+
The folded TTIR matches native Triton's TTIR counts exactly (7 remaining cmpi are
24+
legit `slt` bounds masks; zero overflow guards). The fold is a real semantic-
25+
preserving MLIR-pass equivalence, NOT a dropped check.
26+
27+
(*native cache .ttir is normalized/shorter; the guard COUNTS match, which is the
28+
input-matching invariant the diff proves.)
29+
30+
## MEASURE result — routing the folded TTIR
31+
32+
Routed both TTIRs through the frontend with `TL_FORCE_CP_ASYNC=1` (cp.async live),
33+
PtrAnalysis C++ shim ON (PYTHONPATH=poc/triton_frontend/_cxx/build), executed
34+
cubins, CUDA-event timing N=60/rep x4 reps median-of-medians, parity vs native.
35+
36+
| Routed kernel | for-loops (.cu) | LDGSTS | HMMA | spill STL+LDL | EXEC ms | parity MAXDIFF |
37+
|---------------|----------------:|-------:|-----:|--------------:|--------:|----------------|
38+
| UNFOLD | 84 | 16 | 32 | **272** | **733.29** | **4.882812e-04 PASS** |
39+
| FOLD (raw) | 49 | 16 | 16 | 192 | 638.51 | **1.515449e+03 FAIL** |
40+
| native | — | — | — | — | **1.16–1.20** | reference |
41+
42+
- The fold DID reduce the routed kernel (for-loops 84->49, spills 272->192,
43+
HMMA 32->16, cu_len 22564->16104) and was ~13% faster wall-clock (733->638 ms).
44+
- **But the folded routing is NUMERICALLY WRONG (MAXDIFF 1.5e3, ~half nonzero).**
45+
Per RULE #1 a faster-but-wrong path is forbidden; it is now made to FAIL LOUD.
46+
47+
## Root cause of the fold-routing failure (verifiable)
48+
49+
Triton's canonicalizer FUSES `acc = acc + dot(...)` so the GEMM accumulates
50+
DIRECTLY into the loop-carried `acc` tile (`tensor<64x64xf32>`). The unfolded TTIR
51+
keeps the GEMM writing a FRESH fragment (`dot_c_logical`) and does the accumulation
52+
as a SEPARATE add into a linearly-copied shared carry — which our frontend lowers
53+
correctly.
54+
55+
For the fused (folded) form, the loop-carried accumulator surfaces as a
56+
`LazyTileExpr` and is allocated as a SHARED `carry_tile` (control.py). On CUDA the
57+
tensor-core MMA store layout (`make_mma_store_layout`) requires C in a swizzled
58+
`local.fragment`. Two real frontend bugs were exposed and the third is structural:
59+
60+
1. **lb-name collision (FIXED — control.py `_bound_name`).** The loop-carry
61+
`lb` was stored as a positional SSA name (`%0`). The canonicalizer/CSE renames
62+
the K-loop lower bound to a named constant (`%c0_i32`) and reassigns `%0` to an
63+
unrelated `tt.splat` tile, so `lb` resolved to a float32 tile and aborted with
64+
`Cast(... ffi.OpaquePyObject)`. Fix: bind lb/step by their CONSTANT VALUE
65+
(numbering-independent), identical semantics. Also fixes the standalone matmul
66+
fold-route which hit the same bug.
67+
68+
2. **shared-as-fragment over-accept (FIXED — reduction.py).**
69+
`_is_fragment_scope` lumped `shared` in with fragment scopes (correct for Metal),
70+
so a SHARED carry passed straight to the CUDA MMA as C and aborted at
71+
LayoutInference (`carry_tile must be a fragment, but got shared`). Added
72+
`_is_strict_fragment_scope` and gated the CUDA-fragment copy on it.
73+
74+
3. **swizzle-vs-linear carry round-trip (STRUCTURAL — not safely fixable here).**
75+
With (1)+(2) the fold COMPILES, but the loop-carry snapshot/commit
76+
(`_append_loop_carry_copies` -> `_copy_buffer_stmt`) and the C-seed copy use
77+
SERIAL SCALAR element copies that index the tile LINEARLY. The MMA fragment is
78+
SWIZZLED. So each K-trip reads the carry in mma-swizzle layout but the carry is
79+
written/stored linearly -> the accumulator is silently corrupted
80+
(MAXDIFF ~1.5e3). Correctly lowering the fused form needs a layout-aware
81+
`T.copy` (or a swizzle-correct fragment-resident carry across the K-loop),
82+
which the current frontend does not emit.
83+
84+
## Decision (RULE #1)
85+
86+
- Keep fixes (1) and (2): real correctness improvements, no regression to the
87+
proven UNFOLD path (still 4.882812e-04, 733 ms) or the unit suite (the only
88+
reduction/control test failures are PRE-EXISTING `inspect.getsource` OSErrors,
89+
identical with edits reverted).
90+
- The fused-GEMM-accumulate-into-shared-carry case now RAISES a clear `EmitError`
91+
instead of silently emitting 1.5e3-garbage. **Fold routing = NO-GO** until a
92+
layout-aware loop-carried-fragment lowering lands.
93+
94+
## Remaining gap to native
95+
96+
UNFOLD routed §P1 = 733 ms vs native 1.16–1.20 ms (~611x). The overflow-guard
97+
prologue is NOT the dominant cost on the CURRENT HEAD (the cp.async work already
98+
removed the 64-serial-thread prologue; routed UNFOLD has threadIdx_for=0, 272
99+
spills). The fold's TTIR is correct and matches native, but our frontend cannot
100+
yet exploit the fused accumulate form, so it does not (yet) move §P1 toward native.
101+
The next lever is the layout-aware fragment-resident loop-carry (issue #3 above),
102+
not further TTIR folding.
103+
104+
## Reproduce (gb10, GPU mutex)
105+
106+
```
107+
ssh gb10
108+
source /home/dave/cppmega-venv/bin/activate
109+
cd /home/dave/source/tilelang && git checkout <HEAD>
110+
export PYTHONPATH=/home/dave/source/tilelang/poc/triton_frontend/_cxx/build:$PYTHONPATH
111+
# 1) fold the captured TTIR (BK=32, matching the snapshot):
112+
python /tmp/cap_folded_bk32.py # -> extsi=0 cmpi=7 2147483647=0 andi=3 lines=367
113+
# 2) build/measure UNFOLD (PASS) and FOLD (now fails loud):
114+
python /tmp/em_fold_measure.py unfold # PARITY 4.882812e-04, 733 ms
115+
python /tmp/em_fold_one.py fold # EmitError: CUDA MMA accumulator C is a SHARED loop-carried tile ... (RULE #1)
116+
```

0 commit comments

Comments
 (0)