Phase 5: tiny-gpu kernel optimizations (-O1/-O2)

CLAUDE.md

@@ -19,7 +19,9 @@ source scripts/env.sh
 tinyc examples/mm_relu.tiny --locs          # optimized Graph IR (--emit ast, -O0, --target gpu|tpu,
                                             #   --print-after-all, --verify-each, --time-passes, --disable PASS)
 tinyc examples/mm_relu.tiny --target tpu --run          # compile + run on the real RTL (also --emit tile|ttpu, --stats)
-tinyc examples/mm_relu.tiny --target gpu --run          # compile + run on the simulator (--emit kernel|tgpu|launch, --gpu-regs N)
+tinyc examples/mm_relu.tiny --target gpu --run -O2      # compile + run on the simulator (--emit kernel|tgpu|launch, --gpu-regs N,
+                                                        #   --disable licm|strength-reduce|shared-tile|unroll|fold|kernel-dce)
+python scripts/opt_table.py --write                     # regenerate the measured table in docs/ARCHITECTURE.md
 pytest -q                                   # full suite
 pytest --update-golden                      # rewrite tests/golden/* after reviewing the diff
 cd web && pnpm dev | build | lint | typecheck | test

CONTEXT.md

@@ -21,8 +21,10 @@ faked; host-side work is always labelled as host work.
 - Phase 4 (GPU backend, -O0): DONE. 443 tests green. Graph IR -> memory plan -> Kernel IR -> linear-scan registers -> `.tgpu` + launch JSON
   (`tinyc/gpu/`, schema `docs/launch.schema.json`, spec `docs/TGPU_FORMAT.md`). Every example, including the chained-matmul
   `mlp_gpu`, runs bit-exact on the real simulator; the same program gives identical results on both chips. `tinyc FILE --target gpu --run`.
-- Next: Phase 5 (GPU optimizations -O1/-O2: LICM, strength reduction, unrolling, shared-memory tiling). `-O` currently
-  controls only the graph passes; decide how it maps to kernel optimizations before adding them.
+- Phase 5 (GPU kernel optimizations): DONE. 557 tests green. `-O0/-O1/-O2`: licm, strength-reduce, kernel-dce | + shared-tile, unroll, fold.
+  Measured table in `docs/ARCHITECTURE.md` (generated by `scripts/opt_table.py`, drift-tested). `-O2` has strictly fewer
+  global loads than `-O0` on every matmul example (exact formula tested). Cost model is now an abstract interpretation.
+- Next: Phase 6 (differential fuzzing + hardening; nightly 5000-example job; publishable CLI).
 
 ## Fixed decisions (do not re-litigate; change only via docs/ARCHITECTURE.md)
 | Decision | Choice |
@@ -39,6 +41,9 @@ faked; host-side work is always labelled as host work.
 | TPU host work | `acc`/`epilogue`/`host_add`/`host_relu`/`alloc` are host commands, counted separately from chip passes everywhere |
 | GPU mapping | One thread per output element: block = row (R5), thread = column (R6), tid (R7) = flat index. One kernel per Graph op (fused = one). Kernel code is read-only on R5-R7 (asserted + checked at run time) |
 | GPU cycles | A thread needs taken-branches+1 cycles, a launch the max. Report instructions with cycles. `estimate` gives min/max (relu is data dependent), loads/stores exact; `gpu.interp` counts must equal the simulator's |
+| `-O` levels | -O0 naive; -O1 graph passes + kernel licm/strength-reduce; -O2 + shared-tile, unroll, fold (gpu only; tpu -O2 == -O1). Kernel passes are switchable by name (`--disable`) |
+| Kernel soundness | Lowering assigns every temporary once; only accumulators/pointers/counters are reassigned. Passes may move or rewrite single-assignment values freely |
+| GPU barrier | `SYNCB` only at kernel top level; waiting threads re-execute it each step (counts modelled exactly in `gpu.interp`); out-of-range shared access is a fault in the interpreter |
 | GPU asm rules | Labels are `L<n>` (a label starting with R is parsed as a register); a kernel never ends with an `if` (stray extra cycle); one fresh simulator per launch at its final size, memory carried over |
 | TPU cost | `passes * 15` RTL cycles; utilization = useful MACs (`M*K*N`) / (64 * passes) |
 | Overflow | `range-check` proves every value fits i32 by intervals; otherwise reject (tiny-gpu raises on overflow). i32 *inputs* are assumed within ±2^20 (`PassContext.i32_input_range`); runners must respect it |

