Phase 4: tiny-gpu backend (memory plan, Kernel IR, registers, .tgpu)

CLAUDE.md

@@ -19,6 +19,7 @@ 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)
 pytest -q                                   # full suite
 pytest --update-golden                      # rewrite tests/golden/* after reviewing the diff
 cd web && pnpm dev | build | lint | typecheck | test

CONTEXT.md

@@ -18,7 +18,11 @@ faked; host-side work is always labelled as host work.
 - Phase 3 (TPU backend): DONE. 287 tests green. Graph IR -> Tile IR (`tinyc/tpu/`) -> `.ttpu.json` (schema `docs/ttpu.schema.json`,
   spec `docs/TTPU_FORMAT.md`); `runtime/ttpu-exec.mjs` runs whole programs on the real RTL; cost model == measurement (pinned
   in `tests/perf_baseline.json`). `tinyc FILE --target tpu --run` works.
-- Next: Phase 4 (GPU backend -O0: Kernel IR, memory planner, regalloc, `.tgpu`, gpu_runner).
+- 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.
 
 ## Fixed decisions (do not re-litigate; change only via docs/ARCHITECTURE.md)
 | Decision | Choice |
@@ -33,6 +37,9 @@ faked; host-side work is always labelled as host work.
 | TPU mapping | 4×4 int8 passes, zero padding, K-accumulation + relu/bias as `host.*` ops |
 | TPU executor | Exactly one: `runtime/ttpu-exec.mjs` (Node in tests, browser in UI). `tinyc/tpu/interp.py` is a second, independent *reference* reader of the same JSON spec, never a second executor |
 | 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 |
+| 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 |
 | Fusion | `matmul{epilogue=bias,relu}`; the fused op takes the outermost absorbed op's loc. Backends must lower the epilogue |

README.md

@@ -4,11 +4,17 @@ 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 3.** The frontend, optimizing pass pipeline and the **tiny-tpu backend** work:
-`tinyc examples/mm_relu.tiny --target tpu --run` compiles to tile passes and runs them on the real
-RTL, bit-exact against numpy. The tiny-gpu backend comes next.
+**Status: Phase 4.** 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.
 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).
