Phase 2: pass manager and graph passes

CLAUDE.md

@@ -16,7 +16,8 @@ Verilator→WASM RTL). Pipeline: `.tiny → AST → Graph IR → Tile IR | Kerne
 git submodule update --init
 python3 -m venv .venv && .venv/bin/pip install -e ".[dev]" && .venv/bin/pip install -e targets/tiny-gpu --no-deps
 source scripts/env.sh
-tinyc examples/mm_relu.tiny --locs          # Graph IR (or --emit ast); later: --target both -O2 --run --stats
+tinyc examples/mm_relu.tiny --locs          # optimized Graph IR (--emit ast, -O0, --target gpu|tpu,
+                                            #   --print-after-all, --verify-each, --time-passes, --disable PASS)
 pytest -q                                   # full suite
 pytest --update-golden                      # rewrite tests/golden/* after reviewing the diff
 cd web && pnpm dev | build | lint | typecheck | test

CONTEXT.md

@@ -13,7 +13,9 @@ faked; host-side work is always labelled as host work.
 - Phase 0 (scaffold + drive both targets from pytest): DONE. 17 tests green (runners + target contract), ruff/black clean.
 - Phase 1 (frontend + Graph IR + oracle): DONE. 90 tests green: lexer/parser/diagnostics (20 golden error messages), Graph IR
   printer/parser round trip, verifier, numpy oracle + independent scalar interpreter, 6 examples.
-- Next: Phase 2 (pass manager + graph passes).
+- Phase 2 (pass manager + graph passes): DONE. 183 tests green. Passes: shape-infer, canonicalize, const-fold, cse, dce,
+  fuse (matmul{epilogue=bias,relu}), range-check, legalize-tpu/gpu; `--verify-each`, `--print-after-all`, `--time-passes`.
+- Next: Phase 3 (TPU backend: Tile IR, tiling, `.ttpu.json`, cost model, `ttpu-exec` full command list).
 
 ## Fixed decisions (do not re-litigate; change only via docs/ARCHITECTURE.md)
 | Decision | Choice |
@@ -27,7 +29,8 @@ faked; host-side work is always labelled as host work.
 | GPU regalloc | Linear scan, budget 16 by default, no spilling (compile error) |
 | 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) |
-| Overflow | Range analysis must prove i32 fits; otherwise reject (tiny-gpu raises on overflow) |
+| 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 |
 | Verification | Oracle (numpy) + per-pass translation validation + real targets + Hypothesis differential + golden + perf pins |
 
 ## Target gotchas (see PLAN §2)

README.md

@@ -4,9 +4,10 @@ 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 1.** The frontend and Graph IR work (`tinyc examples/mm_relu.tiny`); no backends yet.
+**Status: Phase 2.** The frontend, Graph IR and optimizing pass pipeline work
+(`tinyc examples/mm_bias_relu.tiny --target gpu --print-after-all`); the chip backends come next.
 Both targets can be driven from pytest. See [docs/PLAN.md](docs/PLAN.md) for the roadmap,
-[docs/LANGUAGE.md](docs/LANGUAGE.md) and [docs/IR.md](docs/IR.md) for the language and IR.
+[docs/LANGUAGE.md](docs/LANGUAGE.md), [docs/IR.md](docs/IR.md) and [docs/PASSES.md](docs/PASSES.md).
 
 ```bash
 git clone --recurse-submodules <repo> && cd tiny-compiler

docs/IR.md

