diff --git a/CLAUDE.md b/CLAUDE.md index 49a4ea8..b8db3f6 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -18,6 +18,7 @@ python3 -m venv .venv && .venv/bin/pip install -e ".[dev]" && .venv/bin/pip inst 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) 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 eaf8ed2..d80977a 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -13,9 +13,12 @@ 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. -- Phase 2 (pass manager + graph passes): DONE. 183 tests green. Passes: shape-infer, canonicalize, const-fold, cse, dce, +- Phase 2 (pass manager + graph passes): DONE. Passes: shape-infer, FIXPOINT{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). +- 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). ## Fixed decisions (do not re-litigate; change only via docs/ARCHITECTURE.md) | Decision | Choice | @@ -28,7 +31,9 @@ faked; host-side work is always labelled as host work. | GPU mapping | One thread per output element: block = row i, thread = column j | | 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) | +| 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 | +| 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 | | Verification | Oracle (numpy) + per-pass translation validation + real targets + Hypothesis differential + golden + perf pins | diff --git a/README.md b/README.md index d12e224..498fa3d 100644 --- a/README.md +++ b/README.md @@ -4,10 +4,11 @@ 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 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. +**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. 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). +[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). ```bash git clone --recurse-submodules && cd tiny-compiler diff --git a/docs/PASSES.md b/docs/PASSES.md index 3256e2c..57ebc02 100644 --- a/docs/PASSES.md +++ b/docs/PASSES.md @@ -3,10 +3,17 @@ `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-] +-O1: shape-infer, FIXPOINT{canonicalize, const-fold, cse, dce}, fuse, range-check, [legalize-] -O0: shape-infer, range-check, [legalize-] ``` +The four simplifications run as a **fixpoint group** (`passes.manager.Fixpoint`): they enable each +other, so the whole group repeats until a full round changes nothing. A single sweep is not enough: +CSE can merge an op into an earlier one and leave a commutative `add` with its operands in +non-canonical order (Hypothesis found this; `tests/passes/test_manager.py` keeps the shrunk +program). Every sub-pass is still verified, recorded and dumped on every round, so +`--print-after-all` shows each round. + | Pass | What it does | |---|---| | `shape-infer` | Recomputes every result type from its operands (repairs parsed or hand-edited IR). | @@ -48,3 +55,11 @@ 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). + +## Tile passes (tiny-tpu backend) +After lowering to Tile IR (`tinyc/tpu/lower.py`) the `-O1` pipeline runs `skip-zero-passes`: +a pass whose A or B window is a constant all-zero tile is removed together with its accumulation +(its product is zero). If that accumulation carried the output tile's epilogue, the epilogue moves +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. diff --git a/docs/PLAN.md b/docs/PLAN.md index 90cb6d8..a7ce42a 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -153,7 +153,7 @@ tpu.schedule @main tile=4x4 { schedule order matches tiny-tpu's `tilingSchedule()` (output tile row-major, K innermost), so the existing tiny-tpu L3 view can display it. - **Tile passes:** skip passes whose A or B tile is all zero (constants only), sink `relu`/bias - into the last `host.acc` of each output tile, and report utilization (useful MACs ÷ 16 per pass). + into the last `host.acc` of each output tile, and report utilization (useful MACs ÷ 64 per pass: a pass is 4×4×4 MACs). - **Emit:** `.ttpu.json` = a buffer table, int8 tile payloads (or slice refs), and a command list (`pass`, `acc`, `relu`, `add_bias`, `store`). The format is versioned and has a JSON Schema in `docs/TTPU_FORMAT.md`. diff --git a/docs/TTPU_FORMAT.md b/docs/TTPU_FORMAT.md new file mode 100644 index 0000000..e259898 --- /dev/null +++ b/docs/TTPU_FORMAT.md @@ -0,0 +1,58 @@ +# `.ttpu.json` (version 1) + +The program a tiny-tpu run executes. Machine-checked by `docs/ttpu.schema.json` +(`tests/tpu/test_ttpu_format.py` validates every emitted program against it). +Consumers: `runtime/ttpu-exec.mjs` (the executor; Node and browser) and `tinyc/tpu/interp.py` +(an independent reference interpreter written from this spec). + +```json +{ "format": "ttpu", "version": 1, "tile": 4, "name": "main", + "buffers": { "A": {"dtype":"i8","shape":[4,8],"kind":"input"}, + "1": {"dtype":"i32","shape":[4,4],"kind":"temp"} }, + "inputs": ["A", "B"], "outputs": ["1"], + "commands": [ ... ], "estimate": { ... } } +``` + +## Buffers +`dtype` is `i8` or `i32`; `shape` has rank 1 or 2 (row-major). `kind`: +`input` (supplied at run time), `const` (carries `data` as nested lists) or `temp` (zero-filled +at start). Buffer names are the Graph IR value names without `%`. + +## Commands, in execution order +Only `pass` runs on the chip. **Everything else is host work**, counted separately and never +attributed to the chip. + +A *window* `{"buf", "r", "c"}` is the 4x4 region of a rank-2 buffer whose top-left element is +`(r, c)`. Reading outside the buffer yields 0 (zero padding); writing outside it is dropped. + +| `op` | Fields | Meaning | +|---|---|---| +| `alloc` | `buf` | **host.** Zero the buffer. | +| `pass` | `id`, `a`, `b` (windows) | **chip.** `tmp[id] = a_window @ b_window` (int8 x int8 -> int32, 4x4). Each pass starts from psum = 0 and reloads its weights. | +| `acc` | `pass`, `dst` (window), optional `bias` (buffer name), optional `relu: true` | **host.** `dst += tmp[pass]` (in-bounds part), then the epilogue: `dst[r][c] += bias[c]`, then `max(dst, 0)`. | +| `epilogue` | `dst` (window), optional `bias`, `relu` | **host.** The epilogue alone (an output tile that no pass feeds). | +| `host_add` | `dst`, `a`, `b` | **host.** Whole-buffer add; same shape, or a rank-1 operand broadcast over the rows of a rank-2 one. | +| `host_relu` | `dst`, `src` | **host.** `dst = max(src, 0)`. | + +Every pass is consumed by exactly one `acc`, after it ran. Every value must fit in i32 or the +executor fails (`RangeError`); `range-check` proves this at compile time for any legal input. + +## Tiling order +For `C = A @ B`, `A: [M,K]`, `B: [K,N]`: output tiles row-major, K innermost, matching tiny-tpu's own +`tilingSchedule()`. `ceil(M/4) * ceil(N/4) * ceil(K/4)` passes before optimization. The last `acc` +of an output tile carries that tile's epilogue. + +## `estimate` +Static cost model; the executor must measure exactly this (`tests/tpu/test_tpu_rtl.py`). + +| Key | Meaning | +|---|---| +| `passes`, `passes_skipped` | chip passes left, and passes removed by `skip-zero-passes` | +| `cycles_per_pass` | 15 (14 from start to done, +1 for the last element to latch; `getCycleCount()`) | +| `rtl_cycles` | `passes * 15` | +| `host_ops` | counts per host command kind | +| `useful_macs` / `capacity_macs` / `utilization` | non-padding multiply-adds (`M*K*N` for the passes kept) over `64 * passes` (a pass is 4x4x4 MACs) | + +## Executor protocol (`runtime/ttpu-cli.mjs`) +Request `{ "program": , "inputs": { name: nested lists } }` -> +`{ "outputs": { name: {shape, data} }, "passes": [{id, cycles}], "totalCycles", "hostOps" }`. diff --git a/docs/ttpu.schema.json b/docs/ttpu.schema.json new file mode 100644 index 0000000..86d0cd9 --- /dev/null +++ b/docs/ttpu.schema.json @@ -0,0 +1,83 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/deaneeth/tiny-compiler/docs/ttpu.schema.json", + "title": "ttpu program (version 1)", + "type": "object", + "required": ["format", "version", "tile", "name", "buffers", "inputs", "outputs", "commands", "estimate"], + "additionalProperties": false, + "properties": { + "format": { "const": "ttpu" }, + "version": { "const": 1 }, + "tile": { "const": 4 }, + "name": { "type": "string" }, + "buffers": { + "type": "object", + "additionalProperties": { "$ref": "#/$defs/buffer" } + }, + "inputs": { "type": "array", "items": { "type": "string" } }, + "outputs": { "type": "array", "items": { "type": "string" }, "minItems": 1 }, + "commands": { "type": "array", "items": { "$ref": "#/$defs/command" } }, + "estimate": { + "type": "object", + "required": ["passes", "passes_skipped", "cycles_per_pass", "rtl_cycles", "host_ops", "useful_macs", "capacity_macs", "utilization"], + "additionalProperties": false, + "properties": { + "passes": { "type": "integer", "minimum": 0 }, + "passes_skipped": { "type": "integer", "minimum": 0 }, + "cycles_per_pass": { "type": "integer", "minimum": 1 }, + "rtl_cycles": { "type": "integer", "minimum": 0 }, + "host_ops": { + "type": "object", + "required": ["alloc", "acc", "epilogue", "host_add", "host_relu"], + "additionalProperties": { "type": "integer", "minimum": 0 } + }, + "useful_macs": { "type": "integer", "minimum": 0 }, + "capacity_macs": { "type": "integer", "minimum": 0 }, + "utilization": { "type": ["number", "null"], "minimum": 0, "maximum": 1 } + } + } + }, + "$defs": { + "buffer": { + "type": "object", + "required": ["dtype", "shape", "kind"], + "additionalProperties": false, + "properties": { + "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" } + } + }, + "ref": { + "type": "object", + "required": ["buf", "r", "c"], + "additionalProperties": false, + "properties": { + "buf": { "type": "string" }, + "r": { "type": "integer", "minimum": 0 }, + "c": { "type": "integer", "minimum": 0 } + } + }, + "command": { + "oneOf": [ + { "type": "object", "required": ["op", "buf"], "additionalProperties": false, + "properties": { "op": { "const": "alloc" }, "buf": { "type": "string" } } }, + { "type": "object", "required": ["op", "id", "a", "b"], "additionalProperties": false, + "properties": { "op": { "const": "pass" }, "id": { "type": "integer", "minimum": 0 }, + "a": { "$ref": "#/$defs/ref" }, "b": { "$ref": "#/$defs/ref" } } }, + { "type": "object", "required": ["op", "pass", "dst"], "additionalProperties": false, + "properties": { "op": { "const": "acc" }, "pass": { "type": "integer", "minimum": 0 }, + "dst": { "$ref": "#/$defs/ref" }, "bias": { "type": "string" }, "relu": { "const": true } } }, + { "type": "object", "required": ["op", "dst"], "additionalProperties": false, + "properties": { "op": { "const": "epilogue" }, "dst": { "$ref": "#/$defs/ref" }, + "bias": { "type": "string" }, "relu": { "const": true } } }, + { "type": "object", "required": ["op", "dst", "a", "b"], "additionalProperties": false, + "properties": { "op": { "const": "host_add" }, "dst": { "type": "string" }, + "a": { "type": "string" }, "b": { "type": "string" } } }, + { "type": "object", "required": ["op", "dst", "src"], "additionalProperties": false, + "properties": { "op": { "const": "host_relu" }, "dst": { "type": "string" }, "src": { "type": "string" } } } + ] + } + } +} diff --git a/pyproject.toml b/pyproject.toml index 97f40fc..05cae6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ requires-python = ">=3.12" dependencies = ["numpy>=2.0"] [project.optional-dependencies] -dev = ["pytest>=8", "hypothesis>=6.100", "black>=24.10", "ruff>=0.6"] +dev = ["pytest>=8", "hypothesis>=6.100", "jsonschema>=4", "black>=24.10", "ruff>=0.6"] [tool.setuptools.packages.find] include = ["tinyc*", "runtime*"] diff --git a/runtime/tpu_runner.py b/runtime/tpu_runner.py index f2af9d5..78a7103 100644 --- a/runtime/tpu_runner.py +++ b/runtime/tpu_runner.py @@ -10,7 +10,7 @@ import subprocess import tempfile from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path import numpy as np @@ -42,20 +42,49 @@ def run_tpu_passes(passes: Sequence[tuple[np.ndarray, np.ndarray]]) -> TpuResult for a, b in passes ] } + out = _call(payload) + return TpuResult( + passes=[ + PassResult(np.asarray(p["c"], dtype=np.int32).reshape(4, 4), p["cycles"]) + for p in out["passes"] + ], + total_cycles=out["totalCycles"], + ) + + +@dataclass +class ProgramResult: + outputs: dict[str, np.ndarray] # int64, original shapes + pass_cycles: list[int] # measured per pass (RTL clock cycles) + total_cycles: int + host_ops: dict[str, int] # counted by the executor, never attributed to the chip + passes: list[dict] = field(default_factory=list) + + +def _call(request: dict) -> dict: with tempfile.TemporaryDirectory() as tmp: - path = os.path.join(tmp, "passes.json") + path = os.path.join(tmp, "request.json") with open(path, "w") as f: - json.dump(payload, f) + json.dump(request, f) proc = subprocess.run( - ["node", str(CLI), path], capture_output=True, text=True, check=False, timeout=120 + ["node", str(CLI), path], capture_output=True, text=True, check=False, timeout=300 ) if proc.returncode != 0: raise RuntimeError(f"ttpu-cli failed ({proc.returncode}): {proc.stderr.strip()}") - out = json.loads(proc.stdout) - return TpuResult( - passes=[ - PassResult(np.asarray(p["c"], dtype=np.int32).reshape(4, 4), p["cycles"]) - for p in out["passes"] - ], + return json.loads(proc.stdout) + + +def run_ttpu(program: dict, inputs: dict) -> ProgramResult: + """Execute a `.ttpu.json` program on the real tiny-tpu RTL through the one JS executor.""" + payload = {k: np.asarray(v, dtype=np.int64).tolist() for k, v in inputs.items()} + out = _call({"program": program, "inputs": payload}) + return ProgramResult( + outputs={ + name: np.asarray(o["data"], dtype=np.int64).reshape(o["shape"]) + for name, o in out["outputs"].items() + }, + pass_cycles=[p["cycles"] for p in out["passes"]], total_cycles=out["totalCycles"], + host_ops=out["hostOps"], + passes=out["passes"], ) diff --git a/runtime/ttpu-cli.mjs b/runtime/ttpu-cli.mjs index 6ef9ed3..9ff2711 100644 --- a/runtime/ttpu-cli.mjs +++ b/runtime/ttpu-cli.mjs @@ -1,10 +1,14 @@ #!/usr/bin/env node -// ttpu-cli.mjs - run tiny-tpu passes from a JSON file and print JSON results. +// ttpu-cli.mjs - run tiny-tpu work from a JSON file and print JSON results. // -// node runtime/ttpu-cli.mjs +// node runtime/ttpu-cli.mjs // -// in : { "passes": [ { "a": [16 ints], "b": [16 ints] }, ... ] } -// out: { "passes": [ { "c": [16 ints], "cycles": n }, ... ], "totalCycles": n } +// request A (raw passes): { "passes": [ { "a": [16 ints], "b": [16 ints] }, ... ] } +// -> { "passes": [ { "c": [16 ints], "cycles": n }, ... ], "totalCycles": n } +// +// request B (a program): { "program": <.ttpu.json>, "inputs": { name: nested lists } } +// -> { "outputs": { name: {shape, data} }, "passes": [ {id, cycles} ], "totalCycles": n, +// "hostOps": { alloc, acc, epilogue, host_add, host_relu } } // // The tiny-tpu build is taken from $TINY_TPU_DIST (a dir holding tiny_tpu.mjs + tiny_tpu.wasm), // defaulting to the pinned submodule's web/public. @@ -19,17 +23,27 @@ const dist = process.env.TINY_TPU_DIST ?? join(here, "..", "targets", "tiny-tpu" const [file] = process.argv.slice(2); if (!file) { - console.error("usage: ttpu-cli.mjs "); + console.error("usage: ttpu-cli.mjs "); process.exit(2); } const { default: factory } = await import(pathToFileURL(resolve(dist, "tiny_tpu.mjs")).href); const tpu = await createTpu(factory); -const { passes } = JSON.parse(readFileSync(file, "utf8")); -const results = passes.map(({ a, b }) => { - const { c, cycles } = tpu.runPass(a, b); - return { c, cycles }; -}); -console.log( - JSON.stringify({ passes: results, totalCycles: results.reduce((n, r) => n + r.cycles, 0) }), -); +const request = JSON.parse(readFileSync(file, "utf8")); + +try { + if (request.program) { + console.log(JSON.stringify(tpu.executeProgram(request.program, request.inputs ?? {}))); + } else { + const results = request.passes.map(({ a, b }) => { + const { c, cycles } = tpu.runPass(a, b); + return { c, cycles }; + }); + console.log( + JSON.stringify({ passes: results, totalCycles: results.reduce((n, r) => n + r.cycles, 0) }), + ); + } +} catch (err) { + console.error(`${err.name}: ${err.message}`); + process.exit(1); +} diff --git a/runtime/ttpu-exec.mjs b/runtime/ttpu-exec.mjs index 102720e..d69040a 100644 --- a/runtime/ttpu-exec.mjs +++ b/runtime/ttpu-exec.mjs @@ -1,37 +1,174 @@ -// ttpu-exec.mjs - the ONE executor for tiny-tpu passes (Node in tests, browser in the UI). +// ttpu-exec.mjs - the ONE executor for tiny-tpu programs (Node in tests, browser in the UI). // -// Phase 0 scope: drive single 4x4 passes on the real tiny-tpu RTL (Verilator -> WASM). -// Phase 3 grows this into the full `.ttpu.json` command list (pass / acc / relu / ...). -// -// Inputs are flat row-major arrays of 16 numbers (int8 for a/b); results are int32[16]. +// A `.ttpu.json` program (docs/TTPU_FORMAT.md) is a list of commands. Only `pass` runs on the +// chip (the real tiny-tpu RTL, Verilator -> WASM); every other command is host work done here +// in plain integers and counted separately, so nothing host-side is passed off as chip work. export const TILE = 4; export const TILE_ELEMS = TILE * TILE; +const I32_MIN = -(2 ** 31); +const I32_MAX = 2 ** 31 - 1; +const DTYPE_RANGE = { i8: [-128, 127], i32: [I32_MIN, I32_MAX] }; + +const size = (shape) => shape.reduce((n, d) => n * d, 1); + +function fits(v, what) { + if (v < I32_MIN || v > I32_MAX) throw new RangeError(`${what} does not fit in i32`); + return v; +} + +/** Flatten nested lists into a row-major array, checking the shape and the dtype range. */ +function flatten(nested, shape, dtype, what) { + const out = []; + const walk = (x, depth) => { + if (depth === shape.length) { + if (!Number.isInteger(x)) throw new Error(`${what}: expected an integer`); + out.push(x); + return; + } + if (!Array.isArray(x) || x.length !== shape[depth]) { + throw new Error(`${what}: expected shape [${shape}]`); + } + for (const item of x) walk(item, depth + 1); + }; + walk(nested, 0); + const [lo, hi] = DTYPE_RANGE[dtype]; + if (out.some((v) => v < lo || v > hi)) throw new Error(`${what}: values outside ${dtype}`); + return out; +} + +/** Zero-padded 4x4 window of a rank-2 buffer at (r0, c0), as a flat row-major array of 16. */ +function getTile(buf, r0, c0) { + const [rows, cols] = buf.shape; + const tile = new Array(TILE_ELEMS).fill(0); + for (let i = 0; i < TILE; i++) { + for (let j = 0; j < TILE; j++) { + if (r0 + i < rows && c0 + j < cols) tile[i * TILE + j] = buf.data[(r0 + i) * cols + c0 + j]; + } + } + return tile; +} + +/** Apply bias then relu to the in-bounds part of dst's window at (r0, c0). */ +function epilogue(dst, r0, c0, cmd, buffers) { + const [rows, cols] = dst.shape; + const bias = cmd.bias === undefined ? null : buffers.get(cmd.bias).data; + for (let i = 0; i < TILE && r0 + i < rows; i++) { + for (let j = 0; j < TILE && c0 + j < cols; j++) { + const k = (r0 + i) * cols + c0 + j; + if (bias) dst.data[k] = fits(dst.data[k] + bias[c0 + j], "bias add"); + if (cmd.relu) dst.data[k] = Math.max(dst.data[k], 0); + } + } +} + /** * @param {() => Promise} factory the `createTinyTpu` default export of tiny_tpu.mjs */ export async function createTpu(factory) { const mod = await factory(); - return { - /** Run one 4x4 pass. Returns { c: number[16], cycles: number, states?: CycleState[] }. */ - runPass(a, b, { states = false } = {}) { - if (a.length !== TILE_ELEMS || b.length !== TILE_ELEMS) { - throw new Error(`a and b must each have ${TILE_ELEMS} elements`); + + /** Run one 4x4 pass on the RTL. Returns { c: number[16], cycles, states? }. */ + function runPass(a, b, { states = false } = {}) { + if (a.length !== TILE_ELEMS || b.length !== TILE_ELEMS) { + throw new Error(`a and b must each have ${TILE_ELEMS} elements`); + } + const sim = new mod.TinyTpuSim(); + try { + sim.reset(); + sim.loadA(a); + sim.loadB(b); + sim.start(); + const raw = sim.run(); + const out = { c: Array.from(sim.getResult(), Number), cycles: sim.getCycleCount() }; + if (states) out.states = Array.from(raw); + return out; + } finally { + sim.delete(); + } + } + + /** + * Execute a `.ttpu.json` program. + * @returns {{outputs: Object, passes: {id:number, cycles:number}[], totalCycles:number, + * hostOps: Object}} + */ + function executeProgram(program, inputs, { states = false } = {}) { + if (program.format !== "ttpu" || program.version !== 1) { + throw new Error(`unsupported program format ${program.format} v${program.version}`); + } + const buffers = new Map(); + for (const [name, b] of Object.entries(program.buffers)) { + const shape = b.shape; + let data; + if (b.kind === "const") data = flatten(b.data, shape, b.dtype, `const '${name}'`); + else if (b.kind === "input") { + if (!(name in inputs)) throw new Error(`missing input '${name}'`); + data = flatten(inputs[name], shape, b.dtype, `input '${name}'`); + } else data = new Array(size(shape)).fill(0); + buffers.set(name, { shape, dtype: b.dtype, data }); + } + + const tmp = new Map(); + const passes = []; + const hostOps = { alloc: 0, acc: 0, epilogue: 0, host_add: 0, host_relu: 0 }; + for (const cmd of program.commands) { + if (cmd.op === "pass") { + const { c, cycles, states: st } = runPass( + getTile(buffers.get(cmd.a.buf), cmd.a.r, cmd.a.c), + getTile(buffers.get(cmd.b.buf), cmd.b.r, cmd.b.c), + { states }, + ); + tmp.set(cmd.id, c); + passes.push(states ? { id: cmd.id, cycles, states: st } : { id: cmd.id, cycles }); + continue; } - const sim = new mod.TinyTpuSim(); - try { - sim.reset(); - sim.loadA(a); - sim.loadB(b); - sim.start(); - const raw = sim.run(); - const out = { c: Array.from(sim.getResult(), Number), cycles: sim.getCycleCount() }; - if (states) out.states = Array.from(raw); - return out; - } finally { - sim.delete(); + hostOps[cmd.op] += 1; + if (cmd.op === "alloc") { + buffers.get(cmd.buf).data.fill(0); + } else if (cmd.op === "acc") { + const dst = buffers.get(cmd.dst.buf); + const [rows, cols] = dst.shape; + const t = tmp.get(cmd.pass); + for (let i = 0; i < TILE && cmd.dst.r + i < rows; i++) { + for (let j = 0; j < TILE && cmd.dst.c + j < cols; j++) { + const k = (cmd.dst.r + i) * cols + cmd.dst.c + j; + dst.data[k] = fits(dst.data[k] + t[i * TILE + j], "accumulate"); + } + } + epilogue(dst, cmd.dst.r, cmd.dst.c, cmd, buffers); + } else if (cmd.op === "epilogue") { + epilogue(buffers.get(cmd.dst.buf), cmd.dst.r, cmd.dst.c, cmd, buffers); + } else if (cmd.op === "host_add") { + const dst = buffers.get(cmd.dst); + const a = buffers.get(cmd.a); + const b = buffers.get(cmd.b); + // same shape, or a rank-1 operand broadcast over the rows of a rank-2 one + const cols = dst.shape[dst.shape.length - 1]; + const at = (x, k) => (x.shape.length === dst.shape.length ? x.data[k] : x.data[k % cols]); + for (let k = 0; k < dst.data.length; k++) dst.data[k] = fits(at(a, k) + at(b, k), "add"); + } else if (cmd.op === "host_relu") { + const dst = buffers.get(cmd.dst); + const src = buffers.get(cmd.src); + for (let k = 0; k < dst.data.length; k++) dst.data[k] = Math.max(src.data[k], 0); + } else { + throw new Error(`unknown command '${cmd.op}'`); } - }, - }; + } + + const outputs = {}; + for (const name of program.outputs) { + const { shape, data } = buffers.get(name); + outputs[name] = { shape, data }; + } + return { + outputs, + passes, + totalCycles: passes.reduce((n, p) => n + p.cycles, 0), + hostOps, + }; + } + + return { runPass, executeProgram }; } diff --git a/tests/conftest.py b/tests/conftest.py index 6197886..1fe6b7c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,6 @@ +import json +from pathlib import Path + import numpy as np import pytest @@ -29,3 +32,27 @@ def check(path, actual: str) -> None: assert actual == path.read_text(), f"{path} differs; review, then `pytest --update-golden`" return check + + +BASELINE = Path(__file__).resolve().parent / "perf_baseline.json" + + +@pytest.fixture +def perf_pin(request): + """`perf_pin(target, data)`: pin cost numbers per target in tests/perf_baseline.json. + + An unexpected change fails; `--update-golden` rewrites only that target's entry. + """ + update = request.config.getoption("--update-golden") + + def check(target: str, data: dict) -> None: + baseline = json.loads(BASELINE.read_text()) if BASELINE.exists() else {} + if update: + baseline[target] = data + BASELINE.write_text(json.dumps(baseline, indent=2, sort_keys=True) + "\n") + return + assert ( + baseline.get(target) == data + ), f"{target} cost numbers changed; if intended run `pytest --update-golden`" + + return check diff --git a/tests/gen.py b/tests/gen.py index c1a42f3..049bc54 100644 --- a/tests/gen.py +++ b/tests/gen.py @@ -14,8 +14,8 @@ def _fmt(dtype, shape): def _nested(draw, shape, lo, hi): - if len(shape) == 1: - return [draw(st.integers(lo, hi)) for _ in range(shape[0])] + if len(shape) == 1: # zeros are common on purpose: they make whole constant tiles zero + return [draw(st.one_of(st.just(0), st.integers(lo, hi))) for _ in range(shape[0])] return [_nested(draw, shape[1:], lo, hi) for _ in range(shape[0])] @@ -86,3 +86,62 @@ def programs(draw, max_ops: int = 6) -> str: outs.append(pool[extra][0]) lines.append("output " + ", ".join(outs)) return "\n".join(lines) + "\n" + + +TPU_DIMS = st.sampled_from([1, 2, 3, 4, 5, 8, 9]) # 1..3 tiles per dimension, some partial +EPILOGUES = [ + "{m}", + "relu({m})", + "{m} + {w}", + "relu({m} + {w})", + "{w} + {m}", + "relu(relu({m}))", + "{m} + {m}", + "relu({m}) + relu({m})", + "relu({m} + {w}) + {m}", +] + + +def _const_with_zero_blocks(draw, rows, cols): + """An i8 matrix whose 4x4 blocks are zero with probability 0.4 (so whole tiles skip).""" + lo, hi = RANGES["i8"] + zero = { + (bi, bj): draw(st.integers(0, 4)) < 2 + for bi in range((rows + 3) // 4) + for bj in range((cols + 3) // 4) + } + return [ + [0 if zero[(i // 4, j // 4)] else draw(st.integers(lo, hi)) for j in range(cols)] + for i in range(rows) + ] + + +@st.composite +def tpu_programs(draw) -> str: + """Programs aimed at tiny-tpu: multi-tile matmuls, constants with zero tiles, every epilogue + shape, and a mix of fused and host-side add/relu. Always legal for the TPU.""" + lines: list[str] = [] + outs: list[str] = [] + for chain in range(draw(st.integers(1, 2))): + m, k, n = draw(TPU_DIMS), draw(TPU_DIMS), draw(TPU_DIMS) + left_const = draw(st.integers(0, 4)) == 0 + right_const = not left_const and draw(st.booleans()) + if left_const: + lines.append(f"const L{chain} : i8[{m}, {k}] = {_const_with_zero_blocks(draw, m, k)}") + lines.append(f"input R{chain} : i8[{k}, {n}]") + a, b = f"L{chain}", f"R{chain}" + elif right_const: + lines.append(f"input L{chain} : i8[{m}, {k}]") + lines.append(f"const R{chain} : i8[{k}, {n}] = {_const_with_zero_blocks(draw, k, n)}") + a, b = f"L{chain}", f"R{chain}" + else: + lines.append(f"input L{chain} : i8[{m}, {k}]") + lines.append(f"input R{chain} : i8[{k}, {n}]") + a, b = f"L{chain}", f"R{chain}" + form = draw(st.sampled_from(EPILOGUES)) + if "{w}" in form: + lines.append(f"input w{chain} : {draw(st.sampled_from(['i8', 'i32']))}[{n}]") + lines.append(f"Y{chain} = " + form.format(m=f"{a} @ {b}", w=f"w{chain}")) + outs.append(f"Y{chain}") + lines.append("output " + ", ".join(outs)) + return "\n".join(lines) + "\n" diff --git a/tests/golden/tile/big_mm.tile b/tests/golden/tile/big_mm.tile new file mode 100644 index 0000000..54729fb --- /dev/null +++ b/tests/golden/tile/big_mm.tile @@ -0,0 +1,57 @@ +tpu.schedule @main(%A: i8[12,12], %B: i8[12,12]) -> (%1) tile=4x4 { + host.alloc %1 : i32[12,12] = 0 + tpu.pass #0 a=%A[0:4,0:4] b=%B[0:4,0:4] // 15 cyc + host.acc %1[0:4,0:4] += #0 + tpu.pass #1 a=%A[0:4,4:8] b=%B[4:8,0:4] // 15 cyc + host.acc %1[0:4,0:4] += #1 + tpu.pass #2 a=%A[0:4,8:12] b=%B[8:12,0:4] // 15 cyc + host.acc %1[0:4,0:4] += #2 {relu} + tpu.pass #3 a=%A[0:4,0:4] b=%B[0:4,4:8] // 15 cyc + host.acc %1[0:4,4:8] += #3 + tpu.pass #4 a=%A[0:4,4:8] b=%B[4:8,4:8] // 15 cyc + host.acc %1[0:4,4:8] += #4 + tpu.pass #5 a=%A[0:4,8:12] b=%B[8:12,4:8] // 15 cyc + host.acc %1[0:4,4:8] += #5 {relu} + tpu.pass #6 a=%A[0:4,0:4] b=%B[0:4,8:12] // 15 cyc + host.acc %1[0:4,8:12] += #6 + tpu.pass #7 a=%A[0:4,4:8] b=%B[4:8,8:12] // 15 cyc + host.acc %1[0:4,8:12] += #7 + tpu.pass #8 a=%A[0:4,8:12] b=%B[8:12,8:12] // 15 cyc + host.acc %1[0:4,8:12] += #8 {relu} + tpu.pass #9 a=%A[4:8,0:4] b=%B[0:4,0:4] // 15 cyc + host.acc %1[4:8,0:4] += #9 + tpu.pass #10 a=%A[4:8,4:8] b=%B[4:8,0:4] // 15 cyc + host.acc %1[4:8,0:4] += #10 + tpu.pass #11 a=%A[4:8,8:12] b=%B[8:12,0:4] // 15 cyc + host.acc %1[4:8,0:4] += #11 {relu} + tpu.pass #12 a=%A[4:8,0:4] b=%B[0:4,4:8] // 15 cyc + host.acc %1[4:8,4:8] += #12 + tpu.pass #13 a=%A[4:8,4:8] b=%B[4:8,4:8] // 15 cyc + host.acc %1[4:8,4:8] += #13 + tpu.pass #14 a=%A[4:8,8:12] b=%B[8:12,4:8] // 15 cyc + host.acc %1[4:8,4:8] += #14 {relu} + tpu.pass #15 a=%A[4:8,0:4] b=%B[0:4,8:12] // 15 cyc + host.acc %1[4:8,8:12] += #15 + tpu.pass #16 a=%A[4:8,4:8] b=%B[4:8,8:12] // 15 cyc + host.acc %1[4:8,8:12] += #16 + tpu.pass #17 a=%A[4:8,8:12] b=%B[8:12,8:12] // 15 cyc + host.acc %1[4:8,8:12] += #17 {relu} + tpu.pass #18 a=%A[8:12,0:4] b=%B[0:4,0:4] // 15 cyc + host.acc %1[8:12,0:4] += #18 + tpu.pass #19 a=%A[8:12,4:8] b=%B[4:8,0:4] // 15 cyc + host.acc %1[8:12,0:4] += #19 + tpu.pass #20 a=%A[8:12,8:12] b=%B[8:12,0:4] // 15 cyc + host.acc %1[8:12,0:4] += #20 {relu} + tpu.pass #21 a=%A[8:12,0:4] b=%B[0:4,4:8] // 15 cyc + host.acc %1[8:12,4:8] += #21 + tpu.pass #22 a=%A[8:12,4:8] b=%B[4:8,4:8] // 15 cyc + host.acc %1[8:12,4:8] += #22 + tpu.pass #23 a=%A[8:12,8:12] b=%B[8:12,4:8] // 15 cyc + host.acc %1[8:12,4:8] += #23 {relu} + tpu.pass #24 a=%A[8:12,0:4] b=%B[0:4,8:12] // 15 cyc + host.acc %1[8:12,8:12] += #24 + tpu.pass #25 a=%A[8:12,4:8] b=%B[4:8,8:12] // 15 cyc + host.acc %1[8:12,8:12] += #25 + tpu.pass #26 a=%A[8:12,8:12] b=%B[8:12,8:12] // 15 cyc + host.acc %1[8:12,8:12] += #26 {relu} +} diff --git a/tests/golden/tile/const_mm.tile b/tests/golden/tile/const_mm.tile new file mode 100644 index 0000000..e20b978 --- /dev/null +++ b/tests/golden/tile/const_mm.tile @@ -0,0 +1,6 @@ +tpu.schedule @main(%X: i8[3,2]) -> (%2) tile=4x4 { + host.const %0 : i8[2,2] = [[1, -2], [3, 4]] + host.alloc %2 : i32[3,2] = 0 + tpu.pass #0 a=%X[0:4,0:4] b=%0[0:4,0:4] // 15 cyc + host.acc %2[0:4,0:4] += #0 {relu} +} diff --git a/tests/golden/tile/mm.tile b/tests/golden/tile/mm.tile new file mode 100644 index 0000000..4b0a185 --- /dev/null +++ b/tests/golden/tile/mm.tile @@ -0,0 +1,5 @@ +tpu.schedule @main(%A: i8[4,4], %B: i8[4,4]) -> (%0) tile=4x4 { + host.alloc %0 : i32[4,4] = 0 + tpu.pass #0 a=%A[0:4,0:4] b=%B[0:4,0:4] // 15 cyc + host.acc %0[0:4,0:4] += #0 +} diff --git a/tests/golden/tile/mm_bias_relu.tile b/tests/golden/tile/mm_bias_relu.tile new file mode 100644 index 0000000..131d52a --- /dev/null +++ b/tests/golden/tile/mm_bias_relu.tile @@ -0,0 +1,7 @@ +tpu.schedule @main(%X: i8[4,8], %W: i8[8,4], %b: i32[4]) -> (%2) tile=4x4 { + host.alloc %2 : i32[4,4] = 0 + tpu.pass #0 a=%X[0:4,0:4] b=%W[0:4,0:4] // 15 cyc + host.acc %2[0:4,0:4] += #0 + tpu.pass #1 a=%X[0:4,4:8] b=%W[4:8,0:4] // 15 cyc + host.acc %2[0:4,0:4] += #1 {bias=%b, relu} +} diff --git a/tests/golden/tile/mm_relu.tile b/tests/golden/tile/mm_relu.tile new file mode 100644 index 0000000..00719b6 --- /dev/null +++ b/tests/golden/tile/mm_relu.tile @@ -0,0 +1,7 @@ +tpu.schedule @main(%A: i8[4,8], %B: i8[8,4]) -> (%1) tile=4x4 { + host.alloc %1 : i32[4,4] = 0 + tpu.pass #0 a=%A[0:4,0:4] b=%B[0:4,0:4] // 15 cyc + host.acc %1[0:4,0:4] += #0 + tpu.pass #1 a=%A[0:4,4:8] b=%B[4:8,0:4] // 15 cyc + host.acc %1[0:4,0:4] += #1 {relu} +} diff --git a/tests/passes/test_manager.py b/tests/passes/test_manager.py index 695ba45..37bc5e7 100644 --- a/tests/passes/test_manager.py +++ b/tests/passes/test_manager.py @@ -4,7 +4,7 @@ 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 +from tinyc.passes import REGISTRY, Fixpoint, PassContext, PassManager, PassVerifyError SRC = "input A : i8[2, 3]\ninput B : i8[3, 2]\nC = relu(A @ B)\noutput C\n" @@ -83,7 +83,10 @@ def test_pipeline_names_per_target_and_level(): def test_disabling_passes_and_unknown_names(): - names = [p.name for p in build_pipeline("gpu", 1, disabled=["fuse", "cse"])] + def flat(pipeline): + return [q.name for p in pipeline for q in (p.passes if isinstance(p, Fixpoint) else [p])] + + names = flat(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"]) @@ -105,3 +108,54 @@ 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) + + +# --- fixpoint group --------------------------------------------------------------------------- + +# Found by Hypothesis (shrunk): CSE merges v3 into the earlier v0, which leaves `v2 + v0` with its +# operands in non-canonical order. One canonicalize-then-cse sweep stops there. +CSE_BREAKS_CANONICAL_ORDER = """\ +input x0 : i8[1] +input x1 : i8[2] +v0 = relu(x0) +v1 = relu(x1) +v2 = x0 + x0 +v3 = relu(x0) +v4 = relu(x0) +v5 = v2 + v3 +output v5 +""" +SIMPLIFY = ("canonicalize", "const-fold", "cse", "dce") + + +def source_func(): + from tinyc.frontend import compile_to_graph + + return compile_to_graph(CSE_BREAKS_CANONICAL_ORDER, "t.tiny") + + +def test_one_sweep_of_the_simplifications_is_not_a_fixpoint(): + f = source_func() + PassManager([REGISTRY[n]() for n in SIMPLIFY]).run(f) + assert PassManager([REGISTRY["canonicalize"]()]).run(f).records[0].changed + + +def test_the_fixpoint_group_reaches_a_stable_program(): + f = source_func() + group = Fixpoint([REGISTRY[n]() for n in SIMPLIFY]) + report = PassManager([group], verify_each=True).run(f) + again = PassManager([group], verify_each=True).run(f) + assert any(r.changed for r in report.records) + assert not any(r.changed for r in again.records) + assert len(report.records) > len(SIMPLIFY) # more than one round, each sub-pass recorded + + +def test_a_fixpoint_group_that_never_settles_is_an_error(): + class Flip: + name = "flip" + + def run(self, func, ctx): + func.name = "a" if func.name != "a" else "b" + + with pytest.raises(RuntimeError, match="did not converge"): + PassManager([Fixpoint([Flip()], max_iterations=3)]).run(func()) diff --git a/tests/perf_baseline.json b/tests/perf_baseline.json new file mode 100644 index 0000000..66600a5 --- /dev/null +++ b/tests/perf_baseline.json @@ -0,0 +1,138 @@ +{ + "tpu": { + "O0": { + "big_mm": { + "host_ops": { + "acc": 27, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 1 + }, + "passes": 27, + "rtl_cycles": 405, + "useful_macs": 1728, + "utilization": 1.0 + }, + "const_mm": { + "host_ops": { + "acc": 1, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 1 + }, + "passes": 1, + "rtl_cycles": 15, + "useful_macs": 12, + "utilization": 0.1875 + }, + "mm": { + "host_ops": { + "acc": 1, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 0 + }, + "passes": 1, + "rtl_cycles": 15, + "useful_macs": 64, + "utilization": 1.0 + }, + "mm_bias_relu": { + "host_ops": { + "acc": 2, + "alloc": 1, + "epilogue": 0, + "host_add": 1, + "host_relu": 1 + }, + "passes": 2, + "rtl_cycles": 30, + "useful_macs": 128, + "utilization": 1.0 + }, + "mm_relu": { + "host_ops": { + "acc": 2, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 1 + }, + "passes": 2, + "rtl_cycles": 30, + "useful_macs": 128, + "utilization": 1.0 + } + }, + "O1": { + "big_mm": { + "host_ops": { + "acc": 27, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 0 + }, + "passes": 27, + "rtl_cycles": 405, + "useful_macs": 1728, + "utilization": 1.0 + }, + "const_mm": { + "host_ops": { + "acc": 1, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 0 + }, + "passes": 1, + "rtl_cycles": 15, + "useful_macs": 12, + "utilization": 0.1875 + }, + "mm": { + "host_ops": { + "acc": 1, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 0 + }, + "passes": 1, + "rtl_cycles": 15, + "useful_macs": 64, + "utilization": 1.0 + }, + "mm_bias_relu": { + "host_ops": { + "acc": 2, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 0 + }, + "passes": 2, + "rtl_cycles": 30, + "useful_macs": 128, + "utilization": 1.0 + }, + "mm_relu": { + "host_ops": { + "acc": 2, + "alloc": 1, + "epilogue": 0, + "host_add": 0, + "host_relu": 0 + }, + "passes": 2, + "rtl_cycles": 30, + "useful_macs": 128, + "utilization": 1.0 + } + } + } +} diff --git a/tests/test_cli.py b/tests/test_cli.py index c030a4f..028ba55 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -70,3 +70,53 @@ def test_front_end_errors_still_work(tmp_path, capsys): src.write_text("C = nope\n") code, _, err = run(capsys, src) assert code == 1 and "undefined name 'nope'" in err + + +# --- tiny-tpu backend --------------------------------------------------------------------- + + +def test_emit_tile_prints_the_schedule(capsys): + code, out, _ = run(capsys, EX / "mm_relu.tiny", "--target", "tpu", "--emit", "tile") + assert code == 0 and out.startswith("tpu.schedule @main") and "host.acc" in out + assert out.count("tpu.pass") == 2 + + +def test_emit_ttpu_prints_valid_json(capsys): + import json + + code, out, _ = run(capsys, EX / "mm_bias_relu.tiny", "--target", "tpu", "--emit", "ttpu") + doc = json.loads(out) + assert code == 0 and doc["format"] == "ttpu" and doc["estimate"]["passes"] == 2 + + +def test_stats_prints_the_cost_estimate(capsys): + import json + + code, out, _ = run(capsys, EX / "big_mm.tiny", "--target", "tpu", "--stats") + est = json.loads(out) + assert code == 0 and est["passes"] == 27 and est["rtl_cycles"] == 405 + + +def test_run_executes_on_the_real_chip_and_reports_measured_vs_estimate(capsys): + code, out, _ = run(capsys, EX / "big_mm.tiny", "--target", "tpu", "--run") + assert code == 0 and "bit-exact vs numpy oracle: yes" in out + assert "measured RTL cycles: 405 (estimate 405)" in out and "not chip work" in out + + +def test_run_is_seeded(capsys): + a = run(capsys, EX / "mm.tiny", "--target", "tpu", "--run", "--seed", "1")[1] + b = run(capsys, EX / "mm.tiny", "--target", "tpu", "--run", "--seed", "1")[1] + c = run(capsys, EX / "mm.tiny", "--target", "tpu", "--run", "--seed", "2")[1] + assert a == b and a != c + + +@pytest.mark.parametrize("flags", [["--emit", "tile"], ["--emit", "ttpu"], ["--run"], ["--stats"]]) +def test_backend_flags_need_a_target(flags): + with pytest.raises(SystemExit) as exc: + main([str(EX / "mm.tiny"), *flags]) + assert exc.value.code == 2 + + +def test_tpu_backend_rejects_what_the_chip_cannot_run(capsys): + code, _, err = run(capsys, EX / "mlp_gpu.tiny", "--target", "tpu", "--run") + assert code == 1 and "tiny-tpu needs i8 matmul operands" in err diff --git a/tests/tpu/helpers.py b/tests/tpu/helpers.py new file mode 100644 index 0000000..37eeefa --- /dev/null +++ b/tests/tpu/helpers.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from tinyc.driver import compile_tpu +from tinyc.ir import ref +from tinyc.tpu import interp + +ROOT = Path(__file__).resolve().parent.parent.parent +TPU_EXAMPLES = sorted(p for p in (ROOT / "examples").glob("*.tiny") if p.stem != "mlp_gpu") + +# W's second K tile (rows 4..8) is all zero: that pass contributes nothing. +ZERO_K_TILE = """\ +input X : i8[4, 8] +const W : i8[8, 4] = [[1, 2, 3, 4], [5, 6, 7, 8], [1, 1, 1, 1], [2, 2, 2, 2], + [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]] +Y = relu(X @ W) +output Y +""" + + +def check_against_oracle(comp, seeds=(1, 2, 3)): + """Python reference interpreter on the emitted JSON vs the numpy oracle.""" + for seed in seeds: + inputs = ref.random_inputs(comp.func, np.random.default_rng(seed)) + got = interp.run_program(comp.program, inputs)["outputs"] + for name, want in zip(comp.schedule.outputs, ref.evaluate(comp.func, inputs)): + np.testing.assert_array_equal(got[name], want) + + +def build(source, opt=1, **kw): + return compile_tpu(source, "t.tiny", opt_level=opt, **kw) diff --git a/tests/tpu/test_tile_ir.py b/tests/tpu/test_tile_ir.py new file mode 100644 index 0000000..368aca6 --- /dev/null +++ b/tests/tpu/test_tile_ir.py @@ -0,0 +1,155 @@ +"""Lowering Graph IR -> Tile IR, the Tile IR verifier, and the cost model's arithmetic.""" + +from pathlib import Path + +import pytest + +from tinyc.driver import compile_graph +from tinyc.tpu import compile_tile_ir +from tinyc.tpu.cost import estimate +from tinyc.tpu.lower import lower +from tinyc.tpu.tile_ir import ( + CYCLES_PER_PASS, + Acc, + Buffer, + Epilogue, + HostAdd, + HostRelu, + Pass, + Schedule, + ScheduleError, + TileRef, + print_schedule, + verify_schedule, +) + +ROOT = Path(__file__).resolve().parent.parent.parent +TPU_EXAMPLES = sorted(p for p in (ROOT / "examples").glob("*.tiny") if p.stem != "mlp_gpu") + + +def tile(source, opt=1, **kw): + func, _ = compile_graph(source, "t.tiny", target="tpu", opt_level=opt) + return compile_tile_ir(func, opt, **kw) + + +def dense(m, k, n, body="C = A @ B", bias=False): + extra = f"input b : i32[{n}]\n" if bias else "" + return f"input A : i8[{m}, {k}]\ninput B : i8[{k}, {n}]\n{extra}{body}\noutput C\n" + + +@pytest.mark.parametrize( + "m,k,n,passes", [(4, 4, 4, 1), (4, 8, 4, 2), (5, 6, 7, 8), (12, 12, 12, 27), (1, 1, 1, 1)] +) +def test_pass_count_is_the_product_of_tile_counts(m, k, n, passes): + assert estimate(tile(dense(m, k, n), opt=0))["passes"] == passes + + +def test_tiling_order_matches_tiny_tpus_tiling_schedule(): + """Output tiles row-major, K innermost (web/src/lib/tiling.ts `tilingSchedule`).""" + s = tile(dense(5, 6, 7), opt=0) + got = [(p.a.r // 4, p.b.c // 4, p.a.c // 4) for p in s.commands if isinstance(p, Pass)] + want = [(oi, oj, k) for oi in range(2) for oj in range(2) for k in range(2)] + assert got == want + first = next(c for c in s.commands if isinstance(c, Pass)) + assert (first.a, first.b) == (TileRef("A", 0, 0), TileRef("B", 0, 0)) + assert first.a.c == first.b.r # A's K offset is B's row offset + + +def test_useful_macs_is_m_k_n_so_padding_is_the_only_waste(): + est = estimate(tile(dense(5, 6, 7), opt=0)) + assert est["useful_macs"] == 5 * 6 * 7 + assert est["capacity_macs"] == 8 * 64 + assert est["utilization"] == round(210 / 512, 4) + assert estimate(tile(dense(12, 12, 12), opt=0))["utilization"] == 1.0 + + +def test_epilogue_rides_on_the_last_accumulation_of_each_output_tile(): + s = tile(dense(8, 8, 4, "C = relu(A @ B + b)", bias=True), opt=1) # -O1 fuses the epilogue + accs = [c for c in s.commands if isinstance(c, Acc)] + assert len(accs) == 4 # 2 output tiles x 2 K tiles + assert [(a.bias, a.relu) for a in accs] == [(None, False), ("b", True)] * 2 + + +def test_unfused_ops_become_whole_buffer_host_ops(): + s = tile(dense(4, 4, 4, "H = A @ B\nC = relu(H + H)"), opt=0) + kinds = [type(c).__name__ for c in s.commands] + assert kinds.count("HostAdd") == 1 and kinds.count("HostRelu") == 1 + assert estimate(s)["host_ops"] == { + "alloc": 1, "acc": 1, "epilogue": 0, "host_add": 1, "host_relu": 1 + } # fmt: skip + + +def test_chip_cycles_are_pass_count_times_the_measured_pass_cost(): + est = estimate(tile(dense(12, 12, 12))) + assert est["cycles_per_pass"] == CYCLES_PER_PASS == 15 + assert est["rtl_cycles"] == 27 * 15 == 405 + + +@pytest.mark.parametrize("path", TPU_EXAMPLES, ids=lambda p: p.stem) +def test_tile_ir_golden(path, golden): + func, _ = compile_graph(path.read_text(), path.name, target="tpu") + golden( + ROOT / "tests" / "golden" / "tile" / f"{path.stem}.tile", + print_schedule(compile_tile_ir(func)), + ) + + +# --- the verifier ----------------------------------------------------------------------------- + + +def base(): + bufs = { + "A": Buffer("A", "i8", (4, 4), "input"), + "B": Buffer("B", "i8", (4, 4), "input"), + "C": Buffer("C", "i32", (4, 4), "temp"), + "v": Buffer("v", "i32", (4,), "input"), + } + cmds = [Pass(0, TileRef("A", 0, 0), TileRef("B", 0, 0)), Acc(0, TileRef("C", 0, 0))] + return Schedule("main", bufs, cmds, ["A", "B", "v"], ["C"]) + + +def test_a_lowered_schedule_verifies(): + verify_schedule(base()) + verify_schedule(lower(compile_graph(dense(5, 6, 7), "t.tiny", target="tpu")[0])) + + +@pytest.mark.parametrize( + "mutate, fragment", + [ + ( + lambda s: s.commands.__setitem__(0, Pass(0, TileRef("Z", 0, 0), TileRef("B", 0, 0))), + "unknown buffer", + ), + ( + lambda s: s.commands.append(Pass(0, TileRef("A", 0, 0), TileRef("B", 0, 0))), + "duplicate pass id", + ), + (lambda s: s.commands.reverse(), "not run yet"), + (lambda s: s.commands.pop(), "never accumulated"), + (lambda s: s.commands.append(Acc(0, TileRef("C", 0, 0))), "accumulated twice"), + (lambda s: s.buffers["A"].__setattr__("dtype", "i32"), "must be i8"), + ( + lambda s: s.commands.__setitem__(0, Pass(0, TileRef("A", 4, 0), TileRef("B", 0, 0))), + "outside", + ), + ( + lambda s: s.commands.__setitem__( + 0, Pass(0, TileRef("A", 0, 4 - 4), TileRef("B", 4, 0)) + ), + "outside", + ), + (lambda s: s.commands.append(Epilogue(TileRef("C", 0, 0), bias="A")), "rank 1"), + (lambda s: s.commands.append(HostAdd("C", "A", "nope")), "unknown buffer"), + (lambda s: s.commands.append(HostRelu("nope", "A")), "unknown buffer"), + (lambda s: s.buffers["C"].__setattr__("dtype", "i8"), "must be i32"), + ( + lambda s: s.commands.__setitem__(0, Pass(0, TileRef("A", 0, 0), TileRef("B", 1, 0))), + "A columns start", + ), + ], +) +def test_verifier_rejects_malformed_schedules(mutate, fragment): + s = base() + mutate(s) + with pytest.raises(ScheduleError, match=fragment): + verify_schedule(s) diff --git a/tests/tpu/test_tile_passes.py b/tests/tpu/test_tile_passes.py new file mode 100644 index 0000000..3e49282 --- /dev/null +++ b/tests/tpu/test_tile_passes.py @@ -0,0 +1,79 @@ +import pytest + +from tests.helpers import expected +from tests.tpu.helpers import ZERO_K_TILE, build, check_against_oracle +from tinyc.tpu.cost import estimate +from tinyc.tpu.tile_ir import Acc, Epilogue, Pass, print_schedule + + +def test_a_zero_constant_k_tile_loses_its_pass_and_the_epilogue_moves_earlier(): + comp = build(ZERO_K_TILE) + assert print_schedule(comp.schedule) == expected(""" + tpu.schedule @main(%X: i8[4,8]) -> (%2) tile=4x4 { + host.const %0 : i8[8,4] = [[1, 2, 3, 4], [5, 6, 7, 8], [1, 1, 1, 1], [2, 2, 2, 2], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]] + host.alloc %2 : i32[4,4] = 0 + tpu.pass #0 a=%X[0:4,0:4] b=%0[0:4,0:4] // 15 cyc + host.acc %2[0:4,0:4] += #0 {relu} + } + """) + assert comp.program["estimate"]["passes"] == 1 + assert comp.program["estimate"]["passes_skipped"] == 1 + check_against_oracle(comp) + + +def test_o0_keeps_every_pass(): + comp = build(ZERO_K_TILE, opt=0) + assert comp.program["estimate"]["passes"] == 2 + assert comp.program["estimate"]["passes_skipped"] == 0 + check_against_oracle(comp) + + +def test_an_all_zero_constant_with_a_bias_becomes_a_standalone_epilogue(): + comp = build( + "input X : i8[4, 4]\nconst Z : i8[4, 4] = [[0,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]]\n" + "input b : i32[4]\nY = relu(X @ Z + b)\noutput Y\n" + ) + cmds = comp.schedule.commands + assert not any(isinstance(c, (Pass, Acc)) for c in cmds) + (epi,) = [c for c in cmds if isinstance(c, Epilogue)] + assert (epi.bias, epi.relu) == ("b", True) + check_against_oracle(comp) + + +def test_an_all_zero_constant_with_only_relu_needs_no_work_at_all(): + comp = build( + "input X : i8[4, 4]\nconst Z : i8[4, 4] = [[0,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]]\n" + "Y = relu(X @ Z)\noutput Y\n" + ) + est = estimate(comp.schedule) + assert est["passes"] == 0 and est["host_ops"]["acc"] == 0 and est["host_ops"]["epilogue"] == 0 + check_against_oracle(comp) # the output is all zeros, straight from the zero-initialized buffer + + +def test_the_zero_tile_can_be_on_the_left_operand_too(): + comp = build( + "const Z : i8[4, 8] = [[1,2,3,4,0,0,0,0],[1,1,1,1,0,0,0,0],[2,2,2,2,0,0,0,0],[0,0,0,0,0,0,0,0]]\n" + "input X : i8[8, 4]\nY = Z @ X\noutput Y\n" + ) + assert comp.program["estimate"]["passes"] == 1 + check_against_oracle(comp) + + +def test_pass_ids_are_renumbered_contiguously_after_skipping(): + # Three K tiles, the middle one all zero. + rows = [[1, 1, 1, 1]] * 4 + [[0, 0, 0, 0]] * 4 + [[2, 2, 2, 2]] * 4 + comp = build(f"input X : i8[4, 12]\nconst W : i8[12, 4] = {rows}\nY = X @ W\noutput Y\n") + assert [c.id for c in comp.schedule.commands if isinstance(c, Pass)] == [0, 1] + assert [c.pass_id for c in comp.schedule.commands if isinstance(c, Acc)] == [0, 1] + check_against_oracle(comp) + + +def test_nothing_is_skipped_when_no_constant_tile_is_zero(): + comp = build("input X : i8[4, 4]\nconst W : i8[4, 4] = [[1,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]]\n" + "Y = X @ W\noutput Y\n") # fmt: skip + assert comp.program["estimate"]["passes_skipped"] == 0 # one non-zero entry keeps the tile + + +@pytest.mark.parametrize("opt", [0, 1]) +def test_verify_each_validates_the_tile_passes(opt): + build(ZERO_K_TILE, opt=opt, verify_each=True) diff --git a/tests/tpu/test_tpu_perf.py b/tests/tpu/test_tpu_perf.py new file mode 100644 index 0000000..a4fcfb5 --- /dev/null +++ b/tests/tpu/test_tpu_perf.py @@ -0,0 +1,16 @@ +"""Pinned cost numbers: an unexpected change in passes/cycles/host ops fails CI.""" + +from tests.tpu.helpers import TPU_EXAMPLES, build + + +def test_tpu_cost_numbers_are_pinned(perf_pin): + pinned = {} + for opt in (0, 1): + pinned[f"O{opt}"] = {} + for path in TPU_EXAMPLES: + est = build(path.read_text(), opt=opt).program["estimate"] + pinned[f"O{opt}"][path.stem] = { + k: est[k] + for k in ("passes", "rtl_cycles", "host_ops", "useful_macs", "utilization") + } + perf_pin("tpu", pinned) diff --git a/tests/tpu/test_tpu_properties.py b/tests/tpu/test_tpu_properties.py new file mode 100644 index 0000000..9a08fc5 --- /dev/null +++ b/tests/tpu/test_tpu_properties.py @@ -0,0 +1,47 @@ +"""Random programs through the TPU backend: reference interpreter, and the real RTL.""" + +import numpy as np +from hypothesis import HealthCheck, event, given, settings +from hypothesis import strategies as st + +from runtime.tpu_runner import run_ttpu +from tests.gen import programs, tpu_programs +from tests.tpu.helpers import build +from tinyc.ir import ref +from tinyc.ir.errors import PassError +from tinyc.tpu import interp + + +def _compile(source, opt): + try: + return build(source, opt=opt, verify_each=True) + except PassError: # not tiny-tpu legal (i32 matmul operand) or overflow not provable + event("rejected by legalization or range-check") + return None + + +@settings(max_examples=150, deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given(st.one_of(tpu_programs(), programs()), st.sampled_from([0, 1]), st.integers(0, 2**32 - 1)) +def test_reference_interpreter_matches_oracle_on_random_programs(source, opt, seed): + comp = _compile(source, opt) + if comp is None: + return + inputs = ref.random_inputs(comp.func, np.random.default_rng(seed)) + got = interp.run_program(comp.program, inputs) + assert got["passes"] == comp.program["estimate"]["passes"] + for name, want in zip(comp.schedule.outputs, ref.evaluate(comp.func, inputs)): + np.testing.assert_array_equal(got["outputs"][name], want) + + +@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given(tpu_programs(), st.sampled_from([0, 1]), st.integers(0, 2**32 - 1)) +def test_real_rtl_matches_oracle_and_the_cost_model_on_random_programs(source, opt, seed): + comp = _compile(source, opt) + if comp is None: + return + est = comp.program["estimate"] + inputs = ref.random_inputs(comp.func, np.random.default_rng(seed)) + got = run_ttpu(comp.program, inputs) + for name, want in zip(comp.schedule.outputs, ref.evaluate(comp.func, inputs)): + np.testing.assert_array_equal(got.outputs[name], want) + assert got.total_cycles == est["rtl_cycles"] and got.host_ops == est["host_ops"] diff --git a/tests/tpu/test_tpu_rtl.py b/tests/tpu/test_tpu_rtl.py new file mode 100644 index 0000000..3237362 --- /dev/null +++ b/tests/tpu/test_tpu_rtl.py @@ -0,0 +1,75 @@ +"""The real chip: every program runs on the Verilated tiny-tpu RTL through the JS executor.""" + +import numpy as np +import pytest + +from runtime.tpu_runner import run_tpu_passes, run_ttpu +from tests.tpu.helpers import TPU_EXAMPLES, ZERO_K_TILE, build +from tinyc.ir import ref +from tinyc.tpu.tile_ir import CYCLES_PER_PASS + + +def test_the_cost_models_pass_cost_is_what_the_rtl_measures(): + eye = np.eye(4, dtype=np.int64) + assert run_tpu_passes([(eye, eye)]).passes[0].cycles == CYCLES_PER_PASS + + +@pytest.mark.parametrize("path", TPU_EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +def test_examples_are_bit_exact_and_the_cost_model_equals_the_measurement(path, opt): + comp = build(path.read_text(), opt=opt) + est = comp.program["estimate"] + for seed in (1, 2): + inputs = ref.random_inputs(comp.func, np.random.default_rng(seed)) + got = run_ttpu(comp.program, inputs) + for name, want in zip(comp.schedule.outputs, ref.evaluate(comp.func, inputs)): + np.testing.assert_array_equal(got.outputs[name], want) + assert len(got.pass_cycles) == est["passes"] + assert set(got.pass_cycles) <= {CYCLES_PER_PASS} + assert got.total_cycles == est["rtl_cycles"] # estimate == measurement, exactly + assert got.host_ops == est["host_ops"] + + +def test_big_matmul_is_27_passes_on_the_chip(): + comp = build((TPU_EXAMPLES[0].parent / "big_mm.tiny").read_text()) + inputs = ref.random_inputs(comp.func, np.random.default_rng(5)) + got = run_ttpu(comp.program, inputs) + assert len(got.pass_cycles) == 27 and got.total_cycles == 405 + assert got.host_ops["acc"] == 27 # 27 accumulations, all done by the host + + +def test_skipping_a_zero_tile_saves_exactly_one_passes_worth_of_cycles(): + inputs = {"X": np.arange(-16, 16).reshape(4, 8)} + slow = run_ttpu(build(ZERO_K_TILE, opt=0).program, inputs) + fast = run_ttpu(build(ZERO_K_TILE, opt=1).program, inputs) + np.testing.assert_array_equal(fast.outputs["2"], slow.outputs["2"]) + assert slow.total_cycles - fast.total_cycles == CYCLES_PER_PASS + assert fast.total_cycles == 15 + + +def test_partial_tiles_are_zero_padded_correctly(): + comp = build("input A : i8[5, 6]\ninput B : i8[6, 7]\nC = relu(A @ B)\noutput C\n") + rng = np.random.default_rng(11) + inputs = {"A": rng.integers(-128, 128, (5, 6)), "B": rng.integers(-128, 128, (6, 7))} + got = run_ttpu(comp.program, inputs) + np.testing.assert_array_equal(got.outputs["1"], np.maximum(inputs["A"] @ inputs["B"], 0)) + assert len(got.pass_cycles) == 8 + + +def test_extreme_values_through_a_full_program(): + comp = build("input A : i8[4, 8]\ninput B : i8[8, 4]\nC = A @ B\noutput C\n") + inputs = {"A": np.full((4, 8), -128), "B": np.full((8, 4), -128)} + got = run_ttpu(comp.program, inputs) + assert (got.outputs["0"] == 8 * 128 * 128).all() + + +def test_the_executor_rejects_bad_inputs(): + comp = build("input A : i8[4, 4]\ninput B : i8[4, 4]\nC = A @ B\noutput C\n") + ok = {"A": np.zeros((4, 4), dtype=int), "B": np.zeros((4, 4), dtype=int)} + run_ttpu(comp.program, ok) + with pytest.raises(RuntimeError, match="missing input 'B'"): + run_ttpu(comp.program, {"A": ok["A"]}) + with pytest.raises(RuntimeError, match="values outside i8"): + run_ttpu(comp.program, {**ok, "A": np.full((4, 4), 200)}) + with pytest.raises(RuntimeError, match="expected shape"): + run_ttpu(comp.program, {**ok, "B": np.zeros((4, 5), dtype=int)}) diff --git a/tests/tpu/test_ttpu_format.py b/tests/tpu/test_ttpu_format.py new file mode 100644 index 0000000..e0d5326 --- /dev/null +++ b/tests/tpu/test_ttpu_format.py @@ -0,0 +1,75 @@ +import copy +import json + +import jsonschema +import numpy as np +import pytest + +from tests.tpu.helpers import ROOT, TPU_EXAMPLES, ZERO_K_TILE, build, check_against_oracle +from tinyc.ir import ref +from tinyc.tpu import interp +from tinyc.tpu.cost import estimate +from tinyc.tpu.emit import dumps + +SCHEMA = json.loads((ROOT / "docs" / "ttpu.schema.json").read_text()) + + +def validate(doc): + jsonschema.Draft202012Validator(SCHEMA).validate(doc) + + +def test_the_schema_itself_is_valid(): + jsonschema.Draft202012Validator.check_schema(SCHEMA) + + +@pytest.mark.parametrize("path", TPU_EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +def test_every_emitted_program_matches_the_schema(path, opt): + comp = build(path.read_text(), opt=opt) + validate(comp.program) + assert comp.program["estimate"] == estimate(comp.schedule) + + +def test_emitted_json_is_deterministic_and_round_trips(): + a = build(ZERO_K_TILE).program + assert dumps(a) == dumps(build(ZERO_K_TILE).program) + assert json.loads(dumps(a)) == a + + +@pytest.mark.parametrize("path", TPU_EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +def test_reference_interpreter_agrees_with_the_oracle(path, opt): + check_against_oracle(build(path.read_text(), opt=opt)) + + +def test_interpreter_counts_what_the_estimate_predicts(): + for path in TPU_EXAMPLES: + comp = build(path.read_text()) + inputs = ref.random_inputs(comp.func, np.random.default_rng(3)) + run = interp.run_program(comp.program, inputs) + assert run["passes"] == comp.program["estimate"]["passes"] + assert run["host_ops"] == comp.program["estimate"]["host_ops"] + + +@pytest.mark.parametrize( + "mutate", + [ + lambda d: d.__setitem__("version", 2), + lambda d: d.__setitem__("tile", 8), + lambda d: d.__setitem__("extra", 1), + lambda d: d["commands"].append({"op": "frobnicate"}), + lambda d: d["commands"][1].pop("id"), + lambda d: d["commands"][1]["a"].__setitem__("r", -1), + lambda d: d["commands"][1].__setitem__("surprise", True), + lambda d: d["buffers"]["A"].__setitem__("dtype", "f32"), + lambda d: d["buffers"]["A"].__setitem__("shape", [1, 2, 3]), + lambda d: d.__setitem__("outputs", []), + lambda d: d["estimate"].pop("rtl_cycles"), + ], +) +def test_the_schema_rejects_malformed_programs(mutate): + doc = copy.deepcopy(build((ROOT / "examples" / "mm_relu.tiny").read_text()).program) + validate(doc) + mutate(doc) + with pytest.raises(jsonschema.ValidationError): + validate(doc) diff --git a/tinyc/cli.py b/tinyc/cli.py index 2066e4d..6c2ca56 100644 --- a/tinyc/cli.py +++ b/tinyc/cli.py @@ -1,25 +1,32 @@ -"""`tinyc FILE.tiny`: compile to Graph IR (optionally through the pass pipeline) and print it.""" +"""`tinyc FILE.tiny`: compile through the pass pipeline; print IR, a tile schedule or a program.""" from __future__ import annotations import argparse +import json import sys from pathlib import Path -from tinyc.driver import build_pipeline +import numpy as np + +from tinyc.driver import build_pipeline, compile_tpu from tinyc.frontend import CompileError, SourceFile, lower, parse from tinyc.frontend.ast import dump from tinyc.frontend.diagnostics import Diagnostic, Span +from tinyc.ir import ref 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 +from tinyc.tpu.emit import dumps +from tinyc.tpu.tile_ir import print_schedule def main(argv: list[str] | None = None) -> int: ap = argparse.ArgumentParser(prog="tinyc", description=__doc__) ap.add_argument("file", type=Path) - ap.add_argument("--emit", choices=["ast", "ir"], default="ir") + ap.add_argument("--emit", choices=["ast", "ir", "tile", "ttpu"], default="ir", + help="tile and ttpu need --target tpu") # fmt: skip 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( @@ -32,8 +39,16 @@ def main(argv: list[str] | None = None) -> int: 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") + ap.add_argument("--run", action="store_true", + help="run on the real chip with seeded random inputs (--target tpu)") # fmt: skip + ap.add_argument("--seed", type=int, default=42, help="seed for --run inputs") + ap.add_argument("--stats", action="store_true", help="print the static cost estimate") args = ap.parse_args(argv) + if (args.emit in ("tile", "ttpu") or args.run or args.stats) and args.target != "tpu": + ap.error( + "--emit tile/ttpu, --run and --stats need --target tpu (the gpu backend comes later)" + ) source = SourceFile(args.file.name, args.file.read_text()) def dump_ir(name: str, text: str) -> None: @@ -53,6 +68,13 @@ def dump_ir(name: str, text: str) -> None: dump=dump_ir, ) report = manager.run(func) + if args.target == "tpu" and (args.emit in ("tile", "ttpu") or args.run or args.stats): + # the backend re-runs the graph pipeline itself, so the dumps above are not repeated + comp = compile_tpu( + source.text, source.name, opt_level=args.opt, disabled=args.disable, + verify_each=args.verify_each, dump=dump_ir if args.print_after_all else None, + ) # fmt: skip + return _tpu_output(args, comp, report) except CompileError as err: print(err.render(), file=sys.stderr) return 1 @@ -74,5 +96,32 @@ def dump_ir(name: str, text: str) -> None: return 0 +def _tpu_output(args, comp, report) -> int: + est = comp.program["estimate"] + if args.time_passes: + print(report.timing_table(), file=sys.stderr) + if args.run: + from runtime.tpu_runner import run_ttpu + + inputs = ref.random_inputs(comp.func, np.random.default_rng(args.seed)) + got = run_ttpu(comp.program, inputs) + want = ref.evaluate(comp.func, inputs) + exact = all((got.outputs[name] == w).all() for name, w in zip(comp.schedule.outputs, want)) + for name, w in zip(comp.schedule.outputs, want): + print(f"%{name} =\n{got.outputs[name]}") + print(f"bit-exact vs numpy oracle: {'yes' if exact else 'NO'}") + print(f"chip passes: {len(got.pass_cycles)} measured RTL cycles: {got.total_cycles}" + f" (estimate {est['rtl_cycles']})") # fmt: skip + print(f"host ops (not chip work): {got.host_ops}") + return 0 if exact else 1 + if args.stats: + print(json.dumps(est, indent=2)) + elif args.emit == "tile": + print(print_schedule(comp.schedule), end="") + else: + print(dumps(comp.program), end="") + return 0 + + if __name__ == "__main__": raise SystemExit(main()) diff --git a/tinyc/driver.py b/tinyc/driver.py index 4ab9429..473ed33 100644 --- a/tinyc/driver.py +++ b/tinyc/driver.py @@ -3,19 +3,27 @@ from __future__ import annotations from collections.abc import Callable, Iterable +from dataclasses import dataclass from tinyc.frontend import compile_to_graph from tinyc.ir.graph import Func -from tinyc.passes import REGISTRY, PassContext, PassManager, PassReport +from tinyc.passes import REGISTRY, Fixpoint, PassContext, PassManager, PassReport +from tinyc.tpu import compile_tile_ir +from tinyc.tpu.emit import to_json +from tinyc.tpu.tile_ir import Schedule -OPTIMIZATIONS = ("canonicalize", "const-fold", "cse", "dce", "fuse") +# These enable each other, so they run as one fixpoint group (see passes.manager.Fixpoint). +SIMPLIFY = ("canonicalize", "const-fold", "cse", "dce") def pipeline_names(target: str | None = None, opt_level: int = 1) -> list[str]: - """-O0 only type-checks and proves ranges; -O1 adds the optimizations.""" + """-O0 only type-checks and proves ranges; -O1 adds the simplifications and fusion. + + The simplify group (`SIMPLIFY`) runs to a fixpoint; names are listed once, in group order. + """ names = ["shape-infer"] if opt_level >= 1: - names += ["canonicalize", "const-fold", "canonicalize", "cse", "dce", "fuse"] + names += [*SIMPLIFY, "fuse"] names.append("range-check") if target in ("tpu", "gpu"): names.append(f"legalize-{target}") @@ -29,7 +37,21 @@ def build_pipeline( 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] + out: list = [] + group: list = [] + for n in pipeline_names(target, opt_level): + if n in off: + continue + if n in SIMPLIFY: + group.append(REGISTRY[n]()) + continue + if group: + out.append(Fixpoint(group)) + group = [] + out.append(REGISTRY[n]()) + if group: # (never reached: fuse/range-check always follow) kept for safety + out.append(Fixpoint(group)) + return out def compile_graph( @@ -49,3 +71,30 @@ def compile_graph( build_pipeline(target, opt_level, disabled), verify_each=verify_each, dump=dump ) return func, manager.run(func, ctx) + + +@dataclass +class TpuCompilation: + func: Func # the optimized, tpu-legalized Graph IR + schedule: Schedule # Tile IR + program: dict # the `.ttpu.json` document + report: PassReport # graph pass records + + +def compile_tpu( + text: str, + filename: str = "", + *, + opt_level: int = 1, + disabled: Iterable[str] = (), + verify_each: bool = False, + ctx: PassContext | None = None, + dump: Callable[[str, str], None] | None = None, +) -> TpuCompilation: + """Source text -> `.ttpu.json` for tiny-tpu (graph passes, then lowering and tile passes).""" + func, report = compile_graph( + text, filename, target="tpu", opt_level=opt_level, disabled=disabled, + verify_each=verify_each, ctx=ctx, dump=dump, + ) # fmt: skip + schedule = compile_tile_ir(func, opt_level, verify_each=verify_each, dump=dump) + return TpuCompilation(func, schedule, to_json(schedule), report) diff --git a/tinyc/passes/__init__.py b/tinyc/passes/__init__.py index 99a4e76..258c3d6 100644 --- a/tinyc/passes/__init__.py +++ b/tinyc/passes/__init__.py @@ -7,7 +7,7 @@ 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.manager import Fixpoint, PassContext, PassManager, PassReport, PassVerifyError from tinyc.passes.range_check import RangeCheck from tinyc.passes.shape_infer import ShapeInfer @@ -18,4 +18,4 @@ ) } # fmt: skip -__all__ = ["REGISTRY", "PassContext", "PassManager", "PassReport", "PassVerifyError"] +__all__ = ["REGISTRY", "Fixpoint", "PassContext", "PassManager", "PassReport", "PassVerifyError"] diff --git a/tinyc/passes/manager.py b/tinyc/passes/manager.py index 2f63a72..de27534 100644 --- a/tinyc/passes/manager.py +++ b/tinyc/passes/manager.py @@ -39,6 +39,20 @@ class Pass(Protocol): def run(self, func: Func, ctx: PassContext) -> None: ... +class Fixpoint: + """Run a group of passes repeatedly until a whole round changes nothing. + + Simplification passes enable each other (CSE can make operands non-canonical, folding can + expose dead code, ...), so running each once leaves work behind and the pipeline would not be + idempotent. Every sub-pass is still verified, recorded and dumped on every round. + """ + + name = "fixpoint" + + def __init__(self, passes: Sequence[Pass], max_iterations: int = 8): + self.passes, self.max_iterations = list(passes), max_iterations + + 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 @@ -108,7 +122,10 @@ def run(self, func: Func, ctx: PassContext | None = None) -> PassReport: expectations = self._expectations(func) if self.verify_each else [] if self.dump: self.dump("input", text) - for p in self.passes: + + def step(p: Pass) -> bool: + """Run one pass with verification, recording and dumping. True if the IR changed.""" + nonlocal text start = time.perf_counter() p.run(func, ctx) seconds = time.perf_counter() - start @@ -119,8 +136,23 @@ def run(self, func: Func, ctx: PassContext | None = None) -> PassReport: 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)) + changed = after != text + report.records.append(PassRecord(p.name, seconds, changed, after)) if self.dump: self.dump(p.name, after) text = after + return changed + + for item in self.passes: + if isinstance(item, Fixpoint): + for _ in range(item.max_iterations): + changed = [step(p) for p in item.passes] # every sub-pass runs each round + if not any(changed): + break + else: + raise RuntimeError( + f"fixpoint group did not converge in {item.max_iterations} rounds" + ) + else: + step(item) return report diff --git a/tinyc/tpu/__init__.py b/tinyc/tpu/__init__.py new file mode 100644 index 0000000..82ac427 --- /dev/null +++ b/tinyc/tpu/__init__.py @@ -0,0 +1,72 @@ +"""tiny-tpu backend: Graph IR -> Tile IR -> `.ttpu.json`.""" + +from __future__ import annotations + +from collections.abc import Callable + +import numpy as np + +from tinyc.ir import ref +from tinyc.ir.errors import RangeError +from tinyc.ir.graph import Func +from tinyc.passes.manager import PassVerifyError +from tinyc.tpu import interp +from tinyc.tpu.emit import to_json +from tinyc.tpu.lower import lower +from tinyc.tpu.tile_ir import Schedule, ScheduleError, print_schedule, verify_schedule +from tinyc.tpu.tile_passes import TILE_PASSES + + +def _expectations(func: Func, n: int, seed: int): + out = [] + for i in range(n): + inputs = ref.random_inputs(func, np.random.default_rng(seed + i)) + try: + out.append((inputs, ref.evaluate(func, inputs))) + except RangeError: + continue + return out + + +def _problem(s: Schedule, expectations) -> str | None: + program = to_json(s) + for inputs, want in expectations: + try: + got = interp.run_program(program, inputs)["outputs"] + except RangeError as e: + return f"now fails on inputs that used to work: {e}" + for name, w in zip(s.outputs, want): + if got[name].shape != w.shape or not (got[name] == w).all(): + return f"output '{name}' differs from the oracle on seeded random inputs" + return None + + +def compile_tile_ir( + func: Func, + opt_level: int = 1, + *, + verify_each: bool = False, + dump: Callable[[str, str], None] | None = None, + seed: int = 42, + num_checks: int = 4, +) -> Schedule: + """Lower a tpu-legalized Graph IR function and run the tile passes.""" + expectations = _expectations(func, num_checks, seed) if verify_each else [] + schedule = lower(func) + stages = [("lower-to-tile", lambda: None)] + if opt_level >= 1: + stages += [(n, (lambda p=p: p().run(schedule))) for n, p in TILE_PASSES.items()] + before = "" + for name, step in stages: + step() + text = print_schedule(schedule) + try: + verify_schedule(schedule) + except ScheduleError as e: + raise PassVerifyError(name, str(e), before, text) from None + if verify_each and (why := _problem(schedule, expectations)): + raise PassVerifyError(name, why, before, text) + if dump: + dump(name, text) + before = text + return schedule diff --git a/tinyc/tpu/cost.py b/tinyc/tpu/cost.py new file mode 100644 index 0000000..27075a0 --- /dev/null +++ b/tinyc/tpu/cost.py @@ -0,0 +1,47 @@ +"""Static cost model for a tile schedule. It must equal what the RTL executor measures.""" + +from __future__ import annotations + +from tinyc.tpu.tile_ir import ( + CAPACITY_MACS_PER_PASS, + CYCLES_PER_PASS, + TILE, + Acc, + Alloc, + Epilogue, + HostAdd, + HostRelu, + Pass, + Schedule, +) + +HOST_OP_KINDS = {Alloc: "alloc", Acc: "acc", Epilogue: "epilogue", HostAdd: "host_add", + HostRelu: "host_relu"} # fmt: skip + + +def useful_macs(s: Schedule, p: Pass) -> int: + """Multiply-adds in this pass that are not zero padding (rows x k x cols actually in range).""" + (m, k), n = s.buffers[p.a.buf].shape, s.buffers[p.b.buf].shape[1] + return ( + max(0, min(TILE, m - p.a.r)) * max(0, min(TILE, k - p.a.c)) * max(0, min(TILE, n - p.b.c)) + ) + + +def estimate(s: Schedule) -> dict: + passes = [c for c in s.commands if isinstance(c, Pass)] + host_ops = {kind: 0 for kind in HOST_OP_KINDS.values()} + for c in s.commands: + if type(c) in HOST_OP_KINDS: + host_ops[HOST_OP_KINDS[type(c)]] += 1 + useful = sum(useful_macs(s, p) for p in passes) + capacity = len(passes) * CAPACITY_MACS_PER_PASS + return { + "passes": len(passes), + "passes_skipped": s.meta.get("passes_skipped", 0), + "cycles_per_pass": CYCLES_PER_PASS, + "rtl_cycles": len(passes) * CYCLES_PER_PASS, + "host_ops": host_ops, + "useful_macs": useful, + "capacity_macs": capacity, + "utilization": round(useful / capacity, 4) if capacity else None, + } diff --git a/tinyc/tpu/emit.py b/tinyc/tpu/emit.py new file mode 100644 index 0000000..400622e --- /dev/null +++ b/tinyc/tpu/emit.py @@ -0,0 +1,72 @@ +"""Tile IR -> `.ttpu.json` (format documented in docs/TTPU_FORMAT.md, schema in docs/ttpu.schema.json).""" + +from __future__ import annotations + +import json + +from tinyc.tpu.cost import estimate +from tinyc.tpu.tile_ir import ( + TILE, + Acc, + Alloc, + Epilogue, + HostAdd, + HostRelu, + Pass, + Schedule, + TileRef, +) + +FORMAT = "ttpu" +VERSION = 1 + + +def _ref(r: TileRef) -> dict: + return {"buf": r.buf, "r": r.r, "c": r.c} + + +def _epilogue(d: dict, bias, relu) -> dict: + if bias is not None: + d["bias"] = bias + if relu: + d["relu"] = True + return d + + +def _command(c) -> dict: + if isinstance(c, Alloc): + return {"op": "alloc", "buf": c.buf} + if isinstance(c, Pass): + return {"op": "pass", "id": c.id, "a": _ref(c.a), "b": _ref(c.b)} + if isinstance(c, Acc): + return _epilogue({"op": "acc", "pass": c.pass_id, "dst": _ref(c.dst)}, c.bias, c.relu) + if isinstance(c, Epilogue): + return _epilogue({"op": "epilogue", "dst": _ref(c.dst)}, c.bias, c.relu) + if isinstance(c, HostAdd): + return {"op": "host_add", "dst": c.dst, "a": c.a, "b": c.b} + assert isinstance(c, HostRelu) + return {"op": "host_relu", "dst": c.dst, "src": c.src} + + +def to_json(s: Schedule) -> dict: + buffers = {} + for name, b in s.buffers.items(): + entry = {"dtype": b.dtype, "shape": list(b.shape), "kind": b.kind} + if b.kind == "const": + entry["data"] = b.data.tolist() + buffers[name] = entry + return { + "format": FORMAT, + "version": VERSION, + "tile": TILE, + "name": s.name, + "buffers": buffers, + "inputs": list(s.inputs), + "outputs": list(s.outputs), + "commands": [_command(c) for c in s.commands], + "estimate": estimate(s), + } + + +def dumps(program: dict) -> str: + return json.dumps(program, indent=2) + "\n" diff --git a/tinyc/tpu/interp.py b/tinyc/tpu/interp.py new file mode 100644 index 0000000..4bd6a18 --- /dev/null +++ b/tinyc/tpu/interp.py @@ -0,0 +1,92 @@ +"""Reference interpreter for `.ttpu.json`, written against the format spec (numpy, no chip). + +Independent of the JS executor (`runtime/ttpu-exec.mjs`): same spec, different language. Together +with the numpy oracle it triangulates the RTL run. The chip's pass is modelled as a plain tile +matmul; the RTL is the only place that actually checks the chip's behaviour. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import numpy as np + +from tinyc.ir.errors import InputError, RangeError +from tinyc.ir.types import RANGES + +I32_LO, I32_HI = RANGES["i32"] + + +def _fit(a: np.ndarray, what: str) -> np.ndarray: + if a.size and (a.min() < I32_LO or a.max() > I32_HI): + raise RangeError(f"{what} does not fit in i32") + return a + + +def _tile(buf: np.ndarray, ref: dict, t: int) -> np.ndarray: + out = np.zeros((t, t), dtype=np.int64) + sub = buf[ref["r"] : ref["r"] + t, ref["c"] : ref["c"] + t] + out[: sub.shape[0], : sub.shape[1]] = sub + return out + + +def _window(buf: np.ndarray, ref: dict, t: int): + r, c = ref["r"], ref["c"] + return buf[r : r + t, c : c + t] # a view, clipped to the buffer + + +def _epilogue(view: np.ndarray, cmd: dict, env, ref: dict) -> None: + if "bias" in cmd: + bias = env[cmd["bias"]][ref["c"] : ref["c"] + view.shape[1]] + view[:] = _fit(view + bias, "bias add") + if cmd.get("relu"): + view[:] = np.maximum(view, 0) + + +def run_program(program: dict, inputs: Mapping[str, object]) -> dict: + """Execute; returns {"outputs": {name: int64 array}, "passes": n, "host_ops": {...}}.""" + t = program["tile"] + env: dict[str, np.ndarray] = {} + for name, b in program["buffers"].items(): + shape, kind = tuple(b["shape"]), b["kind"] + if kind == "const": + env[name] = np.array(b["data"], dtype=np.int64) + elif kind == "input": + if name not in inputs: + raise InputError(f"missing input '{name}'") + arr = np.asarray(inputs[name], dtype=np.int64) + lo, hi = RANGES[b["dtype"]] + if arr.shape != shape or (arr.size and (arr.min() < lo or arr.max() > hi)): + raise InputError(f"input '{name}' must be {b['dtype']}{list(shape)}") + env[name] = arr + else: + env[name] = np.zeros(shape, dtype=np.int64) + + tmp: dict[int, np.ndarray] = {} + passes, host_ops = 0, {"alloc": 0, "acc": 0, "epilogue": 0, "host_add": 0, "host_relu": 0} + for cmd in program["commands"]: + op = cmd["op"] + if op == "pass": + a = _tile(env[cmd["a"]["buf"]], cmd["a"], t) + b = _tile(env[cmd["b"]["buf"]], cmd["b"], t) + tmp[cmd["id"]] = _fit(a @ b, "pass result") + passes += 1 + continue + host_ops[op] += 1 + if op == "alloc": + env[cmd["buf"]][:] = 0 + elif op == "acc": + view = _window(env[cmd["dst"]["buf"]], cmd["dst"], t) + view[:] = _fit(view + tmp[cmd["pass"]][: view.shape[0], : view.shape[1]], "accumulate") + _epilogue(view, cmd, env, cmd["dst"]) + elif op == "epilogue": + _epilogue(_window(env[cmd["dst"]["buf"]], cmd["dst"], t), cmd, env, cmd["dst"]) + elif op == "host_add": + env[cmd["dst"]][:] = _fit(env[cmd["a"]] + env[cmd["b"]], "add") + else: # host_relu + env[cmd["dst"]][:] = np.maximum(env[cmd["src"]], 0) + return { + "outputs": {n: env[n] for n in program["outputs"]}, + "passes": passes, + "host_ops": host_ops, + } diff --git a/tinyc/tpu/lower.py b/tinyc/tpu/lower.py new file mode 100644 index 0000000..74f822f --- /dev/null +++ b/tinyc/tpu/lower.py @@ -0,0 +1,71 @@ +"""Graph IR (legalized for tiny-tpu) -> Tile IR. + +Each matmul becomes one 4x4 pass per (output tile, K tile), in the order tiny-tpu's own +`tilingSchedule()` uses: output tiles row-major, K innermost. After every pass the host +accumulates the tile into the result; the last accumulation of an output tile carries the fused +bias/relu epilogue. `add`/`relu` ops that were not fused are whole-buffer host ops. +""" + +from __future__ import annotations + +from tinyc.ir.graph import Func +from tinyc.ir.types import epilogue +from tinyc.tpu.tile_ir import ( + TILE, + Acc, + Alloc, + Buffer, + HostAdd, + HostRelu, + Pass, + Schedule, + TileRef, +) + + +def lower(func: Func) -> Schedule: + s = Schedule( + name=func.name, + buffers={}, + commands=[], + inputs=[a.name for a in func.args], + outputs=[v.name for v in func.outputs], + ) + for a in func.args: + s.buffers[a.name] = Buffer(a.name, a.type.dtype, a.type.shape, "input") + next_pass = 0 + for op in func.ops: + r = op.result + if op.kind == "const": + s.buffers[r.name] = Buffer(r.name, r.type.dtype, r.type.shape, "const", op.data) + continue + s.buffers[r.name] = Buffer(r.name, r.type.dtype, r.type.shape, "temp") + if op.kind == "matmul": + a, b = op.operands[0], op.operands[1] + stages = epilogue(op.attrs) + bias = op.operands[2].name if "bias" in stages else None + relu = "relu" in stages + (m, k), n = a.type.shape, b.type.shape[1] + s.commands.append(Alloc(r.name)) + for i0 in range(0, m, TILE): + for j0 in range(0, n, TILE): + ks = list(range(0, k, TILE)) + for idx, k0 in enumerate(ks): + last = idx == len(ks) - 1 + s.commands.append( + Pass(next_pass, TileRef(a.name, i0, k0), TileRef(b.name, k0, j0)) + ) + s.commands.append( + Acc( + next_pass, + TileRef(r.name, i0, j0), + bias if last else None, + relu and last, + ) + ) + next_pass += 1 + elif op.kind == "add": + s.commands.append(HostAdd(r.name, op.operands[0].name, op.operands[1].name)) + else: # relu + s.commands.append(HostRelu(r.name, op.operands[0].name)) + return s diff --git a/tinyc/tpu/tile_ir.py b/tinyc/tpu/tile_ir.py new file mode 100644 index 0000000..0095d30 --- /dev/null +++ b/tinyc/tpu/tile_ir.py @@ -0,0 +1,207 @@ +"""Tile IR: the schedule of one program on tiny-tpu. + +A program is a list of commands over named buffers. `Pass` runs one 4x4 int8 tile product on +the chip (the only chip work); everything else (`Acc`, `Epilogue`, `HostAdd`, `HostRelu`, +`Alloc`) is *host* work and is labelled as such everywhere (PLAN §2.2: the chip always starts a +pass from psum = 0, so accumulating over K tiles, bias and relu happen on the host). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + +TILE = 4 +# Measured: TinyTpuSim.getCycleCount() after one pass (14 from start to done + 1 for the last +# element to latch). Pinned by tests/test_targets_contract.py. +CYCLES_PER_PASS = 15 +CAPACITY_MACS_PER_PASS = TILE * TILE * TILE + + +@dataclass(frozen=True) +class TileRef: + """A 4x4 window of `buf` at element offset (r, c); out-of-bounds elements read as 0.""" + + buf: str + r: int + c: int + + +@dataclass +class Buffer: + name: str + dtype: str + shape: tuple[int, ...] + kind: str # "input" | "const" | "temp" + data: np.ndarray | None = None # const only + + +@dataclass +class Alloc: + buf: str + + +@dataclass +class Pass: + id: int + a: TileRef + b: TileRef + + +@dataclass +class Acc: + """dst window += result of pass `pass_id`; then the optional epilogue (bias, relu).""" + + pass_id: int + dst: TileRef + bias: str | None = None + relu: bool = False + + +@dataclass +class Epilogue: + """Apply bias and/or relu to a dst window with no pass feeding it.""" + + dst: TileRef + bias: str | None = None + relu: bool = False + + +@dataclass +class HostAdd: + dst: str + a: str + b: str + + +@dataclass +class HostRelu: + dst: str + src: str + + +Command = Alloc | Pass | Acc | Epilogue | HostAdd | HostRelu + + +@dataclass +class Schedule: + name: str + buffers: dict[str, Buffer] + commands: list[Command] + inputs: list[str] + outputs: list[str] + meta: dict = field(default_factory=dict) # e.g. {"passes_skipped": 3} + + +class ScheduleError(Exception): + pass + + +def verify_schedule(s: Schedule) -> None: + """Structural checks. Raises ScheduleError listing every problem.""" + problems: list[str] = [] + bufs = s.buffers + + def buf(name, what, rank=None, dtype=None): + b = bufs.get(name) + if b is None: + problems.append(f"{what}: unknown buffer '{name}'") + elif rank is not None and len(b.shape) != rank: + problems.append(f"{what}: '{name}' must have rank {rank}, has {len(b.shape)}") + elif dtype is not None and b.dtype != dtype: + problems.append(f"{what}: '{name}' must be {dtype}, is {b.dtype}") + return b + + def window(ref, what): + b = buf(ref.buf, what, rank=2) + if ( + b is not None + and len(b.shape) == 2 + and not (0 <= ref.r < b.shape[0] and 0 <= ref.c < b.shape[1]) + ): + problems.append(f"{what}: window at ({ref.r},{ref.c}) lies outside '{ref.buf}'") + return b + + for name in s.inputs + s.outputs: + buf(name, "interface") + seen_passes: dict[int, Pass] = {} + consumed: set[int] = set() + for i, cmd in enumerate(s.commands): + where = f"command {i} ({type(cmd).__name__})" + if isinstance(cmd, Alloc): + buf(cmd.buf, where) + elif isinstance(cmd, Pass): + if cmd.id in seen_passes: + problems.append(f"{where}: duplicate pass id {cmd.id}") + seen_passes[cmd.id] = cmd + a, b = window(cmd.a, where + " a"), window(cmd.b, where + " b") + for x in (a, b): + if x is not None and x.dtype != "i8": + problems.append(f"{where}: operand '{x.name}' must be i8 (chip is int8 x int8)") + if a is not None and b is not None and cmd.a.c != cmd.b.r: + problems.append(f"{where}: A columns start at {cmd.a.c} but B rows at {cmd.b.r}") + elif isinstance(cmd, (Acc, Epilogue)): + dst = window(cmd.dst, where) + if dst is not None and dst.dtype != "i32": + problems.append(f"{where}: destination '{dst.name}' must be i32") + if cmd.bias is not None: + buf(cmd.bias, where + " bias", rank=1) + if isinstance(cmd, Acc): + if cmd.pass_id not in seen_passes: + problems.append(f"{where}: pass #{cmd.pass_id} not run yet") + elif cmd.pass_id in consumed: + problems.append(f"{where}: pass #{cmd.pass_id} accumulated twice") + consumed.add(cmd.pass_id) + elif isinstance(cmd, HostAdd): + for n in (cmd.dst, cmd.a, cmd.b): + buf(n, where) + elif isinstance(cmd, HostRelu): + for n in (cmd.dst, cmd.src): + buf(n, where) + for pid in seen_passes: + if pid not in consumed: + problems.append(f"pass #{pid} is never accumulated") + if problems: + raise ScheduleError("invalid tile schedule:\n " + "\n ".join(problems)) + + +# --- printing --------------------------------------------------------------------------------- + + +def _win(ref: TileRef) -> str: + return f"%{ref.buf}[{ref.r}:{ref.r + TILE},{ref.c}:{ref.c + TILE}]" + + +def _epi(bias, relu) -> str: + parts = ([f"bias=%{bias}"] if bias else []) + (["relu"] if relu else []) + return f" {{{', '.join(parts)}}}" if parts else "" + + +def print_schedule(s: Schedule) -> str: + args = ", ".join(f"%{n}: {s.buffers[n].dtype}[{','.join(map(str, s.buffers[n].shape))}]" + for n in s.inputs) # fmt: skip + outs = ", ".join(f"%{n}" for n in s.outputs) + lines = [f"tpu.schedule @{s.name}({args}) -> ({outs}) tile={TILE}x{TILE} {{"] + for b in s.buffers.values(): + if b.kind == "const": + t = f"{b.dtype}[{','.join(map(str, b.shape))}]" + lines.append(f" host.const %{b.name} : {t} = {b.data.tolist()!r}") + for cmd in s.commands: + if isinstance(cmd, Alloc): + b = s.buffers[cmd.buf] + lines.append(f" host.alloc %{b.name} : {b.dtype}[{','.join(map(str, b.shape))}] = 0") + elif isinstance(cmd, Pass): + lines.append( + f" tpu.pass #{cmd.id} a={_win(cmd.a)} b={_win(cmd.b)} // {CYCLES_PER_PASS} cyc" + ) + elif isinstance(cmd, Acc): + lines.append(f" host.acc {_win(cmd.dst)} += #{cmd.pass_id}{_epi(cmd.bias, cmd.relu)}") + elif isinstance(cmd, Epilogue): + lines.append(f" host.epilogue {_win(cmd.dst)}{_epi(cmd.bias, cmd.relu)}") + elif isinstance(cmd, HostAdd): + lines.append(f" host.add %{cmd.dst} = %{cmd.a}, %{cmd.b}") + else: + lines.append(f" host.relu %{cmd.dst} = %{cmd.src}") + lines.append("}") + return "\n".join(lines) + "\n" diff --git a/tinyc/tpu/tile_passes.py b/tinyc/tpu/tile_passes.py new file mode 100644 index 0000000..b3195e8 --- /dev/null +++ b/tinyc/tpu/tile_passes.py @@ -0,0 +1,69 @@ +"""Optimizations on the Tile IR.""" + +from __future__ import annotations + +import numpy as np + +from tinyc.tpu.tile_ir import TILE, Acc, Epilogue, Pass, Schedule, TileRef + + +def const_tile(s: Schedule, ref: TileRef) -> np.ndarray | None: + """The zero-padded 4x4 window of a constant buffer, or None if the buffer is not constant.""" + b = s.buffers[ref.buf] + if b.kind != "const": + return None + out = np.zeros((TILE, TILE), dtype=np.int64) + sub = b.data[ref.r : ref.r + TILE, ref.c : ref.c + TILE] + out[: sub.shape[0], : sub.shape[1]] = sub + return out + + +class SkipZeroPasses: + """Drop passes whose A or B tile is a constant all-zero tile: their product is zero. + + The matching accumulation goes too. If it carried the output tile's epilogue, the epilogue + moves to the previous kept accumulation of that tile; if there is none it becomes a + standalone epilogue (or disappears when it is relu only, since relu of the zero-initialized + tile is a no-op). Pass ids are renumbered afterwards. + """ + + name = "skip-zero-passes" + + def run(self, s: Schedule) -> None: + zero = set() + for cmd in s.commands: + if isinstance(cmd, Pass): + for ref in (cmd.a, cmd.b): + t = const_tile(s, ref) + if t is not None and not t.any(): + zero.add(cmd.id) + if not zero: + return + + out: list = [] + for cmd in s.commands: + if isinstance(cmd, Pass) and cmd.id in zero: + continue + if isinstance(cmd, Acc) and cmd.pass_id in zero: + if cmd.bias is not None or cmd.relu: + earlier = next( + (c for c in reversed(out) if isinstance(c, Acc) and c.dst == cmd.dst), None + ) + if earlier is not None: + earlier.bias, earlier.relu = cmd.bias, cmd.relu + elif cmd.bias is not None: + out.append(Epilogue(cmd.dst, cmd.bias, cmd.relu)) + continue + out.append(cmd) + + renumber = {c.id: i for i, c in enumerate(c for c in out if isinstance(c, Pass))} + for c in out: + if isinstance(c, Pass): + c.id = renumber[c.id] + elif isinstance(c, Acc): + c.pass_id = renumber[c.pass_id] + s.commands = out + s.meta["passes_skipped"] = s.meta.get("passes_skipped", 0) + len(zero) + + +TILE_PASSES = {SkipZeroPasses.name: SkipZeroPasses}