README.md

@@ -4,15 +4,16 @@ One tiny tensor language, two real chips: `C = relu(A @ B)` compiles to
 [tiny-gpu](https://github.com/deaneeth/tiny-gpu) assembly or to a
 [tiny-tpu](https://github.com/deaneeth/tiny-tpu) tile schedule, and runs on the real simulator / RTL.
 
-**Status: Phase 4.** The frontend, optimizing pass pipeline and **both backends** work. One program,
+**Status: Phase 5.** The frontend, optimizing pass pipeline and **both backends** work. One program,
 two chips, bit-exact against numpy on each:
 
 ```bash
 tinyc examples/mm_relu.tiny --target tpu --run   # tile passes on the real tiny-tpu RTL
 tinyc examples/mm_relu.tiny --target gpu --run   # SIMT kernels on the tiny-gpu simulator
 ```
 
-GPU kernel optimizations are next.
+`-O2` tiles operands into tiny-gpu's shared memory and unrolls the loops (see
+[docs/ARCHITECTURE.md](docs/ARCHITECTURE.md) for the measured effect of every optimization).
 Both targets can be driven from pytest. See [docs/PLAN.md](docs/PLAN.md) for the roadmap,
 [docs/LANGUAGE.md](docs/LANGUAGE.md), [docs/IR.md](docs/IR.md) and [docs/PASSES.md](docs/PASSES.md) and [docs/TTPU_FORMAT.md](docs/TTPU_FORMAT.md) and [docs/TGPU_FORMAT.md](docs/TGPU_FORMAT.md).
 

docs/ARCHITECTURE.md

@@ -0,0 +1,107 @@
+# Architecture
+
+How the compiler is put together, and the decisions that are not obvious from the code. Change a
+fixed decision here first (see `CONTEXT.md`). Per-pass details are in [PASSES.md](PASSES.md); the
+formats are in [LANGUAGE.md](LANGUAGE.md), [IR.md](IR.md), [TTPU_FORMAT.md](TTPU_FORMAT.md) and
+[TGPU_FORMAT.md](TGPU_FORMAT.md).
+
+## Pipeline
+
+```
+ .tiny -> AST -> Graph IR -> graph passes -+-> Tile IR -> tile passes -> .ttpu.json   (tiny-tpu RTL)
+                                           |
+                                           +-> memory plan -> Kernel IR -> kernel passes
+                                                            -> registers -> .tgpu + launch  (tiny-gpu)
+```
+
+Every level has a printer, a verifier and an interpreter; every pass can be validated against the
+numpy oracle (`--verify-each`); and both backends run on the real chip/simulator, never on a
+re-implementation.
+
+## Optimization levels (`-O`)
+
+| Level | Graph IR | Tile IR (tpu) | Kernel IR (gpu) |
+|---|---|---|---|
+| `-O0` | shape-infer, range-check | none | naive kernels |
+| `-O1` (default) | + simplify fixpoint, fuse | + skip-zero-passes | + licm, strength-reduce, kernel-dce |
+| `-O2` | same as `-O1` | same as `-O1` | + shared-tile, unroll, fold |
+
+Individual passes can be switched off by name (`--disable licm`), which is how the table below
+isolates each optimization.
+
+## Kernel optimizations, measured
+
+Real simulator runs (`tinyc/gpu/report.py`; regenerate with `python scripts/opt_table.py --write`,
+and a test fails if this table drifts). "Cycles" are tiny-gpu cycles: a thread's cycle ends at a taken
+branch, so straight-line code is one cycle however long it is; read them together with instructions.
+`relu` makes the exact counts data dependent, so these use a fixed seed.
+
+<!-- opt-table:start -->
+**`mm_relu` 4x4** (`C = relu(A @ B)`, int8, seed 42)
+
+| configuration | instructions | cycles | global loads | shared loads | registers |
+|---|---:|---:|---:|---:|---:|
+| -O0 (naive) | 1000 | 6 | 144 | 0 | 4 |
+| + licm | 840 (-16%) | 5 | 128 (-11%) | 0 | 5 |
+| + strength-reduce (= -O1) | 728 (-27%) | 5 | 128 (-11%) | 0 | 6 |
+| -O2 with shared-memory tiling only | 728 (-27%) | 6 | 80 (-44%) | 64 | 5 |
+| -O2 with unrolling + folding only | 504 (-50%) | 2 | 128 (-11%) | 0 | 7 |
+| -O2 (all passes) | 504 (-50%) | 3 | 80 (-44%) | 64 | 5 |
+
+**`mm_relu` 8x8** (`C = relu(A @ B)`, int8, seed 42)
+
+| configuration | instructions | cycles | global loads | shared loads | registers |
+|---|---:|---:|---:|---:|---:|
+| -O0 (naive) | 7320 | 10 | 1088 | 0 | 4 |
+| + licm | 6168 (-16%) | 9 | 1024 (-6%) | 0 | 5 |
+| + strength-reduce (= -O1) | 5208 (-29%) | 9 | 1024 (-6%) | 0 | 6 |
+| -O2 with shared-memory tiling only | 4952 (-32%) | 10 | 576 (-47%) | 512 | 5 |
+| -O2 with unrolling + folding only | 3544 (-52%) | 2 | 1024 (-6%) | 0 | 7 |
+| -O2 (all passes) | 3288 (-55%) | 3 | 576 (-47%) | 512 | 5 |
+
+<!-- opt-table:end -->
+
+What each pass does, and what it costs:
+
+- **licm** hoists address arithmetic that does not depend on the loop counter (`row*K + base`).
+  It never hoists loads, so global loads are unchanged. It *raises register pressure*: hoisted values
+  stay live across the loop, which is why `-O0` is the code that fits the smallest register file.
+- **strength-reduce** turns `k*N + col + base` (a multiply and two adds per iteration) into a pointer
+  initialised once and bumped by `N`. Chains of one statement are left alone: an add costs the same as
+  the increment that would replace it.
+- **shared-tile** exploits that every thread of a block reads the same row of A. The block's threads
+  copy it into shared memory cooperatively, meet at a barrier (`SYNCB`), and read it with `SHLD`.
+  Global loads per matmul drop from `2*M*N*K` to `M*K + M*N*K` (about -44% to -47% above). It is roughly
+  instruction-neutral: the copy and the barrier cost about what the cheaper `shld` (no address
+  arithmetic) saves, and in loop form it adds a cycle. Skipped for one-thread blocks.
+- **unroll** removes the loop for trip counts up to 16, so the kernel has no taken back edges and its
+  cycle count collapses. **fold** then folds the constant addresses and copies that unrolling exposes.
+
+### Why -O2 is not simply better
+
+- Unrolling alone is the biggest instruction and cycle win but the most register-hungry (7 registers
+  above: every iteration's temporaries are live at once before the allocator packs them); tiling
+  *reduces* pressure back to 5 because the A address chain disappears. The default register file (13
+  allocatable) has room, but `--gpu-regs 8` only fits `-O0`.
+- It needs a barrier, and tiny-gpu re-executes `SYNCB` on every step a thread spends waiting, so the
+  instruction count includes those repeats (modelled exactly in `gpu/interp.py`).
+- Tiling only pays when a block has at least two threads; with a reduction as short as `K = 2` the
+  copy and barrier are not amortised, so loads and cycles improve more than instructions.
+- A kernel tiles `K` shared words per operand per block; the simulator allocates exactly that.
+
+## Decisions worth knowing
+
+| Decision | Why |
+|---|---|
+| Temporaries are assigned once; only accumulators, pointers and counters are reassigned | It makes every kernel pass sound without dataflow analysis: a single-assignment value can be moved or rewritten freely. |
+| Kernel passes are validated against the oracle after each pass | The structured kernels are interpreted (independently of the allocator and assembler); the final assembly is run on the simulator. |
+| The estimate is an abstract interpretation | Values that depend only on block/thread ids are exact, so loads, stores, shared accesses and the tiled copy's guards are exact; only a branch on loaded data (relu) gives `min`/`max`. |
+| One fresh simulator per kernel launch, memory carried over | tiny-gpu's `set_grid` resize leaves R7 = 0 for threads it adds. |
+| `-O` levels are cumulative, kernel passes are switchable by name | Reproducibility, and so the table above can attribute each gain. |
+
+## Known limits
+
+- Kernel optimization is specific to what `lower.py` emits (matmul-shaped loops); the passes are
+  written generally but only exercised on those kernels.
+- No spilling: a program that needs more registers than the budget is a compile error.
+- The dead final pointer bump of an unrolled loop is not removed (the DCE is not flow-sensitive).

docs/PASSES.md

@@ -66,9 +66,26 @@ validates the schedule after lowering and after each tile pass, by interpreting
 
 ## GPU backend stages
 After the graph pipeline, `compile_gpu_program` runs: memory plan (`gpu/memplan.py`) -> lowering to
-Kernel IR (`gpu/lower.py`, one kernel per op) -> verify -> flatten -> linear-scan register
-allocation (`gpu/regalloc.py`) -> `.tgpu` emission (`gpu/emit.py`). With `--verify-each` the
-structured kernels are interpreted against the numpy oracle after lowering, and the emitted assembly
-is run on the real simulator against the oracle after emission. `--print-after-all` shows
-`lower-to-kernel` and `regalloc-emit`. Phase 4 emits the straightforward (-O0) kernel code; kernel
-optimizations (loop-invariant motion, unrolling, shared-memory tiling) are Phase 5.
+Kernel IR (`gpu/lower.py`, one kernel per op) -> kernel passes (`gpu/kernel_passes.py`) -> verify ->
+flatten -> linear-scan register allocation (`gpu/regalloc.py`) -> `.tgpu` emission (`gpu/emit.py`).
+
+| `-O` | Kernel passes |
+|---|---|
+| 0 | none (the naive kernels) |
+| 1 | `licm`, `strength-reduce`, `kernel-dce` |
+| 2 | `licm`, `strength-reduce`, `shared-tile`, `unroll`, `fold`, `kernel-dce` |
+
+- `licm`: hoist single-assignment arithmetic that does not change inside a loop (never loads).
+- `strength-reduce`: replace a `mul`/`add` chain on the loop counter by a pointer bumped each iteration
+  (chains of two or more statements only; not when the value is read after the loop).
+- `kernel-dce`: remove values that never reach a store, `if` or loop control (liveness-based, so a dead
+  value that only updates itself goes too).
+- `shared-tile`: a loop load whose address is `u + k`, with `u` the same for the whole block, is copied
+  cooperatively into shared memory (`SHST`), a barrier (`SYNCB`) follows, and the loop reads `SHLD`.
+- `unroll`: full unroll of innermost loops of at most 16 iterations (the counter becomes a constant).
+- `fold`: constant propagation, constant folding, add-chain reassociation and copy propagation.
+
+With `--verify-each` the structured kernels are interpreted against the numpy oracle after lowering
+and after *every* kernel pass, and the emitted assembly is run on the real simulator against the
+oracle. `--print-after-all` shows `lower-to-kernel`, each kernel pass, and `regalloc-emit`. Each pass
+can be disabled by name. The measured effect of each is in [ARCHITECTURE.md](ARCHITECTURE.md).

docs/TGPU_FORMAT.md

@@ -27,14 +27,29 @@ Assembler traps found by measurement (and enforced by the emitter):
 - **a kernel must not end with an `if`**: branching to the end of the program costs a stray extra
   cycle; the Kernel IR verifier rejects it.
 
+## Shared memory and barriers (`-O2`)
+`SHST saddr, Rs` / `SHLD Rd, saddr` access the block's own shared memory (`shared_size` words per
+block, declared per kernel in the launch document); `SYNCB` is a block-wide barrier. Measured and
+modelled behaviour:
+- a thread that reaches `SYNCB` stops for that step, and a **waiting thread executes `SYNCB` again on
+  every following step** until all threads of its block are waiting; then the block resumes on the next
+  step. With arrival step `s_t = taken branches before the barrier + 1` and release step `R = max s_t`,
+  a thread executes `SYNCB` `R - s_t + 1` times and finishes at step `R + taken branches after + 1`.
+  `gpu/interp.py` implements this and a test requires it to equal the simulator exactly, including for
+  hand-written kernels whose threads arrive at different times;
+- out-of-range `SHLD` silently reads 0 and out-of-range `SHST` is silently dropped in tiny-gpu; the
+  interpreter treats both as faults, so a compiler bug cannot hide;
+- the Kernel IR only allows a barrier at the top level of a kernel, so every thread reaches it.
+
 ## Cycles vs instructions
 tiny-gpu runs consecutive non-branch instructions of a thread in one step; a thread's step ends at
 a *taken* branch. A thread needs `taken branches + 1` cycles, a launch the maximum over threads. So
 a straight-line kernel is 1 cycle however long it is, and a K-iteration loop is K cycles.
 **Always read cycles together with instructions.** The relu `if` is data dependent: a thread that
 clamps runs one more instruction (`SET`); a thread that does not takes the skip branch, which adds a
-cycle. The static `estimate` therefore gives `min`/`max` for instructions and cycles; loads and
-stores are exact. `gpu.interp` gives exact counts for given inputs, and tests require them to equal
+cycle. The static `estimate` is an abstract interpretation: anything that depends only on block/thread ids is
+exact (loop trip counts, a tiled kernel's copy guards), so loads, stores and shared accesses are exact;
+only a branch on loaded data gives `min`/`max` for instructions and cycles. `gpu.interp` gives exact counts for given inputs, and tests require them to equal
 the simulator's.
 
 ## Launch document

docs/launch.schema.json

@@ -51,30 +51,35 @@
     },
     "estimate": {
       "type": "object",
-      "required": ["launches", "threads", "instructions", "cycles", "loads", "stores", "registers", "per_kernel"],
+      "required": ["launches", "threads", "instructions", "cycles", "loads", "stores", "shared_loads", "shared_stores", "registers", "per_kernel"],
       "additionalProperties": false,
       "properties": {
         "launches": { "type": "integer", "minimum": 0 },
         "threads": { "type": "integer", "minimum": 0 },
         "instructions": { "$ref": "#/$defs/range" },
         "cycles": { "$ref": "#/$defs/range" },
-        "loads": { "type": "integer", "minimum": 0 },
-        "stores": { "type": "integer", "minimum": 0 },
+        "loads": { "type": "integer", "minimum": 0, "description": "global loads, exact" },
+        "stores": { "type": "integer", "minimum": 0, "description": "global stores, exact" },
+        "shared_loads": { "type": "integer", "minimum": 0 },
+        "shared_stores": { "type": "integer", "minimum": 0 },
         "registers": { "type": "integer", "minimum": 0 },
         "per_kernel": {
           "type": "array",
           "items": {
             "type": "object",
-            "required": ["name", "grid", "threads", "instructions", "cycles", "loads", "stores", "registers"],
+            "required": ["name", "grid", "threads", "shared_size", "instructions", "cycles", "loads", "stores", "shared_loads", "shared_stores", "registers"],
             "additionalProperties": false,
             "properties": {
               "name": { "type": "string" },
               "grid": { "type": "array", "items": { "type": "integer" }, "minItems": 2, "maxItems": 2 },
               "threads": { "type": "integer", "minimum": 1 },
+              "shared_size": { "type": "integer", "minimum": 0 },
               "instructions": { "$ref": "#/$defs/range" },
               "cycles": { "$ref": "#/$defs/range" },
               "loads": { "type": "integer", "minimum": 0 },
               "stores": { "type": "integer", "minimum": 0 },
+              "shared_loads": { "type": "integer", "minimum": 0 },
+              "shared_stores": { "type": "integer", "minimum": 0 },
               "registers": { "type": "integer", "minimum": 0 }
             }
           }

scripts/opt_table.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+"""Print the measured optimization table, or rewrite it inside docs/ARCHITECTURE.md.
+
+python scripts/opt_table.py            # print
+python scripts/opt_table.py --write    # update docs/ARCHITECTURE.md between its markers
+"""
+
+import re
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT))
+
+from tinyc.gpu.report import opt_table_markdown
+
+START, END = "<!-- opt-table:start -->", "<!-- opt-table:end -->"
+
+
+def main() -> int:
+    table = opt_table_markdown()
+    if "--write" not in sys.argv:
+        print(table)
+        return 0
+    doc = ROOT / "docs" / "ARCHITECTURE.md"
+    text = doc.read_text()
+    new = re.sub(
+        f"{re.escape(START)}.*?{re.escape(END)}", f"{START}\n{table}\n{END}", text, flags=re.DOTALL
+    )
+    doc.write_text(new)
+    print(f"updated {doc.relative_to(ROOT)}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())