@@ -17,6 +17,11 @@ func @main(%A: i8[4,8], %B: i8[8,4]) -> (i32[4,4]) {
 | `matmul` | `%x = matmul %a, %b : T` |
 | `add` | `%x = add %a, %b : T` |
 | `relu` | `%x = relu %a : T` |
+| fused `matmul` | `%x = matmul{epilogue=bias,relu} %a, %b, %bias : T` (made by the `fuse` pass) |
+
+A fused matmul applies its epilogue stages in order to the matmul result: `bias` adds a rank-1
+`[N]` operand (a third operand) to every row, `relu` clamps at zero. Each stage's result must fit
+in i32, exactly as if the ops were separate. Allowed epilogues: `bias`, `relu`, `bias,relu`.
 
 - `loc("file":line:col)` is optional and points at the start of the source expression. It survives
   every lowering so the UI can map any output line back to the source.

docs/PASSES.md

@@ -0,0 +1,50 @@
+# Graph passes and the pass manager
+
+`tinyc/driver.py` builds the pipeline; `tinyc/passes/manager.py` runs it.
+
+```
+-O1:  shape-infer, canonicalize, const-fold, canonicalize, cse, dce, fuse, range-check, [legalize-<target>]
+-O0:  shape-infer, range-check, [legalize-<target>]
+```
+
+| Pass | What it does |
+|---|---|
+| `shape-infer` | Recomputes every result type from its operands (repairs parsed or hand-edited IR). |
+| `canonicalize` | Broadcast `add` puts the matrix first; same-shape `add` orders operands by definition order; `relu(relu(x))` becomes `relu(x)`. |
+| `const-fold` | Evaluates ops whose operands are all constants, using the oracle's own `eval_op`. A constant expression that overflows i32 is a compile error. |
+| `cse` | Merges identical ops and identical constants (same data **and** dtype). |
+| `dce` | Removes ops that never reach an output. Inputs are always kept. |
+| `fuse` | `matmul -> add(bias[N]) -> relu` becomes one `matmul{epilogue=bias,relu}`; any prefix works (`bias`, `relu`, `bias,relu`). A stage is absorbed only if the previous result has exactly one use and is not an output. The fused op sits where the last absorbed op was and takes **that op's source location** (the outermost expression). |
+| `range-check` | Interval analysis: proves every value fits in i32 or rejects the program (below). |
+| `legalize-tpu` | Every matmul operand must be an i8 input or constant. |
+| `legalize-gpu` | Every op result needs at most `gpu_max_threads` (1024) threads; total buffers must fit `gpu_mem_words` (a conservative bound). |
+
+## Range analysis
+- i8 inputs are exactly `[-128, 127]`. **i32 inputs are assumed to lie in `ctx.i32_input_range`
+  (default ±2^20)**: an i32 input can otherwise hold any 32-bit value, and then nothing involving it
+  could be proven. Whatever runs a compiled program must respect that assumption.
+- Constants are exact. A matmul of K products each in `[pmin, pmax]` lies in `[K*pmin, K*pmax]`;
+  prefix sums stay inside the hull of that interval and 0, so the proof holds for any accumulation
+  order (GPU loop, TPU tiling).
+- The result is stored in `ctx.analysis["ranges"]` as `{value name: (lo, hi)}`.
+
+## Verification
+- After **every** pass the IR is structurally verified (always on).
+- `--verify-each` also *runs* the IR after every pass on seeded random inputs (the scalar
+  interpreter) and compares with the numpy oracle's answers for the program before any pass. A
+  failure raises `PassVerifyError` naming the pass, with the IR before and after.
+- Tests: a directed IR-in/IR-out test per pass; the manager against a deliberately wrong pass;
+  every example through every target and `-O` level; Hypothesis over random programs checks that
+  the pipeline preserves semantics, is idempotent, and that range-check is *sound* (accepted
+  programs never overflow on interval-corner inputs, and every value stays inside its interval).
+- The generator emits dense-layer idioms on purpose: a plain random generator never produced a
+  fusable pattern, so fusion was silently untested.
+
+## CLI
+```
+tinyc examples/mm_bias_relu.tiny --target gpu --print-after-all --verify-each --time-passes
+tinyc examples/mm_relu.tiny -O0                 # no optimizations
+tinyc examples/mm_relu.tiny --disable fuse      # switch a pass off
+```
+Dumps and timings go to stderr; the final IR goes to stdout. Exit codes: 0 ok, 1 compile or
+legalization error, 2 usage error, 3 internal compiler error (a pass broke the program).

tests/gen.py

@@ -0,0 +1,88 @@
+"""Hypothesis strategy for random legal `.tiny` programs (small dims, so ops often line up)."""
+
+from __future__ import annotations
+
+from hypothesis import strategies as st
+
+from tinyc.ir.types import RANGES
+
+DIMS = st.sampled_from([1, 2, 3, 4])
+
+
+def _fmt(dtype, shape):
+    return f"{dtype}[{', '.join(map(str, shape))}]"
+
+
+def _nested(draw, shape, lo, hi):
+    if len(shape) == 1:
+        return [draw(st.integers(lo, hi)) for _ in range(shape[0])]
+    return [_nested(draw, shape[1:], lo, hi) for _ in range(shape[0])]
+
+
+@st.composite
+def programs(draw, max_ops: int = 6) -> str:
+    pool: list[tuple[str, str, tuple[int, ...]]] = []  # (name, dtype, shape)
+    lines: list[str] = []
+    for i in range(draw(st.integers(2, 4))):
+        shape = tuple(draw(DIMS) for _ in range(draw(st.sampled_from([1, 2, 2]))))
+        dtype = draw(st.sampled_from(["i8", "i8", "i32"]))
+        lines.append(f"input x{i} : {_fmt(dtype, shape)}")
+        pool.append((f"x{i}", dtype, shape))
+    for i in range(draw(st.integers(0, 2))):
+        shape = tuple(draw(DIMS) for _ in range(draw(st.sampled_from([1, 2]))))
+        lo, hi = RANGES["i8"]
+        lines.append(f"const c{i} : {_fmt('i8', shape)} = {_nested(draw, shape, lo, hi)}")
+        pool.append((f"c{i}", "i8", shape))
+
+    nbias = 0
+    for j in range(draw(st.integers(1, max_ops))):
+        matmuls = [
+            (a, b)
+            for a in pool
+            for b in pool
+            if len(a[2]) == 2 and len(b[2]) == 2 and a[2][1] == b[2][0]
+        ]
+        adds = []
+        for a in pool:
+            for b in pool:
+                same = a[2] == b[2]
+                bcast = len(a[2]) == 2 and len(b[2]) == 1 and a[2][1] == b[2][0]
+                if same or bcast or (len(b[2]) == 2 and len(a[2]) == 1 and b[2][1] == a[2][0]):
+                    adds.append((a, b))
+        kinds = ["relu"] + (["add"] if adds else []) + (["matmul", "dense"] if matmuls else [])
+        kind = kinds[draw(st.integers(0, len(kinds) - 1))]
+        if kind == "relu":
+            v = pool[draw(st.integers(0, len(pool) - 1))]
+            expr, dtype, shape = f"relu({v[0]})", v[1], v[2]
+        elif kind == "add":
+            a, b = adds[draw(st.integers(0, len(adds) - 1))]
+            expr, dtype = f"{a[0]} + {b[0]}", "i32"
+            shape = a[2] if len(a[2]) >= len(b[2]) else b[2]
+        else:
+            a, b = matmuls[draw(st.integers(0, len(matmuls) - 1))]
+            dtype, shape = "i32", (a[2][0], b[2][1])
+            if kind == "matmul":
+                expr = f"{a[0]} @ {b[0]}"
+            else:  # a dense-layer idiom, so the fuser has something to fuse
+                form = draw(
+                    st.sampled_from(["relu({m})", "{m} + {w}", "relu({m} + {w})", "{w} + {m}"])
+                )
+                m = f"{a[0]} @ {b[0]}"
+                if "{w}" in form:
+                    bias = f"w{nbias}"
+                    nbias += 1
+                    bdtype = draw(st.sampled_from(["i8", "i32"]))
+                    lines.append(f"input {bias} : {_fmt(bdtype, (shape[1],))}")
+                    pool.append((bias, bdtype, (shape[1],)))
+                else:
+                    bias = ""
+                expr = form.format(m=m, w=bias)
+        lines.append(f"v{j} = {expr}")
+        pool.append((f"v{j}", dtype, shape))
+
+    outs = [pool[-1][0]]
+    extra = draw(st.integers(0, len(pool) - 1))
+    if pool[extra][0] not in outs and draw(st.booleans()):
+        outs.append(pool[extra][0])
+    lines.append("output " + ", ".join(outs))
+    return "\n".join(lines) + "\n"

tests/golden/opt/big_mm.ir

@@ -0,0 +1,4 @@
+func @main(%A: i8[12,12], %B: i8[12,12]) -> (i32[12,12]) {
+  %1 = matmul{epilogue=relu} %A, %B : i32[12,12]  loc("big_mm.tiny":4:5)
+  return %1
+}

tests/golden/opt/const_mm.ir

@@ -0,0 +1,5 @@
+func @main(%X: i8[3,2]) -> (i32[3,2]) {
+  %0 = const [[1, -2], [3, 4]] : i8[2,2]        loc("const_mm.tiny":3:1)
+  %2 = matmul{epilogue=relu} %X, %0 : i32[3,2]  loc("const_mm.tiny":5:5)
+  return %2
+}

tests/golden/opt/mlp_gpu.ir

@@ -0,0 +1,5 @@
+func @main(%X: i8[4,8], %W1: i8[8,6], %b1: i32[6], %W2: i8[6,3]) -> (i32[4,3]) {
+  %2 = matmul{epilogue=bias,relu} %X, %W1, %b1 : i32[4,6]  loc("mlp_gpu.tiny":6:5)
+  %3 = matmul %2, %W2 : i32[4,3]                           loc("mlp_gpu.tiny":7:5)
+  return %3
+}

tests/golden/opt/mm.ir

@@ -0,0 +1,4 @@
+func @main(%A: i8[4,4], %B: i8[4,4]) -> (i32[4,4]) {
+  %0 = matmul %A, %B : i32[4,4]  loc("mm.tiny":4:5)
+  return %0
+}

tests/golden/opt/mm_bias_relu.ir

@@ -0,0 +1,4 @@
+func @main(%X: i8[4,8], %W: i8[8,4], %b: i32[4]) -> (i32[4,4]) {
+  %2 = matmul{epilogue=bias,relu} %X, %W, %b : i32[4,4]  loc("mm_bias_relu.tiny":5:5)
+  return %2
+}

tests/golden/opt/mm_relu.ir

@@ -0,0 +1,4 @@
+func @main(%A: i8[4,8], %B: i8[8,4]) -> (i32[4,4]) {
+  %1 = matmul{epilogue=relu} %A, %B : i32[4,4]  loc("mm_relu.tiny":4:5)
+  return %1
+}

tests/helpers.py

@@ -0,0 +1,24 @@
+"""Shared test helpers: run named passes over IR text and compare printed results."""
+
+from __future__ import annotations
+
+import textwrap
+
+from tinyc.ir.parser import parse_func
+from tinyc.ir.printer import print_func
+from tinyc.passes import REGISTRY, PassContext, PassManager
+
+
+def ir(text: str):
+    return parse_func(textwrap.dedent(text).strip() + "\n")
+
+
+def run_passes(text: str, *names: str, ctx: PassContext | None = None) -> str:
+    """Parse `text`, run the named passes with translation validation on, return printed IR."""
+    func = ir(text)
+    PassManager([REGISTRY[n]() for n in names], verify_each=True).run(func, ctx)
+    return print_func(func)
+
+
+def expected(text: str) -> str:
+    return textwrap.dedent(text).strip() + "\n"

tests/passes/test_manager.py

@@ -0,0 +1,107 @@
+import pytest
+
+from tests.helpers import ir
+from tinyc.driver import build_pipeline, compile_graph, pipeline_names
+from tinyc.ir.printer import print_func
+from tinyc.ir.types import TensorType
+from tinyc.passes import REGISTRY, PassContext, PassManager, PassVerifyError
+
+SRC = "input A : i8[2, 3]\ninput B : i8[3, 2]\nC = relu(A @ B)\noutput C\n"
+
+
+class DropRelu:
+    """A deliberately wrong pass: pretends relu is the identity."""
+
+    name = "drop-relu"
+
+    def run(self, func, ctx):
+        for op in func.ops:
+            if op.kind == "relu":
+                func.replace_uses(op.result, op.operands[0])
+
+
+class WrongType:
+    """A pass that corrupts a type: structurally invalid."""
+
+    name = "wrong-type"
+
+    def run(self, func, ctx):
+        func.ops[0].result.type = TensorType("i32", (7, 7))
+
+
+def func():
+    return ir("""
+        func @main(%A: i8[2,3], %B: i8[3,2]) -> (i32[2,2]) {
+          %0 = matmul %A, %B : i32[2,2]
+          %1 = relu %0 : i32[2,2]
+          return %1
+        }
+        """)
+
+
+def test_verify_each_catches_a_semantically_wrong_pass_and_shows_before_and_after():
+    with pytest.raises(PassVerifyError) as exc:
+        PassManager([DropRelu()], verify_each=True).run(func())
+    err = exc.value
+    assert err.pass_name == "drop-relu" and "differs from the oracle" in err.reason
+    assert "return %1" in err.before and "return %0" in err.after
+
+
+def test_without_verify_each_a_semantic_bug_goes_unnoticed():
+    PassManager([DropRelu()]).run(func())  # opt-in: structure is fine, so nothing complains
+
+
+def test_structural_verification_always_runs():
+    with pytest.raises(PassVerifyError, match="operands give"):
+        PassManager([WrongType()]).run(func())
+
+
+def test_report_records_changes_timing_and_ir():
+    f = func()
+    report = PassManager([REGISTRY["fuse"](), REGISTRY["dce"]()]).run(f)
+    assert [r.name for r in report.records] == ["fuse", "dce"]
+    assert [r.changed for r in report.records] == [True, False]
+    assert report.records[-1].ir == print_func(f) and "relu" in report.initial_ir
+    assert all(r.seconds >= 0 for r in report.records)
+    assert "fuse" in report.timing_table()
+
+
+def test_dump_callback_sees_input_then_every_pass():
+    seen = []
+    PassManager([REGISTRY["fuse"](), REGISTRY["dce"]()], dump=lambda n, t: seen.append(n)).run(
+        func()
+    )
+    assert seen == ["input", "fuse", "dce"]
+
+
+def test_pipeline_names_per_target_and_level():
+    assert pipeline_names(None, 0) == ["shape-infer", "range-check"]
+    assert pipeline_names("tpu", 0)[-1] == "legalize-tpu"
+    full = pipeline_names("gpu", 1)
+    assert full[0] == "shape-infer" and full[-2:] == ["range-check", "legalize-gpu"]
+    assert full.index("fuse") < full.index("range-check")
+
+
+def test_disabling_passes_and_unknown_names():
+    names = [p.name for p in build_pipeline("gpu", 1, disabled=["fuse", "cse"])]
+    assert "fuse" not in names and "cse" not in names and "dce" in names
+    with pytest.raises(ValueError, match="unknown passes"):
+        build_pipeline(None, 1, disabled=["nope"])
+
+
+def test_o0_does_not_fuse_but_o1_does():
+    f0, _ = compile_graph(SRC, "t.tiny", opt_level=0)
+    f1, _ = compile_graph(SRC, "t.tiny", opt_level=1)
+    assert [o.kind for o in f0.ops] == ["matmul", "relu"]
+    assert [(o.kind, o.attrs) for o in f1.ops] == [("matmul", {"epilogue": "relu"})]
+
+
+def test_fused_op_keeps_the_outermost_source_location():
+    f, _ = compile_graph(SRC, "t.tiny")
+    assert (f.ops[0].loc.line, f.ops[0].loc.col) == (3, 5)  # the relu( ... ) expression
+
+
+def test_a_pipeline_run_twice_is_a_no_op_the_second_time():
+    f, _ = compile_graph(SRC, "t.tiny", target="gpu")
+    again = PassManager(build_pipeline("gpu", 1), verify_each=True).run(f, PassContext())
+    assert not any(r.changed for r in again.records)