Skip to content

Commit b12bfca

Browse files
apstenku123apstenku123
committed
docs: MEASURED §P1 backend-agnostic guard fold on gb10 + Metal-portability
§P1 dstates EXECUTED on NVIDIA GB10 (sm_121): fold+cp.async drops 3140.2->745.1ms (4.21x), spills 2494->272, .cu guard sites 14->0, parity 4.882812e-04. Fold drops ALL 2147483647 guards to 0 across dc/db/dx/ddAcs CUDA (mma intact). SAME folded prim lowers to valid MSL with guard folded out (Metal-portable, codegen-only). Honest gap: residual 745 vs native 1.19ms is GEMM tiling, not the guard fold. tilelang HEAD c35fea0a.
1 parent 5022edc commit b12bfca

1 file changed

Lines changed: 128 additions & 0 deletions

File tree

‎docs/FOLD-GUARDS-OUR-LAYER.md‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
# Backend-agnostic i32-overflow guard fold — §P1 dstates MEASURED (gb10 CUDA, EXECUTED)
2+
3+
HEAD `46ae6c79de8c1319194a7a0d1cced59228dfec53` (tilelang, branch
4+
`merge/upstream-codegen-reorg`). All CUDA numbers below are EXECUTED on the
5+
NVIDIA GB10 (sm_121, CUDA 13.2) under the gb10 GPU mutex. RULE#1: bit-correct
6+
EXECUTED results only; the Metal check is codegen-only (no Apple-GPU dispatch).
7+
8+
## What the fold is
9+
10+
Triton wraps every i32->i64 address expression in an always-true overflow guard
11+
pair (`addr <= 2147483647` / `-2147483648 <= addr`). For the Tri-Dao chunk shapes
12+
(seqlen 4096, nheads 112, headdim/dstate 64 — all << 2^31) these are statically
13+
TRUE. Un-folded, OUR walker materialized them as serial scalar loops + spills.
14+
15+
The fold lives in OUR backend-agnostic layer (the TTIR->TIR walker), NOT in
16+
libtriton / the NVIDIA pipeline:
17+
18+
- `poc/triton_frontend/op_emitters/arith.py`
19+
- `_is_int32_overflow_guard` (line ~1001) recognizes the `sle 2147483647` /
20+
`sge -2147483648` always-true bound pairs.
21+
- `_emit_cmpi` (line ~1065) folds them to a constant-true tile.
22+
- `_emit_andi` (line ~641) drops `guard_true & x -> x`.
23+
- Gated by `ctx.routed_triton_prologue_opt` (caller passes `prologue_opt=True`).
24+
- VERIFIED backend-agnostic: `grep -ciE 'cuda|nvidia|libtriton|ptx|sass'
25+
op_emitters/arith.py` = **0**. The fold is a pure TIR constant-true
26+
substitution that lowers identically to CUDA and Metal.
27+
28+
## §P1 dstates — EXECUTED on gb10 (em_dstates_cpasync.py)
29+
30+
Input TTIR `/tmp/ttir7/_chunk_scan_bwd_dstates.ttir` (895 lines) carries the full
31+
guard chain: extsi=102, cmpi=109, INT_MAX=66, andi=54.
32+
33+
Interleaved CUDA-event timing, N=60/rep × 4 reps, median-of-medians:
34+
35+
| variant | spill (STL+LDL) | LDGSTS | IMAD | ISETP | parity (MAXDIFF) | ms |
36+
|---|---|---|---|---|---|---|
37+
| OFF (un-routed serial prologue, plain LDG) | 2494 (STL 1688 + LDL 806) | 0 | 1249 | 1280 | 4.882812e-04 | **3140.2** |
38+
| OPT (fold + cp.async, prologue_opt=True) | **272** (STL 136 + LDL 136) | 16 | 803 | 796 | 4.882812e-04 | **745.1** |
39+
| native Triton | — | — | (298) | (69) | reference | **1.19** |
40+
41+
- OFF 3140 ms reproduces the PROVEN ROOT CAUSE serial-prologue baseline (~3144 ms).
42+
- Fold + cp.async: 4.21x faster than OFF; spills 2494 -> 272 (9.2x); IMAD/ISETP halved.
43+
- Parity is byte-identical 4.882812e-04 for BOTH OFF and OPT (guards always-true ->
44+
same values), all 29 360 128 outputs nonzero, allclose(1e-3) PASS.
45+
46+
### Generated §P1 .cu — the guard prologue is GONE
47+
48+
From the EXECUTED tvm_kernels.cu (00000 = OFF, 00001 = OPT):
49+
50+
| | OFF .cu | OPT .cu |
51+
|---|---|---|
52+
| lines | 628 | 336 |
53+
| `2147483647` guard sites | **14** | **0** |
54+
| `for(` loops | 168 | 84 |
55+
56+
The fold drops every overflow-guard comparison (14 -> 0) and halves the loop
57+
count in the consumed/lowered CUDA source. SASS instruction count 35 574 -> 14 809.
58+
59+
## Fold confirmed across the whole routed kernel family (build-only, /tmp/route_popt.py)
60+
61+
`prologue_opt=False` vs `True`, guard count in generated CUDA:
62+
63+
| kernel | guards OFF | guards FOLD | mma | .cu bytes OFF->FOLD |
64+
|---|---|---|---|---|
65+
| _chunk_scan_bwd_dc | 14 | **0** | 2 | 38271 -> 23943 |
66+
| _chunk_state_bwd_db | 15 | **0** | 2 | 42823 -> 27468 |
67+
| _chunk_state_bwd_dx | 20 | **0** | 2 | 54284 -> 39413 |
68+
| _chunk_state_bwd_ddAcs_stable | 16 | **0** | 2 | 42922 -> 31133 |
69+
70+
Fold removes ALL guards in every kernel; `mma=2` (GEMM) and `__global__` intact;
71+
.cu shrinks ~30-37%. No regression: every kernel that routed with popt=False also
72+
routes with popt=True, identical mma count.
73+
74+
## Metal codegen check — CODEGEN ONLY, NO Apple-GPU dispatch
75+
76+
The SAME folded prim (`prologue_opt=True`, NO cp.async / no CUDA-only op) lowers
77+
through the full TileLang Metal pipeline to valid MSL
78+
(`metal_foldcheck2.py` on `_chunk_state_bwd_ddAcs_stable`, 80 guards in):
79+
80+
- `tilelang.lower(pf, target=Target("metal"))` -> CompiledArtifact OK.
81+
- `art.kernel_source`: 43805 chars of MSL — `#include <metal_stdlib>`,
82+
`#include <metal_simdgroup>`, `using namespace metal;`, `kernel`, `threadgroup`,
83+
`device` qualifiers; AtomicAdd lowered to `tl::AtomicAdd`.
84+
- `HAS_GUARD_2147483647 = 0` in the MSL — **the guard chain is folded out in the
85+
Metal output exactly as in CUDA**. No CUDA-only op. The fold is Metal-portable.
86+
87+
(Host is Linux/CUDA; it physically cannot launch an Apple GPU — the check is
88+
inherently watchdog-safe. The dstates kernel's multi-stage copy loop trips
89+
TileLang's Metal `PipelinePlanning` pass — an orthogonal software-pipelining
90+
limitation present with or without the fold, see `metal_foldcheck.py` — so the
91+
fold-portability proof uses the non-pipelined `ddAcs_stable` kernel.)
92+
93+
## Honest gap and the "272 -> ~0 / toward native 1.19 ms" target
94+
95+
- The emitter fold eliminates the DOMINANT i32-overflow serial prologue: 14 -> 0
96+
guard sites in the .cu, 2494 -> 272 spills, 3140 -> 745 ms. This is the working,
97+
EXECUTED, bit-correct path.
98+
- The residual 272 spills are NOT the guard prologue — they are the kernel's
99+
genuine 64×64 GEMM working set. The remaining 745 ms vs native 1.19 ms gap
100+
(626x) is GEMM tiling/scheduling (32 vs 256 HMMA tile, copy staging), NOT the
101+
folded guards.
102+
- The TTIR-level COMPLETE fold (`_fold_ttir` in jit_to_ttir.py) does reach native
103+
guard-chain parity on TEXT (895 -> 362 lines; extsi 102->0, cmpi 109->7,
104+
INT_MAX 66->0, andi 54->3, matching native Triton's folded .ttir). BUT feeding
105+
that pre-folded TTIR back through OUR walker fails
106+
(`EmitError: PtrState references unresolved SSA value '%112'`) because Triton's
107+
`rewrite_tensor_pointer` canonicalizer restructures the addressing into a form
108+
the walker's PtrState resolution cannot consume. So the TTIR-level path is NOT a
109+
working executed route for dstates; the emitter-level fold (OPT above) is.
110+
111+
## GO / NO-GO
112+
113+
GO for the backend-agnostic guard fold: it is real (always-true equivalence,
114+
bit-correct 4.882812e-04), in OUR Metal-portable layer (zero libtriton/CUDA refs),
115+
drops §P1 from 3140 -> 745 ms (4.21x) by eliminating the guard prologue (14 -> 0
116+
.cu guard sites, 2494 -> 272 spills), and the SAME folded prim emits valid MSL with
117+
the guard folded out (Metal-portable). NO-GO on reaching native 1.19 ms: the
118+
remaining 745 ms is the GEMM tiling, a separate tuning axis, not the guard fold.
119+
120+
## Reproduce from HEAD
121+
122+
```
123+
ssh gb10 'mkdir /tmp/gb10_gpu_owner.lock && echo own>/tmp/gb10_gpu_owner.lock/id'
124+
ssh gb10 'source /home/dave/cppmega-venv/bin/activate; cd /home/dave/source/tilelang;
125+
python3 poc/triton_frontend/_test_harness/tridao_parity/em_dstates_cpasync.py' # §P1 EXECUTED
126+
ssh gb10 '... python3 poc/triton_frontend/_test_harness/tridao_parity/metal_foldcheck2.py' # MSL codegen
127+
ssh gb10 'rm -rf /tmp/gb10_gpu_owner.lock'
128+
```

0 commit comments

Comments
 (0)