+[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).
 
 ```bash
 git clone --recurse-submodules <repo> && cd tiny-compiler

docs/PASSES.md

@@ -63,3 +63,12 @@ a pass whose A or B window is a constant all-zero tile is removed together with
 to the previous kept accumulation, becomes a standalone `epilogue` (bias present), or disappears
 (relu only: relu of the zero-initialized tile is a no-op). Pass ids are renumbered. `--verify-each`
 validates the schedule after lowering and after each tile pass, by interpreting the emitted JSON.
+
+## 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.

docs/PLAN.md

@@ -198,7 +198,11 @@ Static estimates printed by `tinyc --stats` and shown before you press Run:
 - GPU: instructions issued, simulator cycles (as tiny-gpu defines them), global/shared loads
   and stores, registers used.
 After running, measured numbers sit next to the estimates. **Estimate and measurement must
-match exactly** for the TPU and for GPU instruction counts. That is a test (§5, layer 6).
+match exactly** for the TPU. On the GPU, loads and stores are exact; instructions and cycles are
+**bounds** (`min`/`max`), because relu's branch is data dependent (a clamping thread runs one more
+instruction, a non-clamping thread takes a skip branch that adds a cycle). The exact numbers for
+given inputs come from the Kernel IR interpreter, whose counts must equal the simulator's (§5,
+layer 6).
 
 ---
 

docs/TGPU_FORMAT.md

@@ -0,0 +1,62 @@
+# tiny-gpu output: `.tgpu` kernels and the launch document
+
+`tinyc FILE --target gpu --emit tgpu` prints the assembly; `--emit launch` prints the launch
+document (machine-checked by `docs/launch.schema.json`); `--emit kernel` prints the Kernel IR.
+
+## Mapping
+One thread per output element. A rank-2 result `[M, N]` launches **M blocks of N threads**:
+`R5` (block id) is the row, `R6` (thread in block) is the column, and `R7` (global thread id) is
+the flat row-major index. Rank-1 results use one block. Each Graph IR op is one kernel launch
+(a fused `matmul{epilogue=...}` is one kernel); launches share global memory.
+
+`R5`, `R6`, `R7` are set by the simulator and **only ever read**. Kernel code uses
+`R0-R4, R8..R(n-1)`; an assertion in the allocator and a runtime check in the runner
+(`reserved_intact`) enforce this. Nothing is passed as kernel arguments: all addresses are
+immediates from the memory plan.
+
+## What tiny-gpu has and does not have
+`SET ADD MUL LD ST CMP BRLT BRGT` (and `JMP BEQ BNE BRZ SYNC SYNCB SHLD SHST CSWAP`). There is no
+SUB, DIV, shift or MAX, so:
+- loops count up: `SET k,0 / L: ... / ADD k,k,1 / CMP k,K / BRLT L`;
+- `relu(x)` is `CMP x,0 / BRGT skip / SET x,0 / skip:`;
+- there is no base+offset addressing: every access computes its address in a register.
+
+Assembler traps found by measurement (and enforced by the emitter):
+- **a label must not start with `R`/`r`**: tiny-gpu's assembler parses it as a register and
+  crashes (`relu_end:` -> `ValueError`). Labels here are `L0`, `L1`, ...;
+- **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.
+
+## 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
+the simulator's.
+
+## Launch document
+```json
+{ "format": "tgpu-launch", "version": 1, "name": "main", "num_registers": 16, "mem_size": 140,
+  "buffers": { "X": {"addr": 0, "dtype": "i8", "shape": [4,8], "kind": "input"},
+               "2": {"addr": 104, "dtype": "i32", "shape": [4,6], "kind": "temp"} },
+  "inputs": ["X"], "outputs": ["3"],
+  "kernels": [ {"name": "k0_matmul", "value": "2", "op": "matmul", "grid": [4,6],
+                "shared_size": 0, "asm": "SET R0, 0 ; ..."} ],
+  "estimate": { ... } }
+```
+Memory is word addressed (`i8` takes one int32 word). Constants (`kind: const`, with `data`) and
+inputs are the initial memory image; everything else starts at 0. Intermediates reuse the words of
+buffers whose last reader has finished (`gpu/memplan.py`); an op's result never overlaps what that
+same op reads.
+
+## Executor
+`runtime/gpu_runner.py::run_launch` runs a launch document on the real simulator: a fresh
+`TinyGPU` per kernel built at its final thread count (the simulator's `set_grid` resize leaves
+R7 = 0 for added threads), global memory carried over. It counts instructions per opcode without
+changing semantics.
+
+## Source mapping
+Each instruction ends with `; line:col`, the source position of the Graph IR op it came from.

docs/launch.schema.json

@@ -0,0 +1,85 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$id": "https://github.com/deaneeth/tiny-compiler/docs/launch.schema.json",
+  "title": "tgpu launch document (version 1)",
+  "type": "object",
+  "required": ["format", "version", "name", "num_registers", "mem_size", "buffers", "inputs", "outputs", "kernels", "estimate"],
+  "additionalProperties": false,
+  "properties": {
+    "format": { "const": "tgpu-launch" },
+    "version": { "const": 1 },
+    "name": { "type": "string" },
+    "num_registers": { "type": "integer", "minimum": 8 },
+    "mem_size": { "type": "integer", "minimum": 0 },
+    "buffers": { "type": "object", "additionalProperties": { "$ref": "#/$defs/buffer" } },
+    "inputs": { "type": "array", "items": { "type": "string" } },
+    "outputs": { "type": "array", "items": { "type": "string" }, "minItems": 1 },
+    "kernels": { "type": "array", "items": { "$ref": "#/$defs/kernel" } },
+    "estimate": { "$ref": "#/$defs/estimate" }
+  },
+  "$defs": {
+    "buffer": {
+      "type": "object",
+      "required": ["addr", "dtype", "shape", "kind"],
+      "additionalProperties": false,
+      "properties": {
+        "addr": { "type": "integer", "minimum": 0 },
+        "dtype": { "enum": ["i8", "i32"] },
+        "shape": { "type": "array", "items": { "type": "integer", "minimum": 1 }, "minItems": 1, "maxItems": 2 },
+        "kind": { "enum": ["input", "const", "temp"] },
+        "data": { "type": "array", "description": "const buffers only: nested lists matching shape" }
+      }
+    },
+    "kernel": {
+      "type": "object",
+      "required": ["name", "value", "op", "grid", "shared_size", "asm"],
+      "additionalProperties": false,
+      "properties": {
+        "name": { "type": "string" },
+        "value": { "type": "string", "description": "the Graph IR value this kernel computes" },
+        "op": { "enum": ["matmul", "add", "relu"] },
+        "grid": { "type": "array", "items": { "type": "integer", "minimum": 1 }, "minItems": 2, "maxItems": 2 },
+        "shared_size": { "type": "integer", "minimum": 0 },
+        "asm": { "type": "string" }
+      }
+    },
+    "range": {
+      "type": "object",
+      "required": ["min", "max"],
+      "additionalProperties": false,
+      "properties": { "min": { "type": "integer", "minimum": 0 }, "max": { "type": "integer", "minimum": 0 } }
+    },
+    "estimate": {
+      "type": "object",
+      "required": ["launches", "threads", "instructions", "cycles", "loads", "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 },
+        "registers": { "type": "integer", "minimum": 0 },
+        "per_kernel": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "required": ["name", "grid", "threads", "instructions", "cycles", "loads", "stores", "registers"],
+            "additionalProperties": false,
+            "properties": {
+              "name": { "type": "string" },
+              "grid": { "type": "array", "items": { "type": "integer" }, "minItems": 2, "maxItems": 2 },
+              "threads": { "type": "integer", "minimum": 1 },
+              "instructions": { "$ref": "#/$defs/range" },
+              "cycles": { "$ref": "#/$defs/range" },
+              "loads": { "type": "integer", "minimum": 0 },
+              "stores": { "type": "integer", "minimum": 0 },
+              "registers": { "type": "integer", "minimum": 0 }
+            }
+          }
+        }
+      }
+    }
+  }
+}

runtime/gpu_runner.py

@@ -28,6 +28,7 @@ class GpuResult:
     cycles: int  # tiny-gpu "cycles": a cycle ends at a branch/jump/sync, not at every instruction
     instructions: int  # thread-instructions executed; use this next to `cycles` (PLAN §2.1)
     op_counts: Counter = field(default_factory=Counter)
+    reserved_intact: bool = True  # R5 (block), R6 (thread), R7 (tid) still hold what the sim set
 
     def read(self, addr: int, length: int) -> list[int]:
         return [int(v) for v in self.memory[addr : addr + length]]
@@ -91,9 +92,76 @@ def run_tgpu(
 
     if gpu.active.any():
         raise RuntimeError(f"kernel still running after {max_cycles} cycles")
+    tid = np.arange(num_blocks * tpb)
+    expected = {5: tid // tpb, 6: tid % tpb, 7: tid}
+    intact = all((gpu.registers[:, r] == want).all() for r, want in expected.items())
     return GpuResult(
         memory=gpu.memory.copy(),
         cycles=len(gpu.history_registers),
         instructions=sum(counts.values()),
         op_counts=counts,
+        reserved_intact=intact,
     )
+
+
+@dataclass
+class LaunchResult:
+    outputs: dict[str, np.ndarray]  # int64, original shapes
+    kernels: list[tuple[str, GpuResult]]  # (kernel name, its measurements), in launch order
+    reserved_intact: bool
+
+    @property
+    def cycles(self) -> int:
+        return sum(r.cycles for _, r in self.kernels)
+
+    @property
+    def instructions(self) -> int:
+        return sum(r.instructions for _, r in self.kernels)
+
+    @property
+    def op_counts(self) -> Counter:
+        total: Counter = Counter()
+        for _, r in self.kernels:
+            total.update(r.op_counts)
+        return total
+
+
+def run_launch(launch: Mapping, inputs: Mapping[str, object]) -> LaunchResult:
+    """Run a launch document (`docs/TGPU_FORMAT.md`) on the real tiny-gpu simulator.
+
+    Global memory persists across kernels; each kernel gets a fresh simulator built at its final
+    thread count (tiny-gpu's `set_grid` resize leaves R7 = 0 for threads it adds).
+    """
+    memory = np.zeros(launch["mem_size"], dtype=np.int64)
+    for name, b in launch["buffers"].items():
+        if b["kind"] == "const":
+            flat = np.asarray(b["data"], dtype=np.int64).reshape(-1)
+        elif b["kind"] == "input":
+            arr = np.asarray(inputs[name], dtype=np.int64)
+            if list(arr.shape) != b["shape"]:
+                raise ValueError(f"input '{name}' must have shape {b['shape']}")
+            flat = arr.reshape(-1)
+        else:
+            continue
+        memory[b["addr"] : b["addr"] + len(flat)] = flat
+
+    results, intact = [], True
+    for k in launch["kernels"]:
+        res = run_tgpu(
+            k["asm"],
+            grid=tuple(k["grid"]),
+            mem_size=launch["mem_size"],
+            memory_init={0: memory},
+            num_registers=launch["num_registers"],
+            shared_size=k["shared_size"],
+        )
+        memory = res.memory.astype(np.int64)
+        intact = intact and res.reserved_intact
+        results.append((k["name"], res))
+
+    outputs = {}
+    for name in launch["outputs"]:
+        b = launch["buffers"][name]
+        size = int(np.prod(b["shape"]))
+        outputs[name] = memory[b["addr"] : b["addr"] + size].reshape(b["shape"])
+    return LaunchResult(outputs, results, intact)

tests/golden/kernel/big_mm.kernel

@@ -0,0 +1,20 @@
+kernel @k0_matmul grid=(12 blocks, 12 threads)  // %1 = matmul at 4:5 {
+  v0 = set 0
+  for v1 in 0..12 {
+    v2 = mul block, 12
+    v2 = add v2, v1
+    v2 = add v2, 0
+    v3 = ld [v2]
+    v4 = mul v1, 12
+    v4 = add v4, thread
+    v4 = add v4, 144
+    v5 = ld [v4]
+    v6 = mul v3, v5
+    v0 = add v0, v6
+  }
+  if le v0, 0 {
+    v0 = set 0
+  }
+  v7 = add tid, 288
+  st [v7], v0
+}

tests/golden/kernel/const_mm.kernel

@@ -0,0 +1,20 @@
+kernel @k0_matmul grid=(3 blocks, 2 threads)  // %2 = matmul at 5:5 {
+  v0 = set 0
+  for v1 in 0..2 {
+    v2 = mul block, 2
+    v2 = add v2, v1
+    v2 = add v2, 0
+    v3 = ld [v2]
+    v4 = mul v1, 2
+    v4 = add v4, thread
+    v4 = add v4, 6
+    v5 = ld [v4]
+    v6 = mul v3, v5
+    v0 = add v0, v6
+  }
+  if le v0, 0 {
+    v0 = set 0
+  }
+  v7 = add tid, 10
+  st [v7], v0
+}

tests/golden/kernel/mlp_gpu.kernel

@@ -0,0 +1,41 @@
+kernel @k0_matmul grid=(4 blocks, 6 threads)  // %2 = matmul at 6:5 {
+  v0 = set 0
+  for v1 in 0..8 {
+    v2 = mul block, 8
+    v2 = add v2, v1
+    v2 = add v2, 0
+    v3 = ld [v2]
+    v4 = mul v1, 6
+    v4 = add v4, thread
+    v4 = add v4, 32
+    v5 = ld [v4]
+    v6 = mul v3, v5
+    v0 = add v0, v6
+  }
+  v7 = add thread, 80
+  v8 = ld [v7]
+  v0 = add v0, v8
+  if le v0, 0 {
+    v0 = set 0
+  }
+  v9 = add tid, 104
+  st [v9], v0
+}
+
+kernel @k1_matmul grid=(4 blocks, 3 threads)  // %3 = matmul at 7:5 {
+  v0 = set 0
+  for v1 in 0..6 {
+    v2 = mul block, 6
+    v2 = add v2, v1
+    v2 = add v2, 104
+    v3 = ld [v2]
+    v4 = mul v1, 3
+    v4 = add v4, thread
+    v4 = add v4, 86
+    v5 = ld [v4]
+    v6 = mul v3, v5
+    v0 = add v0, v6
+  }
+  v7 = add tid, 128
+  st [v7], v0
+}