From 7216700b1dcf47d4a971d0ac48aa2183509eccb6 Mon Sep 17 00:00:00 2001 From: Dineth Hettiarachchi Date: Sun, 25 Oct 2026 00:08:10 +0530 Subject: [PATCH] Phase 2: pass manager and graph passes Add a pass manager with always-on structural verification, optional translation validation (--verify-each) against the numpy oracle, IR dumps and timing. Passes: shape-infer, canonicalize, const-fold, cse, dce, fuse (matmul with bias/relu epilogue), range-check (interval analysis proving i32 safety) and per-target legalization. Add fused-matmul support to the IR (attrs, printer, parser, verifier, oracle, interpreter), driver pipelines per target and -O level, CLI flags, docs, and tests: directed per-pass tests, manager tests, every example through every pipeline, and Hypothesis properties (semantics preserved, idempotent, range-check sound) over generated programs. --- CLAUDE.md | 3 +- CONTEXT.md | 7 +- README.md | 5 +- docs/IR.md | 5 + docs/PASSES.md | 50 ++++ tests/gen.py | 88 +++++++ tests/golden/opt/big_mm.ir | 4 + tests/golden/opt/const_mm.ir | 5 + tests/golden/opt/mlp_gpu.ir | 5 + tests/golden/opt/mm.ir | 4 + tests/golden/opt/mm_bias_relu.ir | 4 + tests/golden/opt/mm_relu.ir | 4 + tests/helpers.py | 24 ++ tests/passes/test_manager.py | 107 ++++++++ tests/passes/test_passes.py | 425 +++++++++++++++++++++++++++++++ tests/passes/test_pipeline.py | 117 +++++++++ tests/test_cli.py | 72 ++++++ tinyc/cli.py | 43 +++- tinyc/driver.py | 51 ++++ tinyc/ir/errors.py | 9 + tinyc/ir/graph.py | 14 + tinyc/ir/interp.py | 61 +++-- tinyc/ir/parser.py | 17 +- tinyc/ir/printer.py | 7 +- tinyc/ir/ref.py | 43 ++-- tinyc/ir/types.py | 28 +- tinyc/ir/verifier.py | 18 +- tinyc/passes/__init__.py | 21 ++ tinyc/passes/canonicalize.py | 28 ++ tinyc/passes/const_fold.py | 21 ++ tinyc/passes/cse.py | 25 ++ tinyc/passes/dce.py | 14 + tinyc/passes/fuse.py | 48 ++++ tinyc/passes/legalize_gpu.py | 26 ++ tinyc/passes/legalize_tpu.py | 31 +++ tinyc/passes/manager.py | 126 +++++++++ tinyc/passes/range_check.py | 73 ++++++ tinyc/passes/shape_infer.py | 17 ++ 38 files changed, 1592 insertions(+), 58 deletions(-) create mode 100644 docs/PASSES.md create mode 100644 tests/gen.py create mode 100644 tests/golden/opt/big_mm.ir create mode 100644 tests/golden/opt/const_mm.ir create mode 100644 tests/golden/opt/mlp_gpu.ir create mode 100644 tests/golden/opt/mm.ir create mode 100644 tests/golden/opt/mm_bias_relu.ir create mode 100644 tests/golden/opt/mm_relu.ir create mode 100644 tests/helpers.py create mode 100644 tests/passes/test_manager.py create mode 100644 tests/passes/test_passes.py create mode 100644 tests/passes/test_pipeline.py create mode 100644 tests/test_cli.py create mode 100644 tinyc/driver.py create mode 100644 tinyc/passes/__init__.py create mode 100644 tinyc/passes/canonicalize.py create mode 100644 tinyc/passes/const_fold.py create mode 100644 tinyc/passes/cse.py create mode 100644 tinyc/passes/dce.py create mode 100644 tinyc/passes/fuse.py create mode 100644 tinyc/passes/legalize_gpu.py create mode 100644 tinyc/passes/legalize_tpu.py create mode 100644 tinyc/passes/manager.py create mode 100644 tinyc/passes/range_check.py create mode 100644 tinyc/passes/shape_infer.py diff --git a/CLAUDE.md b/CLAUDE.md index 3df5844..49a4ea8 100644 --- a/CLAUDE.md +++ b/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 diff --git a/CONTEXT.md b/CONTEXT.md index fb9bcda..eaf8ed2 100644 --- a/CONTEXT.md +++ b/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) diff --git a/README.md b/README.md index c8ac802..d12e224 100644 --- a/README.md +++ b/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 && cd tiny-compiler diff --git a/docs/IR.md b/docs/IR.md index 44651f5..4b020de 100644 --- a/docs/IR.md +++ b/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. diff --git a/docs/PASSES.md b/docs/PASSES.md new file mode 100644 index 0000000..3256e2c --- /dev/null +++ b/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-] +-O0: shape-infer, range-check, [legalize-] +``` + +| 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). diff --git a/tests/gen.py b/tests/gen.py new file mode 100644 index 0000000..c1a42f3 --- /dev/null +++ b/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" diff --git a/tests/golden/opt/big_mm.ir b/tests/golden/opt/big_mm.ir new file mode 100644 index 0000000..fdd4925 --- /dev/null +++ b/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 +} diff --git a/tests/golden/opt/const_mm.ir b/tests/golden/opt/const_mm.ir new file mode 100644 index 0000000..22f1232 --- /dev/null +++ b/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 +} diff --git a/tests/golden/opt/mlp_gpu.ir b/tests/golden/opt/mlp_gpu.ir new file mode 100644 index 0000000..f80ba7d --- /dev/null +++ b/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 +} diff --git a/tests/golden/opt/mm.ir b/tests/golden/opt/mm.ir new file mode 100644 index 0000000..7c208c5 --- /dev/null +++ b/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 +} diff --git a/tests/golden/opt/mm_bias_relu.ir b/tests/golden/opt/mm_bias_relu.ir new file mode 100644 index 0000000..dbe03db --- /dev/null +++ b/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 +} diff --git a/tests/golden/opt/mm_relu.ir b/tests/golden/opt/mm_relu.ir new file mode 100644 index 0000000..0cc9292 --- /dev/null +++ b/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 +} diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..8527c66 --- /dev/null +++ b/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" diff --git a/tests/passes/test_manager.py b/tests/passes/test_manager.py new file mode 100644 index 0000000..695ba45 --- /dev/null +++ b/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) diff --git a/tests/passes/test_passes.py b/tests/passes/test_passes.py new file mode 100644 index 0000000..20f35e9 --- /dev/null +++ b/tests/passes/test_passes.py @@ -0,0 +1,425 @@ +"""Directed tests: IR text in, expected IR text out, one group per pass.""" + +import pytest + +from tests.helpers import expected, ir, run_passes +from tinyc.ir.errors import PassError +from tinyc.passes import PassContext, PassManager +from tinyc.passes.range_check import RangeCheck + +# --- canonicalize ----------------------------------------------------------------------------- + + +def test_canonicalize_puts_the_matrix_first_in_a_broadcast_add(): + out = run_passes( + """ + func @main(%M: i8[2,3], %b: i8[3]) -> (i32[2,3]) { + %0 = add %b, %M : i32[2,3] + return %0 + } + """, + "canonicalize", + ) + assert "add %M, %b" in out + + +def test_canonicalize_orders_same_shape_adds_by_definition_order(): + out = run_passes( + """ + func @main(%x: i8[2], %y: i8[2]) -> (i32[2]) { + %0 = add %y, %x : i32[2] + return %0 + } + """, + "canonicalize", + ) + assert "add %x, %y" in out + + +def test_canonicalize_collapses_double_relu(): + out = run_passes( + """ + func @main(%A: i8[2]) -> (i8[2]) { + %0 = relu %A : i8[2] + %1 = relu %0 : i8[2] + return %1 + } + """, + "canonicalize", + "dce", + ) + assert out == expected(""" + func @main(%A: i8[2]) -> (i8[2]) { + %0 = relu %A : i8[2] + return %0 + } + """) + + +# --- const-fold ------------------------------------------------------------------------------- + + +def test_const_fold_evaluates_constant_chains_and_dce_cleans_up(): + out = run_passes( + """ + func @main(%A: i8[2,2]) -> (i32[2,2]) { + %0 = const [[1, 2], [3, 4]] : i8[2,2] + %1 = const [[1, 0], [0, 1]] : i8[2,2] + %2 = matmul %0, %1 : i32[2,2] + %3 = add %2, %2 : i32[2,2] + return %3 + } + """, + "const-fold", + "dce", + ) + assert out == expected(""" + func @main(%A: i8[2,2]) -> (i32[2,2]) { + %3 = const [[2, 4], [6, 8]] : i32[2,2] + return %3 + } + """) + + +def test_const_fold_folds_relu_of_a_constant_and_keeps_its_dtype(): + out = run_passes( + """ + func @main() -> (i8[3]) { + %0 = const [-5, 0, 7] : i8[3] + %1 = relu %0 : i8[3] + return %1 + } + """, + "const-fold", + ) + assert "%1 = const [0, 0, 7] : i8[3]" in out + + +def test_const_fold_leaves_ops_with_a_non_constant_operand(): + text = """ + func @main(%A: i8[2]) -> (i32[2]) { + %0 = const [1, 2] : i8[2] + %1 = add %A, %0 : i32[2] + return %1 + } + """ + assert run_passes(text, "const-fold") == expected(text) + + +def test_const_fold_reports_overflowing_constant_expressions(): + big = ", ".join(["2000000000"] * 2) + f = ir(f""" + func @main() -> (i32[2]) {{ + %0 = const [{big}] : i32[2] + %1 = add %0, %0 : i32[2] + return %1 + }} + """) + from tinyc.passes import REGISTRY + + with pytest.raises(PassError, match="overflows i32"): + REGISTRY["const-fold"]().run(f, PassContext()) + + +# --- cse -------------------------------------------------------------------------------------- + + +def test_cse_merges_identical_ops(): + out = run_passes( + """ + func @main(%A: i8[2,2], %B: i8[2,2]) -> (i32[2,2]) { + %0 = matmul %A, %B : i32[2,2] + %1 = matmul %A, %B : i32[2,2] + %2 = add %0, %1 : i32[2,2] + return %2 + } + """, + "cse", + ) + assert out == expected(""" + func @main(%A: i8[2,2], %B: i8[2,2]) -> (i32[2,2]) { + %0 = matmul %A, %B : i32[2,2] + %2 = add %0, %0 : i32[2,2] + return %2 + } + """) + + +def test_cse_merges_equal_constants_but_not_different_dtypes(): + out = run_passes( + """ + func @main() -> (i8[2], i8[2], i32[2]) { + %0 = const [1, 2] : i8[2] + %1 = const [1, 2] : i8[2] + %2 = const [1, 2] : i32[2] + return %0, %1, %2 + } + """, + "cse", + ) + assert out.count("const") == 2 and "return %0, %0, %2" in out + + +def test_canonicalize_then_cse_sees_through_commutativity(): + out = run_passes( + """ + func @main(%x: i8[2], %y: i8[2]) -> (i32[2]) { + %0 = add %x, %y : i32[2] + %1 = add %y, %x : i32[2] + %2 = add %0, %1 : i32[2] + return %2 + } + """, + "canonicalize", + "cse", + ) + assert "%2 = add %0, %0" in out and "%1" not in out + + +# --- dce -------------------------------------------------------------------------------------- + + +def test_dce_removes_dead_ops_and_keeps_unused_inputs(): + out = run_passes( + """ + func @main(%A: i8[2], %B: i8[2]) -> (i8[2]) { + %0 = relu %A : i8[2] + %1 = add %A, %B : i32[2] + %2 = relu %1 : i32[2] + return %0 + } + """, + "dce", + ) + assert out == expected(""" + func @main(%A: i8[2], %B: i8[2]) -> (i8[2]) { + %0 = relu %A : i8[2] + return %0 + } + """) + + +# --- fuse ------------------------------------------------------------------------------------- + +MM_BIAS_RELU = """ + func @main(%A: i8[2,3], %B: i8[3,2], %b: i32[2]) -> (i32[2,2]) { + %0 = matmul %A, %B : i32[2,2] + %1 = add %0, %b : i32[2,2] + %2 = relu %1 : i32[2,2] + return %2 + } + """ + + +def test_fuse_matmul_bias_relu(): + assert run_passes(MM_BIAS_RELU, "fuse") == expected(""" + func @main(%A: i8[2,3], %B: i8[3,2], %b: i32[2]) -> (i32[2,2]) { + %2 = matmul{epilogue=bias,relu} %A, %B, %b : i32[2,2] + return %2 + } + """) + + +def test_fuse_matmul_relu_and_matmul_bias_alone(): + out = run_passes( + """ + 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 + } + """, + "fuse", + ) + assert "%1 = matmul{epilogue=relu} %A, %B : i32[2,2]" in out + out = run_passes( + """ + func @main(%A: i8[2,3], %B: i8[3,2], %b: i8[2]) -> (i32[2,2]) { + %0 = matmul %A, %B : i32[2,2] + %1 = add %b, %0 : i32[2,2] + return %1 + } + """, + "fuse", + ) + assert "%1 = matmul{epilogue=bias} %A, %B, %b : i32[2,2]" in out # either operand order + + +def test_fuse_moves_the_op_after_a_late_defined_bias(): + out = run_passes( + """ + func @main(%A: i8[2,3], %B: i8[3,2]) -> (i32[2,2]) { + %0 = matmul %A, %B : i32[2,2] + %1 = const [1, 2] : i32[2] + %2 = add %0, %1 : i32[2,2] + return %2 + } + """, + "fuse", + ) + assert out == expected(""" + func @main(%A: i8[2,3], %B: i8[3,2]) -> (i32[2,2]) { + %1 = const [1, 2] : i32[2] + %2 = matmul{epilogue=bias} %A, %B, %1 : i32[2,2] + return %2 + } + """) + + +@pytest.mark.parametrize( + "body", + [ + # the matmul result has two users + "%0 = matmul %A, %B : i32[2,2]\n%1 = relu %0 : i32[2,2]\n%2 = add %0, %1 : i32[2,2]\nreturn %2", + # the matmul result is also an output + "%0 = matmul %A, %B : i32[2,2]\n%1 = relu %0 : i32[2,2]\nreturn %0, %1", + # add with a full matrix is not a bias + "%0 = matmul %A, %B : i32[2,2]\n%1 = add %0, %0 : i32[2,2]\nreturn %1", + ], +) +def test_fuse_leaves_unfusable_patterns_alone(body): + rets = "i32[2,2], i32[2,2]" if "return %0, %1" in body else "i32[2,2]" + text = f"func @main(%A: i8[2,3], %B: i8[3,2]) -> ({rets}) {{\n{body}\n}}" + assert run_passes(text, "fuse") == run_passes(text) # unchanged + + +def test_fuse_is_idempotent(): + once = run_passes(MM_BIAS_RELU, "fuse") + assert run_passes(once, "fuse") == once + + +# --- range-check ------------------------------------------------------------------------------ + + +def _mm(k, a="i8", b="i8"): + return f""" + func @main(%A: {a}[1,{k}], %B: {b}[{k},1]) -> (i32[1,1]) {{ + %0 = matmul %A, %B : i32[1,1] + return %0 + }} + """ + + +def test_range_check_computes_exact_interval_bounds(): + ctx = PassContext() + f = ir(_mm(4)) + PassManager([RangeCheck()]).run(f, ctx) + # i8 x i8 products lie in [-128*127, 128*128] = [-16256, 16384]; four of them: + assert ctx.analysis["ranges"]["0"] == (-65024, 65536) + assert ctx.analysis["ranges"]["A"] == (-128, 127) + + +def test_range_check_boundary_is_exact(): + PassManager([RangeCheck()]).run(ir(_mm(131071)), PassContext()) # 16384 * 131071 < 2**31 + with pytest.raises(PassError, match="does not fit in i32") as exc: + PassManager([RangeCheck()]).run(ir(_mm(131072)), PassContext()) # exactly 2**31 + assert "matmul" in str(exc.value) + + +def test_range_check_assumes_a_range_for_i32_inputs_and_says_so(): + with pytest.raises(PassError, match="i32 inputs are assumed"): + PassManager([RangeCheck()]).run(ir(_mm(4, "i32", "i32")), PassContext()) + wide = PassContext(i32_input_range=(-1000, 1000)) + PassManager([RangeCheck()]).run(ir(_mm(4, "i32", "i32")), wide) # 4 * 1000 * 1000 fits + + +def test_range_check_tracks_relu_bias_and_constants_exactly(): + ctx = PassContext() + f = ir(""" + func @main(%A: i8[1,2]) -> (i32[1,2]) { + %0 = const [[100, -3], [2, 5]] : i8[2,2] + %1 = matmul{epilogue=relu} %A, %0 : i32[1,2] + return %1 + } + """) + PassManager([RangeCheck()]).run(f, ctx) + ranges = ctx.analysis["ranges"] + assert ranges["0"] == (-3, 100) # a constant is bounded by its own min and max + # products of [-128,127] x [-3,100] lie in [-12800, 12700]; two of them, then relu: + assert ranges["1"] == (0, 25400) + + +def test_range_check_rejects_overflowing_add_chain(): + f = ir(""" + func @main(%a: i32[1], %b: i32[1]) -> (i32[1]) { + %0 = add %a, %b : i32[1] + return %0 + } + """) + PassManager([RangeCheck()]).run(f, PassContext()) # 2 * 2**20 fits + with pytest.raises(PassError, match="add"): + PassManager([RangeCheck()]).run(f, PassContext(i32_input_range=(-(2**31), 2**31 - 1))) + + +# --- legalization ----------------------------------------------------------------------------- + + +def test_tpu_rejects_i32_operands_with_a_clear_message(): + f = ir(""" + func @main(%A: i8[2,2], %B: i8[2,2], %D: i8[2,2]) -> (i32[2,2]) { + %0 = matmul %A, %B : i32[2,2] + %1 = matmul %0, %D : i32[2,2] + return %1 + } + """) + from tinyc.passes import REGISTRY + + with pytest.raises(PassError, match=r"tiny-tpu needs i8 matmul operands, but %0 is i32\[2,2\]"): + REGISTRY["legalize-tpu"]().run(f, PassContext()) + REGISTRY["legalize-gpu"]().run(f, PassContext()) # the GPU runs chained matmuls + + +def test_tpu_rejects_computed_i8_operands_but_accepts_constants(): + from tinyc.passes import REGISTRY + + computed = ir(""" + func @main(%A: i8[2,2], %B: i8[2,2]) -> (i32[2,2]) { + %0 = relu %A : i8[2,2] + %1 = matmul %0, %B : i32[2,2] + return %1 + } + """) + with pytest.raises(PassError, match="computed by 'relu'"): + REGISTRY["legalize-tpu"]().run(computed, PassContext()) + const = ir(""" + func @main(%A: i8[2,2]) -> (i32[2,2]) { + %0 = const [[1, 2], [3, 4]] : i8[2,2] + %1 = matmul %A, %0 : i32[2,2] + return %1 + } + """) + REGISTRY["legalize-tpu"]().run(const, PassContext()) + + +def test_gpu_limits_on_threads_and_memory(): + from tinyc.passes import REGISTRY + + REGISTRY["legalize-gpu"]().run(ir(_mm(4)), PassContext()) + square = ir(""" + func @main(%A: i8[3,2], %B: i8[2,3]) -> (i32[3,3]) { + %0 = matmul %A, %B : i32[3,3] + return %0 + } + """) + with pytest.raises(PassError, match="9 threads"): + REGISTRY["legalize-gpu"]().run(square, PassContext(gpu_max_threads=8)) + with pytest.raises(PassError, match="words of tiny-gpu memory"): + REGISTRY["legalize-gpu"]().run(square, PassContext(gpu_mem_words=20)) + + +# --- shape-infer ------------------------------------------------------------------------------ + + +def test_shape_infer_repairs_stale_types(): + f = ir(""" + func @main(%A: i8[2,3], %B: i8[3,2]) -> (i32[2,2]) { + %0 = matmul %A, %B : i32[2,2] + return %0 + } + """) + from tinyc.ir.types import TensorType + from tinyc.passes import REGISTRY + + f.ops[0].result.type = TensorType("i32", (9, 9)) + REGISTRY["shape-infer"]().run(f, PassContext()) + assert str(f.ops[0].result.type) == "i32[2,2]" diff --git a/tests/passes/test_pipeline.py b/tests/passes/test_pipeline.py new file mode 100644 index 0000000..e34339b --- /dev/null +++ b/tests/passes/test_pipeline.py @@ -0,0 +1,117 @@ +"""The whole pipeline over every example, and over random programs.""" + +from pathlib import Path + +import numpy as np +import pytest +from hypothesis import HealthCheck, event, given, settings +from hypothesis import strategies as st + +from tests.gen import programs +from tinyc.driver import build_pipeline, compile_graph +from tinyc.frontend import compile_to_graph +from tinyc.ir import ref +from tinyc.ir.errors import PassError, RangeError +from tinyc.ir.parser import parse_func +from tinyc.ir.printer import print_func +from tinyc.passes import PassContext, PassManager + +ROOT = Path(__file__).resolve().parent.parent.parent +EXAMPLES = sorted((ROOT / "examples").glob("*.tiny")) +TPU_ILLEGAL = {"mlp_gpu"} # its second matmul has an i32 operand + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +@pytest.mark.parametrize("target", [None, "gpu", "tpu"]) +def test_examples_through_every_pipeline_with_verify_each(path, opt, target): + if target == "tpu" and path.stem in TPU_ILLEGAL: + with pytest.raises(PassError, match="tiny-tpu needs i8"): + compile_graph(path.read_text(), path.name, target=target, opt_level=opt) + return + func, _ = compile_graph( + path.read_text(), path.name, target=target, opt_level=opt, verify_each=True + ) + # the optimized program still computes what the source computes + source = compile_to_graph(path.read_text(), path.name) + rng = np.random.default_rng(7) + for _ in range(5): + inputs = ref.random_inputs(source, rng) + for got, want in zip(ref.evaluate(func, inputs), ref.evaluate(source, inputs)): + np.testing.assert_array_equal(got, want) + # and the optimized IR survives a print/parse round trip, attributes included + text = print_func(func, locs=True) + assert print_func(parse_func(text), locs=True) == text + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_optimized_ir_golden(path, golden): + func, _ = compile_graph(path.read_text(), path.name) + golden(ROOT / "tests" / "golden" / "opt" / f"{path.stem}.ir", print_func(func, locs=True)) + + +def test_dense_layer_fuses_completely_and_mlp_fuses_the_first_layer(): + func, _ = compile_graph((ROOT / "examples" / "mm_bias_relu.tiny").read_text(), "x.tiny") + assert [(o.kind, o.attrs.get("epilogue")) for o in func.ops] == [("matmul", "bias,relu")] + func, _ = compile_graph((ROOT / "examples" / "mlp_gpu.tiny").read_text(), "x.tiny") + assert [(o.kind, o.attrs.get("epilogue")) for o in func.ops] == [ + ("matmul", "bias,relu"), + ("matmul", None), + ] + + +# --- properties over random programs ---------------------------------------------------------- + + +def _try(source: str, **kw): + try: + return compile_graph(source, "gen.tiny", **kw) + except PassError: # range_check proved it could overflow: a legitimate rejection + event("rejected by range-check") + return None + + +@settings(max_examples=200, deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given(programs()) +def test_pipeline_preserves_semantics_on_random_programs(source): + # verify_each re-runs the program after *every* pass and compares with the source's oracle + _try(source, verify_each=True) + + +@settings(max_examples=200, deadline=None) +@given(programs()) +def test_pipeline_is_idempotent(source): + got = _try(source) + if got is None: + return + func, _ = got + again = PassManager(build_pipeline(None, 1), verify_each=True).run(func, PassContext()) + assert not any(r.changed for r in again.records), [r.name for r in again.records if r.changed] + + +@settings(max_examples=200, deadline=None) +@given(programs(), st.integers(0, 2**32 - 1)) +def test_range_check_is_sound_on_extreme_inputs(source, seed): + """If range-check accepts a program, no input in the assumed ranges can overflow, and every + value stays inside its proven interval. Inputs are drawn from the interval corners.""" + ctx = PassContext() + got = _try(source, opt_level=0, ctx=ctx) + if got is None: + return + func, _ = got + rng = np.random.default_rng(seed) + inputs = {} + for a in func.args: + lo, hi = (-128, 127) if a.type.dtype == "i8" else ctx.i32_input_range + inputs[a.name] = rng.choice([lo, hi], size=a.type.shape).astype(np.int64) + env = dict(inputs) + for op in func.ops: + try: + env[op.result.name] = ref.eval_op( + op.kind, [env[v.name] for v in op.operands], op.attrs, op.data, op.result.name + ) + except RangeError as e: # pragma: no cover - would be a soundness bug + pytest.fail(f"accepted program overflowed: {e}\n{source}") + for name, (lo, hi) in ctx.analysis["ranges"].items(): + arr = env[name] + assert lo <= arr.min() and arr.max() <= hi, f"%{name} escaped its interval\n{source}" diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..c030a4f --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,72 @@ +from pathlib import Path + +import pytest + +from tinyc.cli import main + +EX = Path(__file__).resolve().parent.parent / "examples" + + +def run(capsys, *args): + code = main([*map(str, args)]) + out = capsys.readouterr() + return code, out.out, out.err + + +def test_default_run_prints_optimized_ir(capsys): + code, out, _ = run(capsys, EX / "mm_bias_relu.tiny") + assert code == 0 and "matmul{epilogue=bias,relu}" in out + + +def test_o0_and_disable_keep_ops_separate(capsys): + for flags in (["-O", "0"], ["--disable", "fuse"]): + code, out, _ = run(capsys, EX / "mm_relu.tiny", *flags) + assert code == 0 and "epilogue" not in out and "relu %0" in out + + +def test_target_tpu_rejects_chained_matmul_with_a_caret(capsys): + code, out, err = run(capsys, EX / "mlp_gpu.tiny", "--target", "tpu") + assert code == 1 and out == "" + assert "mlp_gpu.tiny:7:5: error: tiny-tpu needs i8 matmul operands" in err + assert "7 | Y = H @ W2" in err and "^" in err + + +def test_target_gpu_accepts_it(capsys): + assert run(capsys, EX / "mlp_gpu.tiny", "--target", "gpu")[0] == 0 + + +def test_print_after_and_time_passes_go_to_stderr(capsys): + code, out, err = run(capsys, EX / "mm_relu.tiny", "--print-after", "fuse", "--time-passes") + assert code == 0 + assert "// ---- after fuse ----" in err and "fuse " in err and "ms" in err + assert "// ----" not in out + + +def test_print_after_all_lists_every_pass(capsys): + _, _, err = run(capsys, EX / "mm_relu.tiny", "--print-after-all", "--target", "gpu") + for name in ("input", "shape-infer", "cse", "fuse", "range-check", "legalize-gpu"): + assert f"// ---- after {name} ----" in err + + +def test_verify_each_flag(capsys): + assert run(capsys, EX / "big_mm.tiny", "--verify-each", "--target", "tpu")[0] == 0 + + +def test_unknown_disable_is_a_usage_error(capsys): + with pytest.raises(SystemExit) as exc: + main([str(EX / "mm.tiny"), "--disable", "nope"]) + assert exc.value.code == 2 + + +def test_range_failure_is_reported_at_the_source_line(tmp_path, capsys): + src = tmp_path / "big.tiny" + src.write_text("input a : i32[1, 4]\ninput b : i32[4, 1]\nC = a @ b\noutput C\n") + code, _, err = run(capsys, src) + assert code == 1 and "big.tiny:3:5: error:" in err and "does not fit in i32" in err + + +def test_front_end_errors_still_work(tmp_path, capsys): + src = tmp_path / "bad.tiny" + src.write_text("C = nope\n") + code, _, err = run(capsys, src) + assert code == 1 and "undefined name 'nope'" in err diff --git a/tinyc/cli.py b/tinyc/cli.py index 1677a9d..2066e4d 100644 --- a/tinyc/cli.py +++ b/tinyc/cli.py @@ -1,4 +1,4 @@ -"""`tinyc FILE.tiny [--emit ast|ir] [--locs]`: print the AST or Graph IR, or a caret diagnostic.""" +"""`tinyc FILE.tiny`: compile to Graph IR (optionally through the pass pipeline) and print it.""" from __future__ import annotations @@ -6,10 +6,14 @@ import sys from pathlib import Path +from tinyc.driver import build_pipeline from tinyc.frontend import CompileError, SourceFile, lower, parse from tinyc.frontend.ast import dump +from tinyc.frontend.diagnostics import Diagnostic, Span +from tinyc.ir.errors import PassError from tinyc.ir.printer import print_func from tinyc.ir.verifier import verify +from tinyc.passes import REGISTRY, PassManager, PassVerifyError def main(argv: list[str] | None = None) -> int: @@ -17,8 +21,25 @@ def main(argv: list[str] | None = None) -> int: ap.add_argument("file", type=Path) ap.add_argument("--emit", choices=["ast", "ir"], default="ir") ap.add_argument("--locs", action="store_true", help="print source locations in the IR") + ap.add_argument("--target", choices=["gpu", "tpu"], help="also run that target's legality pass") + ap.add_argument( + "-O", dest="opt", type=int, choices=[0, 1], default=1, help="optimization level" + ) + ap.add_argument("--disable", action="append", default=[], metavar="PASS", + help=f"skip a pass ({', '.join(REGISTRY)})") # fmt: skip + ap.add_argument("--print-after-all", action="store_true", help="dump the IR after every pass") + ap.add_argument("--print-after", action="append", default=[], metavar="PASS") + ap.add_argument("--verify-each", action="store_true", + help="run the IR after every pass and compare with the numpy oracle") # fmt: skip + ap.add_argument("--time-passes", action="store_true") args = ap.parse_args(argv) + source = SourceFile(args.file.name, args.file.read_text()) + + def dump_ir(name: str, text: str) -> None: + if args.print_after_all or name in args.print_after: + print(f"// ---- after {name} ----\n{text}", file=sys.stderr, end="") + try: program = parse(source) if args.emit == "ast": @@ -26,9 +47,29 @@ def main(argv: list[str] | None = None) -> int: return 0 func = lower(program, source) verify(func) + manager = PassManager( + build_pipeline(args.target, args.opt, args.disable), + verify_each=args.verify_each, + dump=dump_ir, + ) + report = manager.run(func) except CompileError as err: print(err.render(), file=sys.stderr) return 1 + except PassError as err: + if err.loc is None: + print(f"{source.name}: error: {err.message}", file=sys.stderr) + else: + span = Span(err.loc.line, err.loc.col, err.loc.line, err.loc.col + 1) + print(CompileError(source, Diagnostic(err.message, span)).render(), file=sys.stderr) + return 1 + except PassVerifyError as err: + print(f"internal compiler error: {err}", file=sys.stderr) + return 3 + except ValueError as err: # unknown --disable name + ap.error(str(err)) + if args.time_passes: + print(report.timing_table(), file=sys.stderr) print(print_func(func, locs=args.locs), end="") return 0 diff --git a/tinyc/driver.py b/tinyc/driver.py new file mode 100644 index 0000000..4ab9429 --- /dev/null +++ b/tinyc/driver.py @@ -0,0 +1,51 @@ +"""Pipelines: which passes run, in what order, for a target and optimization level.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable + +from tinyc.frontend import compile_to_graph +from tinyc.ir.graph import Func +from tinyc.passes import REGISTRY, PassContext, PassManager, PassReport + +OPTIMIZATIONS = ("canonicalize", "const-fold", "cse", "dce", "fuse") + + +def pipeline_names(target: str | None = None, opt_level: int = 1) -> list[str]: + """-O0 only type-checks and proves ranges; -O1 adds the optimizations.""" + names = ["shape-infer"] + if opt_level >= 1: + names += ["canonicalize", "const-fold", "canonicalize", "cse", "dce", "fuse"] + names.append("range-check") + if target in ("tpu", "gpu"): + names.append(f"legalize-{target}") + return names + + +def build_pipeline( + target: str | None = None, opt_level: int = 1, disabled: Iterable[str] = () +) -> list: + off = set(disabled) + unknown = off - set(REGISTRY) + if unknown: + raise ValueError(f"unknown passes: {sorted(unknown)} (known: {sorted(REGISTRY)})") + return [REGISTRY[n]() for n in pipeline_names(target, opt_level) if n not in off] + + +def compile_graph( + text: str, + filename: str = "", + *, + target: str | None = None, + opt_level: int = 1, + disabled: Iterable[str] = (), + verify_each: bool = False, + ctx: PassContext | None = None, + dump: Callable[[str, str], None] | None = None, +) -> tuple[Func, PassReport]: + """Source text -> optimized, legalized Graph IR (the input to the target backends).""" + func = compile_to_graph(text, filename) + manager = PassManager( + build_pipeline(target, opt_level, disabled), verify_each=verify_each, dump=dump + ) + return func, manager.run(func, ctx) diff --git a/tinyc/ir/errors.py b/tinyc/ir/errors.py index 2ee1bd3..3f259c5 100644 --- a/tinyc/ir/errors.py +++ b/tinyc/ir/errors.py @@ -19,3 +19,12 @@ class InputError(ValueError): class RangeError(ArithmeticError): """A value does not fit in i32. tiny-gpu raises on this, so the semantics reject it.""" + + +class PassError(Exception): + """A pass rejected the program (e.g. a target cannot run it). `loc` points at the source.""" + + def __init__(self, message: str, loc=None): + self.message = message + self.loc = loc + super().__init__(f"{loc}: {message}" if loc else message) diff --git a/tinyc/ir/graph.py b/tinyc/ir/graph.py index 34d3efc..ad33675 100644 --- a/tinyc/ir/graph.py +++ b/tinyc/ir/graph.py @@ -27,6 +27,7 @@ class Op: result: Value loc: Loc | None = None data: np.ndarray | None = None # const only (int64) + attrs: dict[str, str] = field(default_factory=dict) # matmul only: {"epilogue": "bias,relu"} @dataclass(eq=False) @@ -36,6 +37,19 @@ class Func: ops: list[Op] = field(default_factory=list) outputs: list[Value] = field(default_factory=list) + def replace_uses(self, old: Value, new: Value) -> None: + for op in self.ops: + op.operands = [new if v is old else v for v in op.operands] + self.outputs = [new if v is old else v for v in self.outputs] + + def users(self) -> dict[str, list[Op]]: + """value name -> ops using it (an op appears once per operand slot it uses it in).""" + out: dict[str, list[Op]] = {} + for op in self.ops: + for v in op.operands: + out.setdefault(v.name, []).append(op) + return out + class FuncBuilder: """Builds a Func, inferring result types and naming results %0, %1, ...""" diff --git a/tinyc/ir/interp.py b/tinyc/ir/interp.py index 0e3ec47..e52d83a 100644 --- a/tinyc/ir/interp.py +++ b/tinyc/ir/interp.py @@ -14,46 +14,57 @@ from tinyc.ir.errors import RangeError from tinyc.ir.graph import Func from tinyc.ir.ref import I32, validate_inputs +from tinyc.ir.types import epilogue def _fits(v: int) -> bool: return I32[0] <= v <= I32[1] -def _rows(x): # rank-1 -> a single row +def _rows(x): # a rank-1 value is a single row return x if isinstance(x[0], list) else [x] +def _check(r, name: str, what: str): + if not all(_fits(v) for row in _rows(r) for v in row): + raise RangeError(f"%{name} ({what}) does not fit in i32") + return r + + +def _matmul(a, b): + return [[sum(a[i][k] * b[k][j] for k in range(len(b))) for j in range(len(b[0]))] + for i in range(len(a))] # fmt: skip + + +def _add(a, b): # same shape, or a rank-1 operand broadcast over the rows of a rank-2 one + if isinstance(a[0], list) and not isinstance(b[0], list): + return [[x + y for x, y in zip(row, b)] for row in a] + if isinstance(b[0], list) and not isinstance(a[0], list): + return [[x + y for x, y in zip(a, row)] for row in b] + if isinstance(a[0], list): + return [[x + y for x, y in zip(ra, rb)] for ra, rb in zip(a, b)] + return [x + y for x, y in zip(a, b)] + + +def _relu(x): + if isinstance(x[0], list): + return [[max(v, 0) for v in row] for row in x] + return [max(v, 0) for v in x] + + def run(func: Func, inputs: Mapping[str, object]) -> list[np.ndarray]: env = {k: v.tolist() for k, v in validate_inputs(func, inputs).items()} for op in func.ops: - name = op.result.name + name, xs = op.result.name, [env[v.name] for v in op.operands] if op.kind == "const": r = op.data.tolist() elif op.kind == "relu": - x = env[op.operands[0].name] - r = ( - [[max(v, 0) for v in row] for row in x] - if isinstance(x[0], list) - else [max(v, 0) for v in x] - ) + r = _relu(xs[0]) elif op.kind == "matmul": - a, b = (env[v.name] for v in op.operands) - r = [ - [sum(a[i][k] * b[k][j] for k in range(len(b))) for j in range(len(b[0]))] - for i in range(len(a)) - ] - else: # add, with row broadcast of a rank-1 operand - a, b = (env[v.name] for v in op.operands) - if isinstance(a[0], list) and not isinstance(b[0], list): - r = [[x + y for x, y in zip(row, b)] for row in a] - elif isinstance(b[0], list) and not isinstance(a[0], list): - r = [[x + y for x, y in zip(a, row)] for row in b] - elif isinstance(a[0], list): - r = [[x + y for x, y in zip(ra, rb)] for ra, rb in zip(a, b)] - else: - r = [x + y for x, y in zip(a, b)] - if not all(_fits(v) for row in _rows(r) for v in row): - raise RangeError(f"%{name} ({op.kind}) does not fit in i32") + r = _check(_matmul(xs[0], xs[1]), name, "matmul") + for stage in epilogue(op.attrs): + r = _check(_add(r, xs[2]), name, "matmul bias add") if stage == "bias" else _relu(r) + else: + r = _check(_add(xs[0], xs[1]), name, "add") env[name] = r return [np.array(env[v.name], dtype=np.int64) for v in func.outputs] diff --git a/tinyc/ir/parser.py b/tinyc/ir/parser.py index eeea3f7..ade9741 100644 --- a/tinyc/ir/parser.py +++ b/tinyc/ir/parser.py @@ -118,6 +118,20 @@ def loc(self) -> Loc | None: self.take("punct", ")") return Loc(file, line, col) + def attrs(self) -> dict[str, str]: + """`{key=word,word,...}` directly after the op kind (only `epilogue` exists today).""" + if not self.at("punct", "{"): + return {} + self.pos += 1 + key = self.take("word").text + self.take("punct", "=") + words = [self.take("word").text] + while self.at("punct", ","): + self.pos += 1 + words.append(self.take("word").text) + self.take("punct", "}") + return {key: ",".join(words)} + def op(self) -> Op: res_tok = self.take("val") self.take("punct", "=") @@ -125,6 +139,7 @@ def op(self) -> Op: kind = kind_tok.text if kind not in OP_ARITY: raise self.error(f"unknown op '{kind}'", kind_tok.span) + attrs = self.attrs() operands, data, lit_span = [], None, self.tok.span if kind == "const": lit = self.literal() @@ -144,7 +159,7 @@ def op(self) -> Op: raise self.error( f"const data has shape {list(data.shape)}, type says {ty}", lit_span ) - return Op(kind, operands, self.define(res_tok, ty), self.loc(), data) + return Op(kind, operands, self.define(res_tok, ty), self.loc(), data, attrs) def func(self) -> Func: self.take("word", "func") diff --git a/tinyc/ir/printer.py b/tinyc/ir/printer.py index 816e74f..3d55276 100644 --- a/tinyc/ir/printer.py +++ b/tinyc/ir/printer.py @@ -10,7 +10,12 @@ def _op_text(op: Op) -> str: assert op.data is not None body = f"const {op.data.tolist()!r}" else: - body = f"{op.kind} " + ", ".join(f"%{v.name}" for v in op.operands) + attrs = ( + "{" + ",".join(f"{k}={v}" for k, v in sorted(op.attrs.items())) + "}" + if op.attrs + else "" + ) + body = f"{op.kind}{attrs} " + ", ".join(f"%{v.name}" for v in op.operands) return f"%{op.result.name} = {body} : {op.result.type}" diff --git a/tinyc/ir/ref.py b/tinyc/ir/ref.py index 628576e..b66f8be 100644 --- a/tinyc/ir/ref.py +++ b/tinyc/ir/ref.py @@ -2,8 +2,9 @@ Semantics (shared with `ir.interp`, which is an independent scalar implementation): * inputs must fit their declared dtype; every op computes in int64 on widened operands; - * every op result must fit in i32, otherwise RangeError (tiny-gpu raises on overflow). - Only final results are checked here; Phase 2's range analysis bounds partial sums as well. + * every op result must fit in i32, otherwise RangeError (tiny-gpu raises on overflow); + * a fused matmul (`epilogue=bias,relu`) checks each stage's result, exactly like the unfused ops. + Partial sums are not checked here; `passes.range_check` bounds them statically. """ from __future__ import annotations @@ -14,7 +15,7 @@ from tinyc.ir.errors import InputError, RangeError from tinyc.ir.graph import Func -from tinyc.ir.types import RANGES +from tinyc.ir.types import RANGES, epilogue I32 = RANGES["i32"] @@ -38,21 +39,31 @@ def validate_inputs(func: Func, inputs: Mapping[str, object]) -> dict[str, np.nd return out +def _check(r: np.ndarray, name: str, what: str) -> np.ndarray: + if r.size and (r.min() < I32[0] or r.max() > I32[1]): + raise RangeError(f"%{name} ({what}) does not fit in i32") + return r + + +def eval_op(kind, xs, attrs=None, data=None, name="?") -> np.ndarray: + """Evaluate one op on int64 arrays (used by `evaluate` and by constant folding).""" + if kind == "const": + return _check(np.array(data, dtype=np.int64), name, kind) + if kind == "matmul": + r = _check(xs[0] @ xs[1], name, "matmul") + for stage in epilogue(attrs): + r = _check(r + xs[2], name, "matmul bias add") if stage == "bias" else np.maximum(r, 0) + return r + if kind == "add": + return _check(xs[0] + xs[1], name, "add") + return np.maximum(xs[0], 0) # relu + + def evaluate(func: Func, inputs: Mapping[str, object]) -> list[np.ndarray]: - env = {k: v for k, v in validate_inputs(func, inputs).items()} + env = dict(validate_inputs(func, inputs)) for op in func.ops: - x = [env[v.name] for v in op.operands] - if op.kind == "const": - r = op.data - elif op.kind == "matmul": - r = x[0] @ x[1] - elif op.kind == "add": - r = x[0] + x[1] - else: # relu - r = np.maximum(x[0], 0) - if r.size and (r.min() < I32[0] or r.max() > I32[1]): - raise RangeError(f"%{op.result.name} ({op.kind}) does not fit in i32") - env[op.result.name] = r + xs = [env[v.name] for v in op.operands] + env[op.result.name] = eval_op(op.kind, xs, op.attrs, op.data, op.result.name) return [env[v.name] for v in func.outputs] diff --git a/tinyc/ir/types.py b/tinyc/ir/types.py index 3a2d738..8e2e13b 100644 --- a/tinyc/ir/types.py +++ b/tinyc/ir/types.py @@ -6,7 +6,7 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from tinyc.ir.errors import OpTypeError @@ -15,6 +15,8 @@ RANGES = {"i8": (-128, 127), "i32": (-(2**31), 2**31 - 1)} MAX_RANK = 2 OP_ARITY = {"const": 0, "relu": 1, "matmul": 2, "add": 2} +# A fused matmul applies these stages, in this order, to its result: `+ bias` then `relu`. +EPILOGUE_STAGES = ("bias", "relu") @dataclass(frozen=True) @@ -48,12 +50,28 @@ def type_problem(t: TensorType) -> str | None: return None -def infer_result_type(kind: str, operands: Sequence[TensorType]) -> TensorType: +def epilogue(attrs: Mapping[str, str] | None) -> tuple[str, ...]: + """The fused stages of a matmul ('bias', 'relu'), validated and in canonical order.""" + raw = (attrs or {}).get("epilogue", "") + parts = tuple(p for p in raw.split(",") if p) + ordered = tuple(s for s in EPILOGUE_STAGES if s in parts) + if parts != ordered: + raise OpTypeError(f"invalid epilogue '{raw}' (allowed: bias, relu, in that order)") + return parts + + +def expected_arity(kind: str, attrs: Mapping[str, str] | None = None) -> int: + return OP_ARITY[kind] + (kind == "matmul" and "bias" in epilogue(attrs)) + + +def infer_result_type( + kind: str, operands: Sequence[TensorType], attrs: Mapping[str, str] | None = None +) -> TensorType: """Result type of a non-const op, or OpTypeError with a user-facing message.""" if kind == "relu": (x,) = operands return x - a, b = operands + a, b = operands[0], operands[1] if kind == "matmul": if a.rank != 2 or b.rank != 2: raise OpTypeError(f"matmul needs rank-2 operands (got {a} @ {b})") @@ -61,6 +79,10 @@ def infer_result_type(kind: str, operands: Sequence[TensorType]) -> TensorType: raise OpTypeError( f"matmul inner dimensions differ: {a.shape[1]} vs {b.shape[0]} ({a} @ {b})" ) + if "bias" in epilogue(attrs): + bias = operands[2] + if bias.rank != 1 or bias.shape[0] != b.shape[1]: + raise OpTypeError(f"matmul bias must be [{b.shape[1]}], got {bias}") return TensorType("i32", (a.shape[0], b.shape[1])) if kind == "add": if a.shape == b.shape: diff --git a/tinyc/ir/verifier.py b/tinyc/ir/verifier.py index ed39de0..2d2bc59 100644 --- a/tinyc/ir/verifier.py +++ b/tinyc/ir/verifier.py @@ -6,7 +6,7 @@ from tinyc.ir.errors import OpTypeError, VerifyError from tinyc.ir.graph import Func -from tinyc.ir.types import OP_ARITY, RANGES, infer_result_type, type_problem +from tinyc.ir.types import OP_ARITY, RANGES, expected_arity, infer_result_type, type_problem def verify(func: Func) -> None: @@ -29,10 +29,16 @@ def check_type(what: str, t) -> None: if op.kind not in OP_ARITY: problems.append(f"{where}: unknown op") continue - if len(op.operands) != OP_ARITY[op.kind]: - problems.append( - f"{where}: expects {OP_ARITY[op.kind]} operands, got {len(op.operands)}" - ) + try: + arity = expected_arity(op.kind, op.attrs) + except OpTypeError as e: + problems.append(f"{where}: {e}") + continue + if op.attrs and (op.kind != "matmul" or set(op.attrs) != {"epilogue"}): + problems.append(f"{where}: unsupported attributes {sorted(op.attrs)}") + continue + if len(op.operands) != arity: + problems.append(f"{where}: expects {arity} operands, got {len(op.operands)}") continue if r.name in defined: problems.append(f"{where}: name already defined") @@ -57,7 +63,7 @@ def check_type(what: str, t) -> None: problems.append(f"{where}: only const may carry data") if operands_ok: try: - want = infer_result_type(op.kind, [v.type for v in op.operands]) + want = infer_result_type(op.kind, [v.type for v in op.operands], op.attrs) except OpTypeError as e: problems.append(f"{where}: {e}") else: diff --git a/tinyc/passes/__init__.py b/tinyc/passes/__init__.py new file mode 100644 index 0000000..99a4e76 --- /dev/null +++ b/tinyc/passes/__init__.py @@ -0,0 +1,21 @@ +"""Graph-level passes and the pass manager.""" + +from tinyc.passes.canonicalize import Canonicalize +from tinyc.passes.const_fold import ConstFold +from tinyc.passes.cse import CSE +from tinyc.passes.dce import DCE +from tinyc.passes.fuse import Fuse +from tinyc.passes.legalize_gpu import LegalizeGpu +from tinyc.passes.legalize_tpu import LegalizeTpu +from tinyc.passes.manager import PassContext, PassManager, PassReport, PassVerifyError +from tinyc.passes.range_check import RangeCheck +from tinyc.passes.shape_infer import ShapeInfer + +REGISTRY = { + p.name: p + for p in ( + ShapeInfer, Canonicalize, ConstFold, CSE, DCE, Fuse, RangeCheck, LegalizeTpu, LegalizeGpu + ) +} # fmt: skip + +__all__ = ["REGISTRY", "PassContext", "PassManager", "PassReport", "PassVerifyError"] diff --git a/tinyc/passes/canonicalize.py b/tinyc/passes/canonicalize.py new file mode 100644 index 0000000..4112320 --- /dev/null +++ b/tinyc/passes/canonicalize.py @@ -0,0 +1,28 @@ +"""Normal forms that make later passes simpler and let CSE see through operand order. + +* `add`: the rank-2 operand comes first in a broadcast add; same-shape adds order their operands + by definition order, so `a + b` and `b + a` become the same op. +* `relu(relu(x))` -> `relu(x)` (relu is idempotent). The dead outer op is left for DCE. +""" + + +class Canonicalize: + name = "canonicalize" + + def run(self, func, ctx): + position = {v.name: i for i, v in enumerate(func.args)} + for i, op in enumerate(func.ops, start=len(func.args)): + position[op.result.name] = i + defs = {op.result.name: op for op in func.ops} + + for op in func.ops: + if op.kind == "add": + a, b = op.operands + if a.type.rank < b.type.rank or ( + a.type.shape == b.type.shape and position[a.name] > position[b.name] + ): + op.operands = [b, a] + elif op.kind == "relu": + inner = defs.get(op.operands[0].name) + if inner is not None and inner.kind == "relu": + func.replace_uses(op.result, op.operands[0]) diff --git a/tinyc/passes/const_fold.py b/tinyc/passes/const_fold.py new file mode 100644 index 0000000..add8ca4 --- /dev/null +++ b/tinyc/passes/const_fold.py @@ -0,0 +1,21 @@ +"""Evaluate ops whose operands are all constants, using the oracle's own `eval_op`.""" + +from tinyc.ir import ref +from tinyc.ir.errors import PassError, RangeError + + +class ConstFold: + name = "const-fold" + + def run(self, func, ctx): + consts = {op.result.name: op.data for op in func.ops if op.kind == "const"} + for op in func.ops: + if op.kind == "const" or not all(v.name in consts for v in op.operands): + continue + xs = [consts[v.name] for v in op.operands] + try: + data = ref.eval_op(op.kind, xs, op.attrs, op.data, op.result.name) + except RangeError as e: + raise PassError(f"constant expression overflows i32 ({e})", op.loc) from None + op.kind, op.operands, op.attrs, op.data = "const", [], {}, data + consts[op.result.name] = data diff --git a/tinyc/passes/cse.py b/tinyc/passes/cse.py new file mode 100644 index 0000000..3658f10 --- /dev/null +++ b/tinyc/passes/cse.py @@ -0,0 +1,25 @@ +"""Common-subexpression elimination: identical ops (and identical constants) share one result.""" + + +def _key(op): + data = None if op.data is None else (op.data.shape, op.data.tobytes()) + return ( + op.kind, + tuple(id(v) for v in op.operands), + tuple(sorted(op.attrs.items())), + str(op.result.type), + data, + ) + + +class CSE: + name = "cse" + + def run(self, func, ctx): + seen, dead = {}, set() + for op in func.ops: + first = seen.setdefault(_key(op), op.result) + if first is not op.result: + func.replace_uses(op.result, first) # later ops now hash on the survivor + dead.add(id(op)) + func.ops = [op for op in func.ops if id(op) not in dead] diff --git a/tinyc/passes/dce.py b/tinyc/passes/dce.py new file mode 100644 index 0000000..3cee5b8 --- /dev/null +++ b/tinyc/passes/dce.py @@ -0,0 +1,14 @@ +"""Dead-code elimination: drop ops whose results never reach an output. Inputs are kept.""" + + +class DCE: + name = "dce" + + def run(self, func, ctx): + live = {v.name for v in func.outputs} + keep = [] + for op in reversed(func.ops): + if op.result.name in live: + keep.append(op) + live.update(v.name for v in op.operands) + func.ops = keep[::-1] diff --git a/tinyc/passes/fuse.py b/tinyc/passes/fuse.py new file mode 100644 index 0000000..2cdc8b6 --- /dev/null +++ b/tinyc/passes/fuse.py @@ -0,0 +1,48 @@ +"""Fuse `matmul` with the ops that follow it into one matmul with an epilogue. + + matmul -> add(bias[N]) -> relu becomes matmul{epilogue=bias,relu} + +A stage is absorbed only when the previous result has exactly one use and is not an output. +The fused op sits where the last absorbed op was (so a late-defined bias is still defined before +use), keeps that op's result value, and takes its source location (the outermost expression). +""" + +from tinyc.ir.graph import Op +from tinyc.ir.types import epilogue + + +class Fuse: + name = "fuse" + + def run(self, func, ctx): + users = func.users() + out_names = {v.name for v in func.outputs} + + def sole_user(v): + us = users.get(v.name, []) + return us[0] if len(us) == 1 and v.name not in out_names else None + + skip: set[int] = set() + fused_at: dict[int, Op] = {} + for op in func.ops: + if op.kind != "matmul" or epilogue(op.attrs) or id(op) in skip: + continue + cur, last, stages, operands = op.result, op, [], list(op.operands) + nxt = sole_user(cur) + if nxt is not None and nxt.kind == "add": + other = nxt.operands[1] if nxt.operands[0] is cur else nxt.operands[0] + if other.type.rank == 1 and other.type.shape[0] == cur.type.shape[1]: + stages.append("bias") + operands.append(other) + skip.add(id(last)) + cur, last = nxt.result, nxt + nxt = sole_user(cur) + if nxt is not None and nxt.kind == "relu": + stages.append("relu") + skip.add(id(last)) + cur, last = nxt.result, nxt + if stages: + fused_at[id(last)] = Op( + "matmul", operands, last.result, last.loc, None, {"epilogue": ",".join(stages)} + ) + func.ops = [fused_at.get(id(o), o) for o in func.ops if id(o) not in skip] diff --git a/tinyc/passes/legalize_gpu.py b/tinyc/passes/legalize_gpu.py new file mode 100644 index 0000000..0cb5457 --- /dev/null +++ b/tinyc/passes/legalize_gpu.py @@ -0,0 +1,26 @@ +"""tiny-gpu legality: launch size and memory budget (one thread per output element). + +The memory check is a conservative upper bound (every value gets its own buffer); the memory +planner in the GPU backend reuses buffers and will never need more. +""" + +from tinyc.ir.errors import PassError + + +class LegalizeGpu: + name = "legalize-gpu" + + def run(self, func, ctx): + for op in func.ops: + if op.kind != "const" and op.result.type.size > ctx.gpu_max_threads: + raise PassError( + f"{op.kind} result {op.result.type} needs {op.result.type.size} threads, " + f"more than tiny-gpu launches ({ctx.gpu_max_threads})", + op.loc, + ) + words = sum(v.type.size for v in func.args) + sum(o.result.type.size for o in func.ops) + if words > ctx.gpu_mem_words: + raise PassError( + f"program needs up to {words} words of tiny-gpu memory, budget is " + f"{ctx.gpu_mem_words}" + ) diff --git a/tinyc/passes/legalize_tpu.py b/tinyc/passes/legalize_tpu.py new file mode 100644 index 0000000..73d2c82 --- /dev/null +++ b/tinyc/passes/legalize_tpu.py @@ -0,0 +1,31 @@ +"""tiny-tpu legality: every matmul operand must be an i8 graph input or constant. + +The array multiplies int8 x int8 only, there is no requantization yet, and the hardware silently +truncates wider operands, so anything else is rejected here rather than miscomputed. +""" + +from tinyc.ir.errors import PassError + + +class LegalizeTpu: + name = "legalize-tpu" + + def run(self, func, ctx): + producer = {op.result.name: op for op in func.ops} + for op in func.ops: + if op.kind != "matmul": + continue + for v in op.operands[:2]: + if v.type.dtype != "i8": + raise PassError( + f"tiny-tpu needs i8 matmul operands, but %{v.name} is {v.type} " + "(requantization is not supported yet)", + op.loc, + ) + made_by = producer.get(v.name) + if made_by is not None and made_by.kind != "const": + raise PassError( + f"tiny-tpu matmul operands must be i8 inputs or constants, but " + f"%{v.name} is computed by '{made_by.kind}'", + op.loc, + ) diff --git a/tinyc/passes/manager.py b/tinyc/passes/manager.py new file mode 100644 index 0000000..2f63a72 --- /dev/null +++ b/tinyc/passes/manager.py @@ -0,0 +1,126 @@ +"""Pass manager: runs passes in order, verifies after each, and can translation-validate them. + +* After every pass the IR is structurally verified (`ir.verifier`), always. +* With `verify_each`, the IR is also *run* after every pass on seeded random inputs (the + independent scalar interpreter) and compared with the numpy oracle's answers for the program + as it was before any pass. A failure names the pass and shows the IR before and after. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import Protocol + +import numpy as np + +from tinyc.ir import interp, ref +from tinyc.ir.errors import RangeError, VerifyError +from tinyc.ir.graph import Func +from tinyc.ir.printer import print_func +from tinyc.ir.verifier import verify + + +@dataclass +class PassContext: + """Options and results shared by passes.""" + + # i32 *inputs* are assumed to lie in this range; range_check relies on it (i8 inputs are exact). + i32_input_range: tuple[int, int] = (-(2**20), 2**20) + gpu_max_threads: int = 1024 # tiny-gpu simulator speed cap + gpu_mem_words: int = 65536 # global memory budget for the memory planner + analysis: dict = field(default_factory=dict) # e.g. analysis["ranges"] from range_check + + +class Pass(Protocol): + name: str + + def run(self, func: Func, ctx: PassContext) -> None: ... + + +class PassVerifyError(Exception): + def __init__(self, pass_name: str, reason: str, before: str, after: str): + self.pass_name, self.reason, self.before, self.after = pass_name, reason, before, after + super().__init__( + f"pass '{pass_name}' broke the program: {reason}\n" + f"--- IR before ---\n{before}--- IR after ---\n{after}" + ) + + +@dataclass +class PassRecord: + name: str + seconds: float + changed: bool + ir: str # IR text after this pass + + +@dataclass +class PassReport: + initial_ir: str + records: list[PassRecord] = field(default_factory=list) + + def timing_table(self) -> str: + rows = [f"{r.name:<14} {r.seconds * 1000:8.2f} ms {'changed' if r.changed else '-'}" + for r in self.records] # fmt: skip + return "\n".join(rows) + + +class PassManager: + def __init__( + self, + passes: Sequence[Pass], + *, + verify_each: bool = False, + num_checks: int = 4, + seed: int = 42, + dump: Callable[[str, str], None] | None = None, + ): + self.passes, self.verify_each = list(passes), verify_each + self.num_checks, self.seed, self.dump = num_checks, seed, dump + + def _expectations(self, func: Func): + out = [] + for i in range(self.num_checks): + inputs = ref.random_inputs(func, np.random.default_rng(self.seed + i)) + try: + out.append((inputs, ref.evaluate(func, inputs))) + except RangeError: # this input overflows in the source program itself: not usable + continue + return out + + def _semantic_problem(self, func: Func, expectations) -> str | None: + for inputs, want in expectations: + try: + got = interp.run(func, inputs) + except RangeError as e: + return f"now fails on inputs that used to work: {e}" + for i, (g, w) in enumerate(zip(got, want)): + if g.shape != w.shape or not (g == w).all(): + return f"output {i} differs from the oracle on seeded random inputs" + return None + + def run(self, func: Func, ctx: PassContext | None = None) -> PassReport: + ctx = ctx or PassContext() + text = print_func(func) + report = PassReport(initial_ir=text) + expectations = self._expectations(func) if self.verify_each else [] + if self.dump: + self.dump("input", text) + for p in self.passes: + start = time.perf_counter() + p.run(func, ctx) + seconds = time.perf_counter() - start + after = print_func(func) + try: + verify(func) + except VerifyError as e: + raise PassVerifyError(p.name, "; ".join(e.problems), text, after) from None + if self.verify_each and (why := self._semantic_problem(func, expectations)): + raise PassVerifyError(p.name, why, text, after) + report.records.append(PassRecord(p.name, seconds, after != text, after)) + if self.dump: + self.dump(p.name, after) + text = after + return report diff --git a/tinyc/passes/range_check.py b/tinyc/passes/range_check.py new file mode 100644 index 0000000..3f39d96 --- /dev/null +++ b/tinyc/passes/range_check.py @@ -0,0 +1,73 @@ +"""Prove every value fits in i32 by interval analysis, or reject the program. + +tiny-gpu raises on i32 overflow and tiny-tpu accumulates in 32 bits, so the compiler must show +no intermediate can overflow for *any* legal input. Bounds: + * i8 inputs are exactly [-128, 127]; i32 inputs are assumed to lie in `ctx.i32_input_range`; + constants are exact; + * a matmul sum of K products, each within [pmin, pmax], lies in [K*pmin, K*pmax]; every + *partial* sum lies in [min(0, K*pmin), max(0, K*pmax)], and that envelope must fit too, + so the proof holds for any accumulation order (tiling on the TPU, the loop on the GPU). +Results go to `ctx.analysis["ranges"]` as {value name: (lo, hi)}. +""" + +from __future__ import annotations + +from tinyc.ir.errors import PassError +from tinyc.ir.types import RANGES, epilogue + +I32_LO, I32_HI = RANGES["i32"] +Interval = tuple[int, int] + + +def _mul(a: Interval, b: Interval) -> Interval: + corners = [x * y for x in a for y in b] + return min(corners), max(corners) + + +def _need_fit(op, lo: int, hi: int, what: str, ctx, has_i32_input: bool) -> None: + if I32_LO <= lo and hi <= I32_HI: + return + note = ( + f" (i32 inputs are assumed to lie in {list(ctx.i32_input_range)})" if has_i32_input else "" + ) + raise PassError( + f"%{op.result.name} ({what}) may reach {lo}..{hi}, which does not fit in i32{note}", + op.loc, + ) + + +class RangeCheck: + name = "range-check" + + def run(self, func, ctx): + ranges: dict[str, Interval] = {} + for a in func.args: + ranges[a.name] = RANGES["i8"] if a.type.dtype == "i8" else ctx.i32_input_range + has_i32 = any(a.type.dtype == "i32" for a in func.args) + + for op in func.ops: + xs = [ranges[v.name] for v in op.operands] + + def fit(r: Interval, what: str, op=op) -> None: + _need_fit(op, *r, what, ctx, has_i32) + + if op.kind == "const": + r = (int(op.data.min()), int(op.data.max())) + elif op.kind == "relu": + r = (max(xs[0][0], 0), max(xs[0][1], 0)) + elif op.kind == "add": + r = (xs[0][0] + xs[1][0], xs[0][1] + xs[1][1]) + fit(r, "add") + else: # matmul, possibly fused + k = op.operands[0].type.shape[1] + pmin, pmax = _mul(xs[0], xs[1]) + fit((min(0, k * pmin), max(0, k * pmax)), "matmul partial sums") + r = (k * pmin, k * pmax) + for stage in epilogue(op.attrs): + if stage == "bias": + r = (r[0] + xs[2][0], r[1] + xs[2][1]) + fit(r, "matmul bias add") + else: + r = (max(r[0], 0), max(r[1], 0)) + ranges[op.result.name] = r + ctx.analysis["ranges"] = ranges diff --git a/tinyc/passes/shape_infer.py b/tinyc/passes/shape_infer.py new file mode 100644 index 0000000..d7c10cd --- /dev/null +++ b/tinyc/passes/shape_infer.py @@ -0,0 +1,17 @@ +"""Recompute every result type from its operands (fixes IR that was parsed or hand-edited).""" + +from tinyc.ir.errors import OpTypeError, PassError +from tinyc.ir.types import infer_result_type + + +class ShapeInfer: + name = "shape-infer" + + def run(self, func, ctx): + for op in func.ops: + if op.kind == "const": + continue + try: + op.result.type = infer_result_type(op.kind, [v.type for v in op.operands], op.attrs) + except OpTypeError as e: + raise PassError(str(e), op.loc) from None