From 1de65b2c7c3ff1bfb92ae41c198ddeecfddfc1ad Mon Sep 17 00:00:00 2001 From: Dineth Hettiarachchi Date: Sun, 25 Oct 2026 01:43:25 +0530 Subject: [PATCH] Phase 4: tiny-gpu backend (memory plan, Kernel IR, registers, .tgpu) Lower the legalized Graph IR to tiny-gpu SIMT kernels: one thread per output element (block = row, thread = column, tid = flat index), one kernel per op with fused epilogues staying a single launch. Add a liveness-based memory planner, the Kernel IR with printer, verifier and flatten-to-instructions, a linear-scan register allocator that never writes R5-R7 (asserted, and checked on the simulator), `.tgpu` assembly with source positions, a launch document with JSON Schema and spec, a static cost model, and an independent Kernel IR interpreter whose counts must equal the simulator's exactly. `tinyc --target gpu` gains --emit kernel|tgpu|launch, --run (real simulator, seeded inputs, measured vs estimated), --stats and --gpu-regs. The runner runs multi-launch programs with memory carried over and a fresh simulator per kernel. All examples, including the chained-matmul mlp_gpu, run bit-exact, and the same program gives identical results on tiny-gpu and tiny-tpu. Measured tiny-gpu facts are pinned in the target contract test: a label starting with R is parsed as a register, a loop of K iterations costs K cycles, relu costs a cycle only when a thread does not clamp, and branching to the end of a program costs a stray cycle. Tests: kernel IR and verifier, memory plan (incl. random programs), register allocator (loop-aware liveness, pressure errors), schema validation, every example on the real simulator, estimate bounds attained at data extremes, pinned cost numbers, and Hypothesis properties over generated programs. --- CLAUDE.md | 1 + CONTEXT.md | 9 +- README.md | 14 +- docs/PASSES.md | 9 + docs/PLAN.md | 6 +- docs/TGPU_FORMAT.md | 62 ++++++ docs/launch.schema.json | 85 ++++++++ runtime/gpu_runner.py | 68 ++++++ tests/golden/kernel/big_mm.kernel | 20 ++ tests/golden/kernel/const_mm.kernel | 20 ++ tests/golden/kernel/mlp_gpu.kernel | 41 ++++ tests/golden/kernel/mm.kernel | 17 ++ tests/golden/kernel/mm_bias_relu.kernel | 23 +++ tests/golden/kernel/mm_relu.kernel | 20 ++ tests/golden/tgpu/big_mm.tgpu | 24 +++ tests/golden/tgpu/const_mm.tgpu | 24 +++ tests/golden/tgpu/mlp_gpu.tgpu | 48 +++++ tests/golden/tgpu/mm.tgpu | 20 ++ tests/golden/tgpu/mm_bias_relu.tgpu | 27 +++ tests/golden/tgpu/mm_relu.tgpu | 24 +++ tests/gpu/helpers.py | 32 +++ tests/gpu/test_gpu_emit.py | 86 ++++++++ tests/gpu/test_gpu_perf.py | 24 +++ tests/gpu/test_gpu_properties.py | 49 +++++ tests/gpu/test_gpu_sim.py | 129 ++++++++++++ tests/gpu/test_kernel_ir.py | 142 +++++++++++++ tests/gpu/test_memplan.py | 100 +++++++++ tests/gpu/test_regalloc.py | 117 +++++++++++ tests/perf_baseline.json | 186 +++++++++++++++++ tests/test_cli.py | 54 +++++ tests/test_targets_contract.py | 48 +++++ tinyc/cli.py | 95 ++++++--- tinyc/driver.py | 32 +++ tinyc/gpu/__init__.py | 115 +++++++++++ tinyc/gpu/cost.py | 82 ++++++++ tinyc/gpu/emit.py | 80 +++++++ tinyc/gpu/interp.py | 144 +++++++++++++ tinyc/gpu/kernel_ir.py | 263 ++++++++++++++++++++++++ tinyc/gpu/lower.py | 123 +++++++++++ tinyc/gpu/memplan.py | 120 +++++++++++ tinyc/gpu/regalloc.py | 112 ++++++++++ tinyc/ir/ref.py | 13 ++ tinyc/passes/legalize_gpu.py | 4 + tinyc/passes/manager.py | 12 +- tinyc/tpu/__init__.py | 15 +- 45 files changed, 2684 insertions(+), 55 deletions(-) create mode 100644 docs/TGPU_FORMAT.md create mode 100644 docs/launch.schema.json create mode 100644 tests/golden/kernel/big_mm.kernel create mode 100644 tests/golden/kernel/const_mm.kernel create mode 100644 tests/golden/kernel/mlp_gpu.kernel create mode 100644 tests/golden/kernel/mm.kernel create mode 100644 tests/golden/kernel/mm_bias_relu.kernel create mode 100644 tests/golden/kernel/mm_relu.kernel create mode 100644 tests/golden/tgpu/big_mm.tgpu create mode 100644 tests/golden/tgpu/const_mm.tgpu create mode 100644 tests/golden/tgpu/mlp_gpu.tgpu create mode 100644 tests/golden/tgpu/mm.tgpu create mode 100644 tests/golden/tgpu/mm_bias_relu.tgpu create mode 100644 tests/golden/tgpu/mm_relu.tgpu create mode 100644 tests/gpu/helpers.py create mode 100644 tests/gpu/test_gpu_emit.py create mode 100644 tests/gpu/test_gpu_perf.py create mode 100644 tests/gpu/test_gpu_properties.py create mode 100644 tests/gpu/test_gpu_sim.py create mode 100644 tests/gpu/test_kernel_ir.py create mode 100644 tests/gpu/test_memplan.py create mode 100644 tests/gpu/test_regalloc.py create mode 100644 tinyc/gpu/__init__.py create mode 100644 tinyc/gpu/cost.py create mode 100644 tinyc/gpu/emit.py create mode 100644 tinyc/gpu/interp.py create mode 100644 tinyc/gpu/kernel_ir.py create mode 100644 tinyc/gpu/lower.py create mode 100644 tinyc/gpu/memplan.py create mode 100644 tinyc/gpu/regalloc.py diff --git a/CLAUDE.md b/CLAUDE.md index b8db3f6..5737de9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -19,6 +19,7 @@ source scripts/env.sh tinyc examples/mm_relu.tiny --locs # optimized Graph IR (--emit ast, -O0, --target gpu|tpu, # --print-after-all, --verify-each, --time-passes, --disable PASS) tinyc examples/mm_relu.tiny --target tpu --run # compile + run on the real RTL (also --emit tile|ttpu, --stats) +tinyc examples/mm_relu.tiny --target gpu --run # compile + run on the simulator (--emit kernel|tgpu|launch, --gpu-regs N) pytest -q # full suite pytest --update-golden # rewrite tests/golden/* after reviewing the diff cd web && pnpm dev | build | lint | typecheck | test diff --git a/CONTEXT.md b/CONTEXT.md index d80977a..949d769 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -18,7 +18,11 @@ faked; host-side work is always labelled as host work. - Phase 3 (TPU backend): DONE. 287 tests green. Graph IR -> Tile IR (`tinyc/tpu/`) -> `.ttpu.json` (schema `docs/ttpu.schema.json`, spec `docs/TTPU_FORMAT.md`); `runtime/ttpu-exec.mjs` runs whole programs on the real RTL; cost model == measurement (pinned in `tests/perf_baseline.json`). `tinyc FILE --target tpu --run` works. -- Next: Phase 4 (GPU backend -O0: Kernel IR, memory planner, regalloc, `.tgpu`, gpu_runner). +- Phase 4 (GPU backend, -O0): DONE. 443 tests green. Graph IR -> memory plan -> Kernel IR -> linear-scan registers -> `.tgpu` + launch JSON + (`tinyc/gpu/`, schema `docs/launch.schema.json`, spec `docs/TGPU_FORMAT.md`). Every example, including the chained-matmul + `mlp_gpu`, runs bit-exact on the real simulator; the same program gives identical results on both chips. `tinyc FILE --target gpu --run`. +- Next: Phase 5 (GPU optimizations -O1/-O2: LICM, strength reduction, unrolling, shared-memory tiling). `-O` currently + controls only the graph passes; decide how it maps to kernel optimizations before adding them. ## Fixed decisions (do not re-litigate; change only via docs/ARCHITECTURE.md) | Decision | Choice | @@ -33,6 +37,9 @@ faked; host-side work is always labelled as host work. | TPU mapping | 4×4 int8 passes, zero padding, K-accumulation + relu/bias as `host.*` ops | | TPU executor | Exactly one: `runtime/ttpu-exec.mjs` (Node in tests, browser in UI). `tinyc/tpu/interp.py` is a second, independent *reference* reader of the same JSON spec, never a second executor | | TPU host work | `acc`/`epilogue`/`host_add`/`host_relu`/`alloc` are host commands, counted separately from chip passes everywhere | +| GPU mapping | One thread per output element: block = row (R5), thread = column (R6), tid (R7) = flat index. One kernel per Graph op (fused = one). Kernel code is read-only on R5-R7 (asserted + checked at run time) | +| GPU cycles | A thread needs taken-branches+1 cycles, a launch the max. Report instructions with cycles. `estimate` gives min/max (relu is data dependent), loads/stores exact; `gpu.interp` counts must equal the simulator's | +| GPU asm rules | Labels are `L` (a label starting with R is parsed as a register); a kernel never ends with an `if` (stray extra cycle); one fresh simulator per launch at its final size, memory carried over | | TPU cost | `passes * 15` RTL cycles; utilization = useful MACs (`M*K*N`) / (64 * passes) | | Overflow | `range-check` proves every value fits i32 by intervals; otherwise reject (tiny-gpu raises on overflow). i32 *inputs* are assumed within ±2^20 (`PassContext.i32_input_range`); runners must respect it | | Fusion | `matmul{epilogue=bias,relu}`; the fused op takes the outermost absorbed op's loc. Backends must lower the epilogue | diff --git a/README.md b/README.md index 498fa3d..3806031 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,17 @@ One tiny tensor language, two real chips: `C = relu(A @ B)` compiles to [tiny-gpu](https://github.com/deaneeth/tiny-gpu) assembly or to a [tiny-tpu](https://github.com/deaneeth/tiny-tpu) tile schedule, and runs on the real simulator / RTL. -**Status: Phase 3.** The frontend, optimizing pass pipeline and the **tiny-tpu backend** work: -`tinyc examples/mm_relu.tiny --target tpu --run` compiles to tile passes and runs them on the real -RTL, bit-exact against numpy. The tiny-gpu backend comes next. +**Status: Phase 4.** The frontend, optimizing pass pipeline and **both backends** work. One program, +two chips, bit-exact against numpy on each: + +```bash +tinyc examples/mm_relu.tiny --target tpu --run # tile passes on the real tiny-tpu RTL +tinyc examples/mm_relu.tiny --target gpu --run # SIMT kernels on the tiny-gpu simulator +``` + +GPU kernel optimizations are next. Both targets can be driven from pytest. See [docs/PLAN.md](docs/PLAN.md) for the roadmap, -[docs/LANGUAGE.md](docs/LANGUAGE.md), [docs/IR.md](docs/IR.md) and [docs/PASSES.md](docs/PASSES.md) and [docs/TTPU_FORMAT.md](docs/TTPU_FORMAT.md). +[docs/LANGUAGE.md](docs/LANGUAGE.md), [docs/IR.md](docs/IR.md) and [docs/PASSES.md](docs/PASSES.md) and [docs/TTPU_FORMAT.md](docs/TTPU_FORMAT.md) and [docs/TGPU_FORMAT.md](docs/TGPU_FORMAT.md). ```bash git clone --recurse-submodules && cd tiny-compiler diff --git a/docs/PASSES.md b/docs/PASSES.md index 57ebc02..c88e41d 100644 --- a/docs/PASSES.md +++ b/docs/PASSES.md @@ -63,3 +63,12 @@ a pass whose A or B window is a constant all-zero tile is removed together with to the previous kept accumulation, becomes a standalone `epilogue` (bias present), or disappears (relu only: relu of the zero-initialized tile is a no-op). Pass ids are renumbered. `--verify-each` validates the schedule after lowering and after each tile pass, by interpreting the emitted JSON. + +## GPU backend stages +After the graph pipeline, `compile_gpu_program` runs: memory plan (`gpu/memplan.py`) -> lowering to +Kernel IR (`gpu/lower.py`, one kernel per op) -> verify -> flatten -> linear-scan register +allocation (`gpu/regalloc.py`) -> `.tgpu` emission (`gpu/emit.py`). With `--verify-each` the +structured kernels are interpreted against the numpy oracle after lowering, and the emitted assembly +is run on the real simulator against the oracle after emission. `--print-after-all` shows +`lower-to-kernel` and `regalloc-emit`. Phase 4 emits the straightforward (-O0) kernel code; kernel +optimizations (loop-invariant motion, unrolling, shared-memory tiling) are Phase 5. diff --git a/docs/PLAN.md b/docs/PLAN.md index a7ce42a..6649bf1 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -198,7 +198,11 @@ Static estimates printed by `tinyc --stats` and shown before you press Run: - GPU: instructions issued, simulator cycles (as tiny-gpu defines them), global/shared loads and stores, registers used. After running, measured numbers sit next to the estimates. **Estimate and measurement must -match exactly** for the TPU and for GPU instruction counts. That is a test (§5, layer 6). +match exactly** for the TPU. On the GPU, loads and stores are exact; instructions and cycles are +**bounds** (`min`/`max`), because relu's branch is data dependent (a clamping thread runs one more +instruction, a non-clamping thread takes a skip branch that adds a cycle). The exact numbers for +given inputs come from the Kernel IR interpreter, whose counts must equal the simulator's (§5, +layer 6). --- diff --git a/docs/TGPU_FORMAT.md b/docs/TGPU_FORMAT.md new file mode 100644 index 0000000..576ac7b --- /dev/null +++ b/docs/TGPU_FORMAT.md @@ -0,0 +1,62 @@ +# tiny-gpu output: `.tgpu` kernels and the launch document + +`tinyc FILE --target gpu --emit tgpu` prints the assembly; `--emit launch` prints the launch +document (machine-checked by `docs/launch.schema.json`); `--emit kernel` prints the Kernel IR. + +## Mapping +One thread per output element. A rank-2 result `[M, N]` launches **M blocks of N threads**: +`R5` (block id) is the row, `R6` (thread in block) is the column, and `R7` (global thread id) is +the flat row-major index. Rank-1 results use one block. Each Graph IR op is one kernel launch +(a fused `matmul{epilogue=...}` is one kernel); launches share global memory. + +`R5`, `R6`, `R7` are set by the simulator and **only ever read**. Kernel code uses +`R0-R4, R8..R(n-1)`; an assertion in the allocator and a runtime check in the runner +(`reserved_intact`) enforce this. Nothing is passed as kernel arguments: all addresses are +immediates from the memory plan. + +## What tiny-gpu has and does not have +`SET ADD MUL LD ST CMP BRLT BRGT` (and `JMP BEQ BNE BRZ SYNC SYNCB SHLD SHST CSWAP`). There is no +SUB, DIV, shift or MAX, so: +- loops count up: `SET k,0 / L: ... / ADD k,k,1 / CMP k,K / BRLT L`; +- `relu(x)` is `CMP x,0 / BRGT skip / SET x,0 / skip:`; +- there is no base+offset addressing: every access computes its address in a register. + +Assembler traps found by measurement (and enforced by the emitter): +- **a label must not start with `R`/`r`**: tiny-gpu's assembler parses it as a register and + crashes (`relu_end:` -> `ValueError`). Labels here are `L0`, `L1`, ...; +- **a kernel must not end with an `if`**: branching to the end of the program costs a stray extra + cycle; the Kernel IR verifier rejects it. + +## Cycles vs instructions +tiny-gpu runs consecutive non-branch instructions of a thread in one step; a thread's step ends at +a *taken* branch. A thread needs `taken branches + 1` cycles, a launch the maximum over threads. So +a straight-line kernel is 1 cycle however long it is, and a K-iteration loop is K cycles. +**Always read cycles together with instructions.** The relu `if` is data dependent: a thread that +clamps runs one more instruction (`SET`); a thread that does not takes the skip branch, which adds a +cycle. The static `estimate` therefore gives `min`/`max` for instructions and cycles; loads and +stores are exact. `gpu.interp` gives exact counts for given inputs, and tests require them to equal +the simulator's. + +## Launch document +```json +{ "format": "tgpu-launch", "version": 1, "name": "main", "num_registers": 16, "mem_size": 140, + "buffers": { "X": {"addr": 0, "dtype": "i8", "shape": [4,8], "kind": "input"}, + "2": {"addr": 104, "dtype": "i32", "shape": [4,6], "kind": "temp"} }, + "inputs": ["X"], "outputs": ["3"], + "kernels": [ {"name": "k0_matmul", "value": "2", "op": "matmul", "grid": [4,6], + "shared_size": 0, "asm": "SET R0, 0 ; ..."} ], + "estimate": { ... } } +``` +Memory is word addressed (`i8` takes one int32 word). Constants (`kind: const`, with `data`) and +inputs are the initial memory image; everything else starts at 0. Intermediates reuse the words of +buffers whose last reader has finished (`gpu/memplan.py`); an op's result never overlaps what that +same op reads. + +## Executor +`runtime/gpu_runner.py::run_launch` runs a launch document on the real simulator: a fresh +`TinyGPU` per kernel built at its final thread count (the simulator's `set_grid` resize leaves +R7 = 0 for added threads), global memory carried over. It counts instructions per opcode without +changing semantics. + +## Source mapping +Each instruction ends with `; line:col`, the source position of the Graph IR op it came from. diff --git a/docs/launch.schema.json b/docs/launch.schema.json new file mode 100644 index 0000000..c62e0ef --- /dev/null +++ b/docs/launch.schema.json @@ -0,0 +1,85 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/deaneeth/tiny-compiler/docs/launch.schema.json", + "title": "tgpu launch document (version 1)", + "type": "object", + "required": ["format", "version", "name", "num_registers", "mem_size", "buffers", "inputs", "outputs", "kernels", "estimate"], + "additionalProperties": false, + "properties": { + "format": { "const": "tgpu-launch" }, + "version": { "const": 1 }, + "name": { "type": "string" }, + "num_registers": { "type": "integer", "minimum": 8 }, + "mem_size": { "type": "integer", "minimum": 0 }, + "buffers": { "type": "object", "additionalProperties": { "$ref": "#/$defs/buffer" } }, + "inputs": { "type": "array", "items": { "type": "string" } }, + "outputs": { "type": "array", "items": { "type": "string" }, "minItems": 1 }, + "kernels": { "type": "array", "items": { "$ref": "#/$defs/kernel" } }, + "estimate": { "$ref": "#/$defs/estimate" } + }, + "$defs": { + "buffer": { + "type": "object", + "required": ["addr", "dtype", "shape", "kind"], + "additionalProperties": false, + "properties": { + "addr": { "type": "integer", "minimum": 0 }, + "dtype": { "enum": ["i8", "i32"] }, + "shape": { "type": "array", "items": { "type": "integer", "minimum": 1 }, "minItems": 1, "maxItems": 2 }, + "kind": { "enum": ["input", "const", "temp"] }, + "data": { "type": "array", "description": "const buffers only: nested lists matching shape" } + } + }, + "kernel": { + "type": "object", + "required": ["name", "value", "op", "grid", "shared_size", "asm"], + "additionalProperties": false, + "properties": { + "name": { "type": "string" }, + "value": { "type": "string", "description": "the Graph IR value this kernel computes" }, + "op": { "enum": ["matmul", "add", "relu"] }, + "grid": { "type": "array", "items": { "type": "integer", "minimum": 1 }, "minItems": 2, "maxItems": 2 }, + "shared_size": { "type": "integer", "minimum": 0 }, + "asm": { "type": "string" } + } + }, + "range": { + "type": "object", + "required": ["min", "max"], + "additionalProperties": false, + "properties": { "min": { "type": "integer", "minimum": 0 }, "max": { "type": "integer", "minimum": 0 } } + }, + "estimate": { + "type": "object", + "required": ["launches", "threads", "instructions", "cycles", "loads", "stores", "registers", "per_kernel"], + "additionalProperties": false, + "properties": { + "launches": { "type": "integer", "minimum": 0 }, + "threads": { "type": "integer", "minimum": 0 }, + "instructions": { "$ref": "#/$defs/range" }, + "cycles": { "$ref": "#/$defs/range" }, + "loads": { "type": "integer", "minimum": 0 }, + "stores": { "type": "integer", "minimum": 0 }, + "registers": { "type": "integer", "minimum": 0 }, + "per_kernel": { + "type": "array", + "items": { + "type": "object", + "required": ["name", "grid", "threads", "instructions", "cycles", "loads", "stores", "registers"], + "additionalProperties": false, + "properties": { + "name": { "type": "string" }, + "grid": { "type": "array", "items": { "type": "integer" }, "minItems": 2, "maxItems": 2 }, + "threads": { "type": "integer", "minimum": 1 }, + "instructions": { "$ref": "#/$defs/range" }, + "cycles": { "$ref": "#/$defs/range" }, + "loads": { "type": "integer", "minimum": 0 }, + "stores": { "type": "integer", "minimum": 0 }, + "registers": { "type": "integer", "minimum": 0 } + } + } + } + } + } + } +} diff --git a/runtime/gpu_runner.py b/runtime/gpu_runner.py index 20c3b91..eed0eb0 100644 --- a/runtime/gpu_runner.py +++ b/runtime/gpu_runner.py @@ -28,6 +28,7 @@ class GpuResult: cycles: int # tiny-gpu "cycles": a cycle ends at a branch/jump/sync, not at every instruction instructions: int # thread-instructions executed; use this next to `cycles` (PLAN §2.1) op_counts: Counter = field(default_factory=Counter) + reserved_intact: bool = True # R5 (block), R6 (thread), R7 (tid) still hold what the sim set def read(self, addr: int, length: int) -> list[int]: return [int(v) for v in self.memory[addr : addr + length]] @@ -91,9 +92,76 @@ def run_tgpu( if gpu.active.any(): raise RuntimeError(f"kernel still running after {max_cycles} cycles") + tid = np.arange(num_blocks * tpb) + expected = {5: tid // tpb, 6: tid % tpb, 7: tid} + intact = all((gpu.registers[:, r] == want).all() for r, want in expected.items()) return GpuResult( memory=gpu.memory.copy(), cycles=len(gpu.history_registers), instructions=sum(counts.values()), op_counts=counts, + reserved_intact=intact, ) + + +@dataclass +class LaunchResult: + outputs: dict[str, np.ndarray] # int64, original shapes + kernels: list[tuple[str, GpuResult]] # (kernel name, its measurements), in launch order + reserved_intact: bool + + @property + def cycles(self) -> int: + return sum(r.cycles for _, r in self.kernels) + + @property + def instructions(self) -> int: + return sum(r.instructions for _, r in self.kernels) + + @property + def op_counts(self) -> Counter: + total: Counter = Counter() + for _, r in self.kernels: + total.update(r.op_counts) + return total + + +def run_launch(launch: Mapping, inputs: Mapping[str, object]) -> LaunchResult: + """Run a launch document (`docs/TGPU_FORMAT.md`) on the real tiny-gpu simulator. + + Global memory persists across kernels; each kernel gets a fresh simulator built at its final + thread count (tiny-gpu's `set_grid` resize leaves R7 = 0 for threads it adds). + """ + memory = np.zeros(launch["mem_size"], dtype=np.int64) + for name, b in launch["buffers"].items(): + if b["kind"] == "const": + flat = np.asarray(b["data"], dtype=np.int64).reshape(-1) + elif b["kind"] == "input": + arr = np.asarray(inputs[name], dtype=np.int64) + if list(arr.shape) != b["shape"]: + raise ValueError(f"input '{name}' must have shape {b['shape']}") + flat = arr.reshape(-1) + else: + continue + memory[b["addr"] : b["addr"] + len(flat)] = flat + + results, intact = [], True + for k in launch["kernels"]: + res = run_tgpu( + k["asm"], + grid=tuple(k["grid"]), + mem_size=launch["mem_size"], + memory_init={0: memory}, + num_registers=launch["num_registers"], + shared_size=k["shared_size"], + ) + memory = res.memory.astype(np.int64) + intact = intact and res.reserved_intact + results.append((k["name"], res)) + + outputs = {} + for name in launch["outputs"]: + b = launch["buffers"][name] + size = int(np.prod(b["shape"])) + outputs[name] = memory[b["addr"] : b["addr"] + size].reshape(b["shape"]) + return LaunchResult(outputs, results, intact) diff --git a/tests/golden/kernel/big_mm.kernel b/tests/golden/kernel/big_mm.kernel new file mode 100644 index 0000000..d79a058 --- /dev/null +++ b/tests/golden/kernel/big_mm.kernel @@ -0,0 +1,20 @@ +kernel @k0_matmul grid=(12 blocks, 12 threads) // %1 = matmul at 4:5 { + v0 = set 0 + for v1 in 0..12 { + v2 = mul block, 12 + v2 = add v2, v1 + v2 = add v2, 0 + v3 = ld [v2] + v4 = mul v1, 12 + v4 = add v4, thread + v4 = add v4, 144 + v5 = ld [v4] + v6 = mul v3, v5 + v0 = add v0, v6 + } + if le v0, 0 { + v0 = set 0 + } + v7 = add tid, 288 + st [v7], v0 +} diff --git a/tests/golden/kernel/const_mm.kernel b/tests/golden/kernel/const_mm.kernel new file mode 100644 index 0000000..05393e8 --- /dev/null +++ b/tests/golden/kernel/const_mm.kernel @@ -0,0 +1,20 @@ +kernel @k0_matmul grid=(3 blocks, 2 threads) // %2 = matmul at 5:5 { + v0 = set 0 + for v1 in 0..2 { + v2 = mul block, 2 + v2 = add v2, v1 + v2 = add v2, 0 + v3 = ld [v2] + v4 = mul v1, 2 + v4 = add v4, thread + v4 = add v4, 6 + v5 = ld [v4] + v6 = mul v3, v5 + v0 = add v0, v6 + } + if le v0, 0 { + v0 = set 0 + } + v7 = add tid, 10 + st [v7], v0 +} diff --git a/tests/golden/kernel/mlp_gpu.kernel b/tests/golden/kernel/mlp_gpu.kernel new file mode 100644 index 0000000..4f9ee57 --- /dev/null +++ b/tests/golden/kernel/mlp_gpu.kernel @@ -0,0 +1,41 @@ +kernel @k0_matmul grid=(4 blocks, 6 threads) // %2 = matmul at 6:5 { + v0 = set 0 + for v1 in 0..8 { + v2 = mul block, 8 + v2 = add v2, v1 + v2 = add v2, 0 + v3 = ld [v2] + v4 = mul v1, 6 + v4 = add v4, thread + v4 = add v4, 32 + v5 = ld [v4] + v6 = mul v3, v5 + v0 = add v0, v6 + } + v7 = add thread, 80 + v8 = ld [v7] + v0 = add v0, v8 + if le v0, 0 { + v0 = set 0 + } + v9 = add tid, 104 + st [v9], v0 +} + +kernel @k1_matmul grid=(4 blocks, 3 threads) // %3 = matmul at 7:5 { + v0 = set 0 + for v1 in 0..6 { + v2 = mul block, 6 + v2 = add v2, v1 + v2 = add v2, 104 + v3 = ld [v2] + v4 = mul v1, 3 + v4 = add v4, thread + v4 = add v4, 86 + v5 = ld [v4] + v6 = mul v3, v5 + v0 = add v0, v6 + } + v7 = add tid, 128 + st [v7], v0 +} diff --git a/tests/golden/kernel/mm.kernel b/tests/golden/kernel/mm.kernel new file mode 100644 index 0000000..d9d7c24 --- /dev/null +++ b/tests/golden/kernel/mm.kernel @@ -0,0 +1,17 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %0 = matmul at 4:5 { + v0 = set 0 + for v1 in 0..4 { + v2 = mul block, 4 + v2 = add v2, v1 + v2 = add v2, 0 + v3 = ld [v2] + v4 = mul v1, 4 + v4 = add v4, thread + v4 = add v4, 16 + v5 = ld [v4] + v6 = mul v3, v5 + v0 = add v0, v6 + } + v7 = add tid, 32 + st [v7], v0 +} diff --git a/tests/golden/kernel/mm_bias_relu.kernel b/tests/golden/kernel/mm_bias_relu.kernel new file mode 100644 index 0000000..4e615f9 --- /dev/null +++ b/tests/golden/kernel/mm_bias_relu.kernel @@ -0,0 +1,23 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %2 = matmul at 5:5 { + v0 = set 0 + for v1 in 0..8 { + v2 = mul block, 8 + v2 = add v2, v1 + v2 = add v2, 0 + v3 = ld [v2] + v4 = mul v1, 4 + v4 = add v4, thread + v4 = add v4, 32 + v5 = ld [v4] + v6 = mul v3, v5 + v0 = add v0, v6 + } + v7 = add thread, 64 + v8 = ld [v7] + v0 = add v0, v8 + if le v0, 0 { + v0 = set 0 + } + v9 = add tid, 68 + st [v9], v0 +} diff --git a/tests/golden/kernel/mm_relu.kernel b/tests/golden/kernel/mm_relu.kernel new file mode 100644 index 0000000..478be65 --- /dev/null +++ b/tests/golden/kernel/mm_relu.kernel @@ -0,0 +1,20 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %1 = matmul at 4:5 { + v0 = set 0 + for v1 in 0..8 { + v2 = mul block, 8 + v2 = add v2, v1 + v2 = add v2, 0 + v3 = ld [v2] + v4 = mul v1, 4 + v4 = add v4, thread + v4 = add v4, 32 + v5 = ld [v4] + v6 = mul v3, v5 + v0 = add v0, v6 + } + if le v0, 0 { + v0 = set 0 + } + v7 = add tid, 64 + st [v7], v0 +} diff --git a/tests/golden/tgpu/big_mm.tgpu b/tests/golden/tgpu/big_mm.tgpu new file mode 100644 index 0000000..e652426 --- /dev/null +++ b/tests/golden/tgpu/big_mm.tgpu @@ -0,0 +1,24 @@ +; kernel k0_matmul: %1 = matmul (t.tiny:4:5) +; grid: 12 blocks x 12 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +SET R1, 0 ; 4:5 +L0: +MUL R2, R5, 12 ; 4:5 +ADD R2, R2, R1 ; 4:5 +ADD R2, R2, 0 ; 4:5 +LD R2, R2 ; 4:5 +MUL R3, R1, 12 ; 4:5 +ADD R3, R3, R6 ; 4:5 +ADD R3, R3, 144 ; 4:5 +LD R3, R3 ; 4:5 +MUL R2, R2, R3 ; 4:5 +ADD R0, R0, R2 ; 4:5 +ADD R1, R1, 1 ; 4:5 +CMP R1, 12 ; 4:5 +BRLT L0 ; 4:5 +CMP R0, 0 ; 4:5 +BRGT L1 ; 4:5 +SET R0, 0 ; 4:5 +L1: +ADD R1, R7, 288 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/const_mm.tgpu b/tests/golden/tgpu/const_mm.tgpu new file mode 100644 index 0000000..363035e --- /dev/null +++ b/tests/golden/tgpu/const_mm.tgpu @@ -0,0 +1,24 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:5:5) +; grid: 3 blocks x 2 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:5 +SET R1, 0 ; 5:5 +L0: +MUL R2, R5, 2 ; 5:5 +ADD R2, R2, R1 ; 5:5 +ADD R2, R2, 0 ; 5:5 +LD R2, R2 ; 5:5 +MUL R3, R1, 2 ; 5:5 +ADD R3, R3, R6 ; 5:5 +ADD R3, R3, 6 ; 5:5 +LD R3, R3 ; 5:5 +MUL R2, R2, R3 ; 5:5 +ADD R0, R0, R2 ; 5:5 +ADD R1, R1, 1 ; 5:5 +CMP R1, 2 ; 5:5 +BRLT L0 ; 5:5 +CMP R0, 0 ; 5:5 +BRGT L1 ; 5:5 +SET R0, 0 ; 5:5 +L1: +ADD R1, R7, 10 ; 5:5 +ST R1, R0 ; 5:5 diff --git a/tests/golden/tgpu/mlp_gpu.tgpu b/tests/golden/tgpu/mlp_gpu.tgpu new file mode 100644 index 0000000..0f2369b --- /dev/null +++ b/tests/golden/tgpu/mlp_gpu.tgpu @@ -0,0 +1,48 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:6:5) +; grid: 4 blocks x 6 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 6:5 +SET R1, 0 ; 6:5 +L0: +MUL R2, R5, 8 ; 6:5 +ADD R2, R2, R1 ; 6:5 +ADD R2, R2, 0 ; 6:5 +LD R2, R2 ; 6:5 +MUL R3, R1, 6 ; 6:5 +ADD R3, R3, R6 ; 6:5 +ADD R3, R3, 32 ; 6:5 +LD R3, R3 ; 6:5 +MUL R2, R2, R3 ; 6:5 +ADD R0, R0, R2 ; 6:5 +ADD R1, R1, 1 ; 6:5 +CMP R1, 8 ; 6:5 +BRLT L0 ; 6:5 +ADD R1, R6, 80 ; 6:5 +LD R1, R1 ; 6:5 +ADD R0, R0, R1 ; 6:5 +CMP R0, 0 ; 6:5 +BRGT L1 ; 6:5 +SET R0, 0 ; 6:5 +L1: +ADD R1, R7, 104 ; 6:5 +ST R1, R0 ; 6:5 + +; kernel k1_matmul: %3 = matmul (t.tiny:7:5) +; grid: 4 blocks x 3 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 7:5 +SET R1, 0 ; 7:5 +L0: +MUL R2, R5, 6 ; 7:5 +ADD R2, R2, R1 ; 7:5 +ADD R2, R2, 104 ; 7:5 +LD R2, R2 ; 7:5 +MUL R3, R1, 3 ; 7:5 +ADD R3, R3, R6 ; 7:5 +ADD R3, R3, 86 ; 7:5 +LD R3, R3 ; 7:5 +MUL R2, R2, R3 ; 7:5 +ADD R0, R0, R2 ; 7:5 +ADD R1, R1, 1 ; 7:5 +CMP R1, 6 ; 7:5 +BRLT L0 ; 7:5 +ADD R1, R7, 128 ; 7:5 +ST R1, R0 ; 7:5 diff --git a/tests/golden/tgpu/mm.tgpu b/tests/golden/tgpu/mm.tgpu new file mode 100644 index 0000000..f08787a --- /dev/null +++ b/tests/golden/tgpu/mm.tgpu @@ -0,0 +1,20 @@ +; kernel k0_matmul: %0 = matmul (t.tiny:4:5) +; grid: 4 blocks x 4 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +SET R1, 0 ; 4:5 +L0: +MUL R2, R5, 4 ; 4:5 +ADD R2, R2, R1 ; 4:5 +ADD R2, R2, 0 ; 4:5 +LD R2, R2 ; 4:5 +MUL R3, R1, 4 ; 4:5 +ADD R3, R3, R6 ; 4:5 +ADD R3, R3, 16 ; 4:5 +LD R3, R3 ; 4:5 +MUL R2, R2, R3 ; 4:5 +ADD R0, R0, R2 ; 4:5 +ADD R1, R1, 1 ; 4:5 +CMP R1, 4 ; 4:5 +BRLT L0 ; 4:5 +ADD R1, R7, 32 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/mm_bias_relu.tgpu b/tests/golden/tgpu/mm_bias_relu.tgpu new file mode 100644 index 0000000..3a22e4e --- /dev/null +++ b/tests/golden/tgpu/mm_bias_relu.tgpu @@ -0,0 +1,27 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:5:5) +; grid: 4 blocks x 4 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:5 +SET R1, 0 ; 5:5 +L0: +MUL R2, R5, 8 ; 5:5 +ADD R2, R2, R1 ; 5:5 +ADD R2, R2, 0 ; 5:5 +LD R2, R2 ; 5:5 +MUL R3, R1, 4 ; 5:5 +ADD R3, R3, R6 ; 5:5 +ADD R3, R3, 32 ; 5:5 +LD R3, R3 ; 5:5 +MUL R2, R2, R3 ; 5:5 +ADD R0, R0, R2 ; 5:5 +ADD R1, R1, 1 ; 5:5 +CMP R1, 8 ; 5:5 +BRLT L0 ; 5:5 +ADD R1, R6, 64 ; 5:5 +LD R1, R1 ; 5:5 +ADD R0, R0, R1 ; 5:5 +CMP R0, 0 ; 5:5 +BRGT L1 ; 5:5 +SET R0, 0 ; 5:5 +L1: +ADD R1, R7, 68 ; 5:5 +ST R1, R0 ; 5:5 diff --git a/tests/golden/tgpu/mm_relu.tgpu b/tests/golden/tgpu/mm_relu.tgpu new file mode 100644 index 0000000..9b98046 --- /dev/null +++ b/tests/golden/tgpu/mm_relu.tgpu @@ -0,0 +1,24 @@ +; kernel k0_matmul: %1 = matmul (t.tiny:4:5) +; grid: 4 blocks x 4 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +SET R1, 0 ; 4:5 +L0: +MUL R2, R5, 8 ; 4:5 +ADD R2, R2, R1 ; 4:5 +ADD R2, R2, 0 ; 4:5 +LD R2, R2 ; 4:5 +MUL R3, R1, 4 ; 4:5 +ADD R3, R3, R6 ; 4:5 +ADD R3, R3, 32 ; 4:5 +LD R3, R3 ; 4:5 +MUL R2, R2, R3 ; 4:5 +ADD R0, R0, R2 ; 4:5 +ADD R1, R1, 1 ; 4:5 +CMP R1, 8 ; 4:5 +BRLT L0 ; 4:5 +CMP R0, 0 ; 4:5 +BRGT L1 ; 4:5 +SET R0, 0 ; 4:5 +L1: +ADD R1, R7, 64 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/gpu/helpers.py b/tests/gpu/helpers.py new file mode 100644 index 0000000..fdedc2e --- /dev/null +++ b/tests/gpu/helpers.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from tinyc.driver import compile_gpu +from tinyc.ir import ref + +ROOT = Path(__file__).resolve().parent.parent.parent +EXAMPLES = sorted((ROOT / "examples").glob("*.tiny")) # every example is legal on tiny-gpu + + +def build(source: str, opt: int = 1, regs: int = 16, **kw): + from tinyc.passes import PassContext + + return compile_gpu( + source, "t.tiny", opt_level=opt, ctx=PassContext(gpu_num_registers=regs), **kw + ) + + +def run_on_simulator(comp, seed: int = 1, inputs=None): + """Run on the real tiny-gpu simulator. Returns (LaunchResult, inputs, oracle outputs).""" + from runtime.gpu_runner import run_launch + + inputs = inputs or ref.random_inputs(comp.func, np.random.default_rng(seed)) + return run_launch(comp.launch, inputs), inputs, ref.evaluate(comp.func, inputs) + + +def assert_bit_exact(comp, result, want): + for name, w in zip(comp.launch["outputs"], want): + np.testing.assert_array_equal(result.outputs[name], w) diff --git a/tests/gpu/test_gpu_emit.py b/tests/gpu/test_gpu_emit.py new file mode 100644 index 0000000..0123b85 --- /dev/null +++ b/tests/gpu/test_gpu_emit.py @@ -0,0 +1,86 @@ +import copy +import json +import re + +import jsonschema +import pytest + +from tests.gpu.helpers import EXAMPLES, ROOT, build +from tinyc.gpu.emit import dumps + +SCHEMA = json.loads((ROOT / "docs" / "launch.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", EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +def test_every_launch_document_matches_the_schema(path, opt): + comp = build(path.read_text(), opt=opt) + validate(comp.launch) + assert comp.launch["estimate"] == comp.program.estimate + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_tgpu_assembly_golden(path, golden): + golden( + ROOT / "tests" / "golden" / "tgpu" / f"{path.stem}.tgpu", + build(path.read_text()).program.asm_text(), + ) + + +def test_emission_is_deterministic(): + src = (ROOT / "examples" / "mlp_gpu.tiny").read_text() + assert dumps(build(src).launch) == dumps(build(src).launch) + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_every_instruction_carries_its_source_position(path): + for asm in build(path.read_text()).program.asm.values(): + for line in asm.splitlines(): + if line and not line.startswith(";") and not line.endswith(":"): + assert re.search(r"; \d+:\d+$", line), line + + +def test_header_comments_and_labels_are_valid_for_the_assembler(): + asm = build("input A : i8[2, 2]\ninput B : i8[2, 2]\nC = relu(A @ B)\noutput C\n").program.asm[ + "k0_matmul" + ] + assert asm.startswith("; kernel k0_matmul") and "reserved" in asm.splitlines()[1] + labels = [ln for ln in asm.splitlines() if ln.endswith(":")] + assert labels == ["L0:", "L1:"] + + +@pytest.mark.parametrize( + "mutate", + [ + lambda d: d.__setitem__("version", 2), + lambda d: d.__setitem__("num_registers", 4), + lambda d: d.__setitem__("extra", 1), + lambda d: d["kernels"][0].__setitem__("op", "conv"), + lambda d: d["kernels"][0].__setitem__("grid", [4]), + lambda d: d["kernels"][0].pop("asm"), + lambda d: d["buffers"]["A"].__setitem__("dtype", "f32"), + lambda d: d["buffers"]["A"].__setitem__("addr", -1), + lambda d: d.__setitem__("outputs", []), + lambda d: d["estimate"]["cycles"].pop("max"), + ], +) +def test_the_schema_rejects_malformed_documents(mutate): + doc = copy.deepcopy(build((ROOT / "examples" / "mm_relu.tiny").read_text()).launch) + validate(doc) + mutate(doc) + with pytest.raises(jsonschema.ValidationError): + validate(doc) + + +def test_a_program_with_no_ops_is_a_valid_launch_with_zero_kernels(): + comp = build("const W : i8[2, 2] = [[1, 2], [3, 4]]\noutput W\n") + validate(comp.launch) + assert comp.launch["kernels"] == [] and comp.launch["buffers"]["0"]["kind"] == "const" diff --git a/tests/gpu/test_gpu_perf.py b/tests/gpu/test_gpu_perf.py new file mode 100644 index 0000000..1a02700 --- /dev/null +++ b/tests/gpu/test_gpu_perf.py @@ -0,0 +1,24 @@ +"""Pinned static cost numbers: an unexpected change in launches/instructions/loads fails CI.""" + +from tests.gpu.helpers import EXAMPLES, build + + +def test_gpu_cost_numbers_are_pinned(perf_pin): + pinned = {} + for opt in (0, 1): + pinned[f"O{opt}"] = {} + for path in EXAMPLES: + est = build(path.read_text(), opt=opt).program.estimate + pinned[f"O{opt}"][path.stem] = { + k: est[k] + for k in ( + "launches", + "threads", + "instructions", + "cycles", + "loads", + "stores", + "registers", + ) + } + perf_pin("gpu", pinned) diff --git a/tests/gpu/test_gpu_properties.py b/tests/gpu/test_gpu_properties.py new file mode 100644 index 0000000..3410d81 --- /dev/null +++ b/tests/gpu/test_gpu_properties.py @@ -0,0 +1,49 @@ +"""Random programs through the GPU backend: the interpreter, and the real simulator.""" + +import numpy as np +from hypothesis import HealthCheck, event, given, settings +from hypothesis import strategies as st + +from tests.gen import programs, tpu_programs +from tests.gpu.helpers import build, run_on_simulator +from tinyc.gpu import interp +from tinyc.ir import ref +from tinyc.ir.errors import PassError + +SOURCES = st.one_of(programs(), tpu_programs()) + + +def _compile(source, opt): + try: + return build( + source, opt=opt, verify_each=True + ) # kernel interpreter AND simulator each stage + except PassError: + event("rejected by range-check or legalization") + return None + + +@settings(max_examples=120, deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given(SOURCES, st.sampled_from([0, 1]), st.integers(0, 2**32 - 1)) +def test_kernel_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)["outputs"] + for name, want in zip(comp.launch["outputs"], ref.evaluate(comp.func, inputs)): + np.testing.assert_array_equal(got[name], want) + + +@settings(max_examples=40, deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given(SOURCES, st.sampled_from([0, 1]), st.integers(0, 2**32 - 1)) +def test_simulator_matches_oracle_and_interpreter_counts_on_random_programs(source, opt, seed): + comp = _compile(source, opt) + if comp is None: + return + res, inputs, want = run_on_simulator(comp, seed) + assert res.reserved_intact + for name, w in zip(comp.launch["outputs"], want): + np.testing.assert_array_equal(res.outputs[name], w) + for (_, sim), counts in zip(res.kernels, interp.run_program(comp.program, inputs)["kernels"]): + assert dict(sim.op_counts) == dict(counts.ops) and sim.cycles == counts.cycles diff --git a/tests/gpu/test_gpu_sim.py b/tests/gpu/test_gpu_sim.py new file mode 100644 index 0000000..ef585f7 --- /dev/null +++ b/tests/gpu/test_gpu_sim.py @@ -0,0 +1,129 @@ +"""The real thing: every program runs on the unmodified tiny-gpu simulator.""" + +import numpy as np +import pytest + +from runtime.gpu_runner import run_launch +from runtime.tpu_runner import run_ttpu +from tests.gpu.helpers import EXAMPLES, ROOT, assert_bit_exact, build, run_on_simulator +from tinyc.driver import compile_tpu +from tinyc.gpu import interp +from tinyc.ir import ref +from tinyc.ir.errors import PassError + +MM_RELU = "input A : i8[4, 8]\ninput B : i8[8, 4]\nC = relu(A @ B)\noutput C\n" + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +def test_examples_are_bit_exact_and_never_touch_the_reserved_registers(path, opt): + comp = build(path.read_text(), opt=opt) + for seed in (1, 2): + res, _, want = run_on_simulator(comp, seed) + assert_bit_exact(comp, res, want) + assert res.reserved_intact + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +def test_estimate_bounds_contain_the_measurement_and_loads_stores_are_exact(path, opt): + comp = build(path.read_text(), opt=opt) + est = comp.program.estimate + res, _, _ = run_on_simulator(comp, 7) + assert est["instructions"]["min"] <= res.instructions <= est["instructions"]["max"] + assert est["cycles"]["min"] <= res.cycles <= est["cycles"]["max"] + assert res.op_counts["LD"] == est["loads"] and res.op_counts["ST"] == est["stores"] + for (name, r), per in zip(res.kernels, est["per_kernel"]): + assert name == per["name"] + assert per["cycles"]["min"] <= r.cycles <= per["cycles"]["max"] + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +@pytest.mark.parametrize("opt", [0, 1]) +def test_interpreter_counts_equal_simulator_counts_exactly(path, opt): + """Checks lowering, register allocation and emission together: the structured Kernel IR, + interpreted without any registers, must execute exactly what the simulator executes.""" + comp = build(path.read_text(), opt=opt) + res, inputs, _ = run_on_simulator(comp, 3) + got = interp.run_program(comp.program, inputs) + for (name, sim), counts in zip(res.kernels, got["kernels"]): + assert name == counts.name + assert dict(sim.op_counts) == dict(counts.ops) + assert (sim.instructions, sim.cycles) == (counts.instructions, counts.cycles) + for n in comp.launch["outputs"]: + np.testing.assert_array_equal(got["outputs"][n], res.outputs[n]) + + +def test_the_estimate_bounds_are_attained_at_the_data_extremes(): + comp = build(MM_RELU) + est = comp.program.estimate + ones = np.ones((4, 8), dtype=np.int64) + # every product positive: nothing clamps -> fewest instructions, but every thread takes the skip branch + keep = run_launch(comp.launch, {"A": ones, "B": np.ones((8, 4), dtype=np.int64)}) + assert keep.instructions == est["instructions"]["min"] and keep.cycles == est["cycles"]["max"] + # every product negative: everything clamps -> one more instruction per thread, one cycle fewer + clamp = run_launch(comp.launch, {"A": ones, "B": -np.ones((8, 4), dtype=np.int64)}) + assert clamp.instructions == est["instructions"]["max"] and clamp.cycles == est["cycles"]["min"] + assert (clamp.outputs["1"] == 0).all() + + +def test_a_straight_line_kernel_is_one_cycle_however_many_instructions(): + comp = build("input a : i8[8]\ninput b : i8[8]\nC = a + b\noutput C\n") + res, _, _ = run_on_simulator(comp) + assert res.cycles == 1 and res.instructions == 8 * 7 # 7 instructions per thread, 8 threads + + +def test_matmul_cycles_are_the_loop_trip_count_plus_the_relu_cycle(): + comp = build(MM_RELU) # K = 8 + res, _, _ = run_on_simulator(comp, 1) + assert res.cycles in (8, 9) # 8 iterations; +1 when some thread does not clamp + + +def test_global_memory_persists_across_launches(): + comp = build((ROOT / "examples" / "mlp_gpu.tiny").read_text()) + assert len(comp.launch["kernels"]) == 2 + res, _, want = run_on_simulator(comp) # layer 2 reads what layer 1 wrote + assert_bit_exact(comp, res, want) + + +def test_unfused_o0_chain_matches_the_fused_o1_program(): + src = (ROOT / "examples" / "mm_bias_relu.tiny").read_text() + o0, o1 = build(src, opt=0), build(src, opt=1) + assert len(o0.launch["kernels"]) == 3 and len(o1.launch["kernels"]) == 1 + inputs = ref.random_inputs(o0.func, np.random.default_rng(9)) + a, b = run_launch(o0.launch, inputs), run_launch(o1.launch, inputs) + for n0, n1 in zip(o0.launch["outputs"], o1.launch["outputs"]): + np.testing.assert_array_equal(a.outputs[n0], b.outputs[n1]) + assert b.instructions < a.instructions # fusion removes launches and memory traffic + + +def test_the_smallest_legal_register_budget_still_works(): + comp = build(MM_RELU, regs=8) + res, _, want = run_on_simulator(comp) + assert_bit_exact(comp, res, want) + + +def test_fewer_than_eight_registers_is_rejected_at_compile_time(): + with pytest.raises(PassError, match="at least 8 registers"): + build(MM_RELU, regs=7) + + +def test_extreme_int8_values_do_not_overflow_the_simulator(): + 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)} + res = run_launch(comp.launch, inputs) + assert (res.outputs["0"] == 8 * 128 * 128).all() + + +# --- the series' point: one program, two chips ------------------------------------------------ + + +@pytest.mark.parametrize("path", [p for p in EXAMPLES if p.stem != "mlp_gpu"], ids=lambda p: p.stem) +def test_tiny_gpu_and_tiny_tpu_agree_on_the_same_program_and_inputs(path): + gpu = build(path.read_text()) + tpu = compile_tpu(path.read_text(), path.name) + inputs = ref.random_inputs(gpu.func, np.random.default_rng(21)) + on_gpu = run_launch(gpu.launch, inputs) + on_tpu = run_ttpu(tpu.program, inputs) + for g, t in zip(gpu.launch["outputs"], tpu.schedule.outputs): + np.testing.assert_array_equal(on_gpu.outputs[g], on_tpu.outputs[t]) diff --git a/tests/gpu/test_kernel_ir.py b/tests/gpu/test_kernel_ir.py new file mode 100644 index 0000000..9eafc85 --- /dev/null +++ b/tests/gpu/test_kernel_ir.py @@ -0,0 +1,142 @@ +import re + +import pytest + +from tests.gpu.helpers import EXAMPLES, ROOT, build +from tinyc.gpu.kernel_ir import ( + Assign, + For, + If, + Kernel, + KernelError, + Load, + Special, + Store, + VReg, + flatten, + print_kernel, + verify_kernel, +) + +v = [VReg(i) for i in range(8)] + + +def kernel(body, grid=(1, 4)): + return Kernel("k", grid, body, "0", "relu") + + +def kernels(source, opt=1): + return build(source, opt=opt).program.kernels + + +# --- lowering --------------------------------------------------------------------------------- + + +def test_each_graph_op_is_one_kernel_and_fusion_removes_launches(): + src = "input A : i8[4, 8]\ninput B : i8[8, 4]\nC = relu(A @ B)\noutput C\n" + assert [k.kind for k in kernels(src, opt=0)] == ["matmul", "relu"] + assert [k.kind for k in kernels(src, opt=1)] == ["matmul"] # fused: one launch + + +def test_grid_is_rows_by_columns_and_rank1_uses_one_block(): + (k,) = kernels("input A : i8[4, 8]\ninput B : i8[8, 6]\nC = A @ B\noutput C\n") + assert k.grid == (4, 6) and k.threads == 24 + (k,) = kernels("input a : i8[5]\ninput b : i8[5]\nC = a + b\noutput C\n") + assert k.grid == (1, 5) + + +def test_broadcast_add_indexes_the_bias_by_column_and_same_shape_by_flat_id(): + text = print_kernel(kernels("input M : i8[3, 4]\ninput b : i8[4]\nY = M + b\noutput Y\n")[0]) + assert "add thread," in text and "add tid," in text + + +def test_chained_matmuls_are_two_launches_over_shared_memory(): + ks = kernels((ROOT / "examples" / "mlp_gpu.tiny").read_text()) + assert [(k.kind, k.grid) for k in ks] == [("matmul", (4, 6)), ("matmul", (4, 3))] + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_kernel_ir_golden(path, golden): + prog = build(path.read_text()).program + text = "\n".join(print_kernel(k) for k in prog.kernels) + golden(ROOT / "tests" / "golden" / "kernel" / f"{path.stem}.kernel", text) + + +def test_matmul_kernel_has_the_naive_o0_shape(): + text = print_kernel(kernels("input A : i8[4, 8]\ninput B : i8[8, 4]\nC = A @ B\noutput C\n")[0]) + assert "for v1 in 0..8 {" in text + assert ( + text.count("= ld [") == 2 and text.count("= mul ") == 3 + ) # A and B recomputed every iteration + + +# --- flattening idioms -------------------------------------------------------------------------- + + +def test_for_and_if_flatten_to_tiny_gpus_idioms(): + k = kernel( + [ + Assign("set", v[0], (0,)), + For(v[1], 0, 3, [Assign("add", v[0], (v[0], v[1]))]), + If(v[0], 0, [Assign("set", v[0], (0,))]), + Store(Special("tid"), v[0]), + ] + ) + ops = [(i.op, i.label) for i in flatten(k)] + assert ops == [ + ("SET", None), ("SET", None), ("LABEL", "L0"), ("ADD", None), ("ADD", None), ("CMP", None), + ("BRLT", "L0"), ("CMP", None), ("BRGT", "L1"), ("SET", None), ("LABEL", "L1"), ("ST", None), + ] # fmt: skip + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_labels_never_start_with_r_because_the_assembler_would_parse_them_as_registers(path): + for k in build(path.read_text()).program.kernels: + for inst in flatten(k): + if inst.label: + assert re.match(r"L\d+$", inst.label), inst.label + + +# --- the verifier ----------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "body, fragment", + [ + ([Assign("add", v[0], (v[1], 1)), Store(Special("tid"), v[0])], "v1 used before"), + ([Store(Special("tid"), v[0])], "v0 used before"), + ([For(v[0], 3, 3, [Assign("set", v[1], (0,))]), Store(0, v[1])], "stop > start"), + ([Assign("set", v[0], (2**31,)), Store(0, v[0])], "does not fit in i32"), + ([Assign("set", v[0], (1, 2)), Store(0, v[0])], "'set' with 2 operands"), + ([Assign("neg", v[0], (1,)), Store(0, v[0])], "'neg'"), + ([Assign("set", v[0], (Special("warp"),)), Store(0, v[0])], "unknown special"), + ( + [Assign("set", v[0], (0,)), If(v[0], 0, [Assign("set", v[0], (1,))])], + "must not end with an `if`", + ), + ( # assigned only inside an `if`: not definitely assigned afterwards + [Assign("set", v[0], (0,)), If(v[0], 0, [Assign("set", v[1], (1,))]), Store(0, v[1])], + "v1 used before", + ), + ], +) +def test_kernel_verifier_rejects(body, fragment): + with pytest.raises(KernelError, match=fragment): + verify_kernel(kernel(body)) + + +def test_kernel_verifier_rejects_oversized_grids(): + with pytest.raises(KernelError, match="1024-thread"): + verify_kernel(kernel([Store(0, 0)], grid=(33, 32))) + with pytest.raises(KernelError, match="positive"): + verify_kernel(kernel([Store(0, 0)], grid=(0, 4))) + + +def test_values_assigned_in_a_for_body_are_definite_after_it_because_it_runs_at_least_once(): + verify_kernel(kernel([For(v[0], 0, 2, [Assign("set", v[1], (5,))]), Store(0, v[1])])) + + +def test_load_and_store_operands_are_checked(): + verify_kernel(kernel([Load(v[0], 3), Store(Special("tid"), v[0])])) + with pytest.raises(KernelError, match="used before"): + verify_kernel(kernel([Load(v[0], v[5]), Store(0, v[0])])) diff --git a/tests/gpu/test_memplan.py b/tests/gpu/test_memplan.py new file mode 100644 index 0000000..2b09d41 --- /dev/null +++ b/tests/gpu/test_memplan.py @@ -0,0 +1,100 @@ +import pytest +from hypothesis import HealthCheck, event, given, settings + +from tests.gen import programs +from tests.helpers import ir +from tinyc.driver import compile_graph +from tinyc.gpu.memplan import MemPlan, Region, check_plan, plan_memory +from tinyc.ir.errors import PassError + + +def plan(text): + f = ir(text) + p = plan_memory(f) + check_plan(f, p) + return p + + +CHAIN = """ + func @main(%A: i8[2,2]) -> (i8[2,2]) { + %0 = relu %A : i8[2,2] + %1 = relu %0 : i8[2,2] + %2 = relu %1 : i8[2,2] + %3 = relu %2 : i8[2,2] + return %3 + } + """ + + +def test_inputs_then_constants_are_laid_out_first_and_pinned(): + p = plan(""" + func @main(%A: i8[2,3], %b: i32[3]) -> (i32[2,3]) { + %0 = const [[1, 2, 3], [4, 5, 6]] : i8[2,3] + %1 = add %A, %0 : i32[2,3] + return %1 + } + """) + assert (p.buffers["A"].addr, p.buffers["b"].addr, p.buffers["0"].addr) == (0, 6, 9) + assert p.buffers["0"].kind == "const" and p.buffers["A"].kind == "input" + + +def test_dead_intermediates_are_reused_but_outputs_are_not(): + p = plan(CHAIN) + # %0 and %1 cannot share (both live at op 1); %2 reuses %0's words; %3 (output) is pinned + assert p.buffers["0"].addr == 4 and p.buffers["1"].addr == 8 + assert p.buffers["2"].addr == p.buffers["0"].addr + assert p.mem_size == 12 # 4 input + 3 live temps... not the naive 4 + 4*4 = 20 + + +def test_a_result_never_overlaps_what_its_own_op_reads(): + p = plan(CHAIN) + for dst, src in (("1", "0"), ("2", "1"), ("3", "2")): + a, b = p.buffers[dst], p.buffers[src] + assert a.addr >= b.addr + b.size or b.addr >= a.addr + a.size + + +def test_an_output_that_is_an_input_or_constant_needs_no_new_memory(): + p = plan("func @main(%A: i8[3]) -> (i8[3]) {\n return %A\n}") + assert p.mem_size == 3 + + +def test_check_plan_detects_overlap_of_simultaneously_live_buffers(): + f = ir(CHAIN) + p = plan_memory(f) + p.buffers["1"] = Region("1", p.buffers["0"].addr, 4, "i8", (2, 2), "temp") # collide with %0 + with pytest.raises(AssertionError, match="overlap"): + check_plan(f, p) + + +def test_the_memory_budget_is_enforced_by_legalization(): + from tinyc.passes import PassContext + + with pytest.raises(PassError, match="words of tiny-gpu memory"): + compile_graph( + "input A : i8[30, 30]\nB = relu(A)\noutput B\n", "t.tiny", target="gpu", + ctx=PassContext(gpu_mem_words=100), + ) # fmt: skip + + +def test_the_backend_checks_the_planned_size_too_when_legalization_is_disabled(): + from tinyc.driver import compile_gpu + from tinyc.passes import PassContext + + with pytest.raises(PassError, match="words of tiny-gpu memory"): + compile_gpu( + "input A : i8[30, 30]\nB = relu(A)\noutput B\n", "t.tiny", + disabled=["legalize-gpu"], ctx=PassContext(gpu_mem_words=100), + ) # fmt: skip + + +@settings(max_examples=150, deadline=None, suppress_health_check=[HealthCheck.too_slow]) +@given(programs()) +def test_plans_for_random_programs_never_overlap_live_buffers(source): + try: + func, _ = compile_graph(source, "g.tiny", target="gpu") + except PassError: + event("rejected") + return + p = plan_memory(func) + check_plan(func, p) + assert isinstance(p, MemPlan) and p.mem_size >= sum(a.type.size for a in func.args) diff --git a/tests/gpu/test_regalloc.py b/tests/gpu/test_regalloc.py new file mode 100644 index 0000000..3787773 --- /dev/null +++ b/tests/gpu/test_regalloc.py @@ -0,0 +1,117 @@ +import pytest + +from tests.gpu.helpers import EXAMPLES, build +from tinyc.gpu.kernel_ir import ( + Assign, + For, + Kernel, + Special, + Store, + VInst, + VReg, + flatten, +) +from tinyc.gpu.regalloc import ( + RegisterPressure, + allocatable, + allocate, + assert_reserved_untouched, + live_intervals, +) + +v = [VReg(i) for i in range(12)] +TID = Special("tid") + + +def alloc(body, regs=16): + insts = flatten(Kernel("k", (1, 4), body, "0", "relu")) + return insts, allocate(insts, regs, "k") + + +def test_the_pool_excludes_the_reserved_registers(): + assert allocatable(16) == [0, 1, 2, 3, 4, 8, 9, 10, 11, 12, 13, 14, 15] + assert allocatable(8) == [0, 1, 2, 3, 4] + + +def test_fewer_than_8_registers_is_rejected_because_r7_holds_the_thread_id(): + with pytest.raises(ValueError, match="at least 8"): + allocatable(7) + + +def test_sequential_temporaries_share_one_register(): + body = [ + Assign("set", v[0], (1,)), + Assign("add", v[1], (v[0], 1)), # v0 dies here, v1 may reuse its register + Assign("add", v[2], (v[1], 1)), + Store(TID, v[2]), + ] + _, a = alloc(body) + assert a.registers_used == 1 and a.max_live == 1 + + +def test_values_live_at_the_same_time_get_distinct_registers(): + body = [Assign("set", v[i], (i,)) for i in range(4)] + body += [Assign("add", v[4], (v[0], v[1])), Assign("add", v[5], (v[2], v[3])), + Assign("add", v[6], (v[4], v[5])), Store(TID, v[6])] # fmt: skip + _, a = alloc(body) + assert len({a.assignment[v[i]] for i in range(4)}) == 4 + + +def test_a_value_used_inside_a_loop_stays_live_until_the_back_edge(): + body = [ + Assign("set", v[0], (0,)), # accumulator: defined before, read and written every iteration + For(v[1], 0, 3, [Assign("set", v[2], (7,)), Assign("add", v[0], (v[0], v[2]))]), + Store(TID, v[0]), + ] + insts = flatten(Kernel("k", (1, 4), body, "0", "relu")) + spans = live_intervals(insts) + back_edge = max(i for i, x in enumerate(insts) if x.op == "BRLT") + store = max(i for i, x in enumerate(insts) if x.op == "ST") + assert spans[v[0]][1] == store >= back_edge # live through the loop and after it + # the loop-local temp must not steal the accumulator's or the counter's register + _, a = alloc(body) + assert len({a.assignment[v[0]], a.assignment[v[1]], a.assignment[v[2]]}) == 3 + + +def test_a_value_defined_before_a_loop_and_last_used_inside_survives_to_the_back_edge(): + body = [ + Assign("set", v[0], (5,)), # read once per iteration, never after the loop + Assign("set", v[3], (0,)), + For(v[1], 0, 3, [Assign("add", v[3], (v[3], v[0]))]), + Store(TID, v[3]), + ] + insts = flatten(Kernel("k", (1, 4), body, "0", "relu")) + back_edge = max(i for i, x in enumerate(insts) if x.op == "BRLT") + assert live_intervals(insts)[v[0]][1] == back_edge + + +def test_running_out_of_registers_is_an_error_that_says_how_many_were_needed(): + body = [Assign("set", v[i], (i,)) for i in range(6)] + body += [Assign("add", v[6], (v[i], v[i + 1])) for i in range(5)] + body += [Assign("add", v[7], (v[0], v[5])), Store(TID, v[7])] + with pytest.raises(RegisterPressure, match=r"needs more than 5 registers.*--gpu-regs"): + alloc(body, regs=8) + alloc(body, regs=16) # the same kernel fits with the default budget + + +def test_assertion_catches_writes_to_reserved_registers(): + insts = [VInst("SET", v[0], (1,))] + with pytest.raises(AssertionError, match="reserved"): + assert_reserved_untouched(insts, {v[0]: 7}) + assert_reserved_untouched(insts, {v[0]: 4}) + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_real_kernels_never_use_r5_to_r7_as_destinations(path): + import re + + for asm in build(path.read_text()).program.asm.values(): + for line in asm.splitlines(): + m = re.match(r"(SET|ADD|MUL|LD) R(\d+),", line) + if m: + assert int(m.group(2)) not in (5, 6, 7), line + + +def test_matmul_needs_only_four_registers_and_fits_the_smallest_legal_budget(): + comp = build("input A : i8[4, 8]\ninput B : i8[8, 4]\nC = relu(A @ B)\noutput C\n", regs=8) + assert comp.program.registers == {"k0_matmul": 4} diff --git a/tests/perf_baseline.json b/tests/perf_baseline.json index 66600a5..bea457c 100644 --- a/tests/perf_baseline.json +++ b/tests/perf_baseline.json @@ -1,4 +1,190 @@ { + "gpu": { + "O0": { + "big_mm": { + "cycles": { + "max": 14, + "min": 13 + }, + "instructions": { + "max": 24048, + "min": 23904 + }, + "launches": 2, + "loads": 3600, + "registers": 4, + "stores": 288, + "threads": 288 + }, + "const_mm": { + "cycles": { + "max": 4, + "min": 3 + }, + "instructions": { + "max": 222, + "min": 216 + }, + "launches": 2, + "loads": 30, + "registers": 4, + "stores": 12, + "threads": 12 + }, + "mlp_gpu": { + "cycles": { + "max": 17, + "min": 16 + }, + "instructions": { + "max": 3912, + "min": 3888 + }, + "launches": 4, + "loads": 600, + "registers": 4, + "stores": 84, + "threads": 84 + }, + "mm": { + "cycles": { + "max": 4, + "min": 4 + }, + "instructions": { + "max": 896, + "min": 896 + }, + "launches": 1, + "loads": 128, + "registers": 4, + "stores": 16, + "threads": 16 + }, + "mm_bias_relu": { + "cycles": { + "max": 11, + "min": 10 + }, + "instructions": { + "max": 1952, + "min": 1936 + }, + "launches": 3, + "loads": 304, + "registers": 4, + "stores": 48, + "threads": 48 + }, + "mm_relu": { + "cycles": { + "max": 10, + "min": 9 + }, + "instructions": { + "max": 1840, + "min": 1824 + }, + "launches": 2, + "loads": 272, + "registers": 4, + "stores": 32, + "threads": 32 + } + }, + "O1": { + "big_mm": { + "cycles": { + "max": 13, + "min": 12 + }, + "instructions": { + "max": 23472, + "min": 23328 + }, + "launches": 1, + "loads": 3456, + "registers": 4, + "stores": 144, + "threads": 144 + }, + "const_mm": { + "cycles": { + "max": 3, + "min": 2 + }, + "instructions": { + "max": 198, + "min": 192 + }, + "launches": 1, + "loads": 24, + "registers": 4, + "stores": 6, + "threads": 6 + }, + "mlp_gpu": { + "cycles": { + "max": 15, + "min": 14 + }, + "instructions": { + "max": 3720, + "min": 3696 + }, + "launches": 2, + "loads": 552, + "registers": 4, + "stores": 36, + "threads": 36 + }, + "mm": { + "cycles": { + "max": 4, + "min": 4 + }, + "instructions": { + "max": 896, + "min": 896 + }, + "launches": 1, + "loads": 128, + "registers": 4, + "stores": 16, + "threads": 16 + }, + "mm_bias_relu": { + "cycles": { + "max": 9, + "min": 8 + }, + "instructions": { + "max": 1824, + "min": 1808 + }, + "launches": 1, + "loads": 272, + "registers": 4, + "stores": 16, + "threads": 16 + }, + "mm_relu": { + "cycles": { + "max": 9, + "min": 8 + }, + "instructions": { + "max": 1776, + "min": 1760 + }, + "launches": 1, + "loads": 256, + "registers": 4, + "stores": 16, + "threads": 16 + } + } + }, "tpu": { "O0": { "big_mm": { diff --git a/tests/test_cli.py b/tests/test_cli.py index 028ba55..8070825 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -120,3 +120,57 @@ def test_backend_flags_need_a_target(flags): 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 + + +# --- tiny-gpu backend --------------------------------------------------------------------- + + +def test_emit_tgpu_prints_assembly_with_source_positions(capsys): + code, out, _ = run(capsys, EX / "mm_relu.tiny", "--target", "gpu", "--emit", "tgpu") + assert code == 0 and out.startswith("; kernel k0_matmul") and "BRLT L0" in out + assert "; 4:5" in out + + +def test_emit_kernel_and_launch(capsys): + import json + + code, out, _ = run(capsys, EX / "mlp_gpu.tiny", "--target", "gpu", "--emit", "kernel") + assert code == 0 and out.count("kernel @") == 2 and "for v1 in 0..8" in out + code, out, _ = run(capsys, EX / "mlp_gpu.tiny", "--target", "gpu", "--emit", "launch") + doc = json.loads(out) + assert code == 0 and doc["format"] == "tgpu-launch" and len(doc["kernels"]) == 2 + + +def test_gpu_stats_prints_the_estimate_bounds(capsys): + import json + + code, out, _ = run(capsys, EX / "mm_relu.tiny", "--target", "gpu", "--stats") + est = json.loads(out) + assert code == 0 and est["launches"] == 1 and est["cycles"] == {"min": 8, "max": 9} + + +def test_gpu_run_executes_on_the_simulator_and_reports_measured_vs_estimate(capsys): + code, out, _ = run(capsys, EX / "mlp_gpu.tiny", "--target", "gpu", "--run") + assert code == 0 and "bit-exact vs numpy oracle: yes" in out + assert "R5-R7 untouched: yes" in out and "measured instructions:" in out + assert "straight-line code in one cycle" in out + + +def test_the_same_program_runs_on_both_chips_from_the_same_seed(capsys): + gpu = run(capsys, EX / "mm_relu.tiny", "--target", "gpu", "--run", "--seed", "5")[1] + tpu = run(capsys, EX / "mm_relu.tiny", "--target", "tpu", "--run", "--seed", "5")[1] + matrix = lambda out: out.split("bit-exact")[0] + assert matrix(gpu) == matrix(tpu) + + +def test_gpu_regs_budget_and_its_error(capsys): + assert run(capsys, EX / "mm_relu.tiny", "--target", "gpu", "--run", "--gpu-regs", "8")[0] == 0 + code, _, err = run(capsys, EX / "mm_relu.tiny", "--target", "gpu", "--run", "--gpu-regs", "7") + assert code == 1 and "at least 8 registers" in err + + +def test_gpu_emit_flags_need_a_gpu_target(capsys): + for flags in (["--emit", "tgpu"], ["--emit", "kernel", "--target", "tpu"]): + with pytest.raises(SystemExit) as exc: + main([str(EX / "mm.tiny"), *flags]) + assert exc.value.code == 2 diff --git a/tests/test_targets_contract.py b/tests/test_targets_contract.py index aa0c73b..f91708e 100644 --- a/tests/test_targets_contract.py +++ b/tests/test_targets_contract.py @@ -95,3 +95,51 @@ def test_submodules_are_pinned_and_checked_out(): names = {line.split()[1] for line in status} assert names == {"targets/tiny-gpu", "targets/tiny-tpu"} assert not any(line.startswith(("-", "+", "U")) for line in status), status + + +# --- tiny-gpu facts the GPU backend's code generation depends on (measured, PLAN §2.1) ------------ + + +def _run(src, **kw): + from runtime.gpu_runner import run_tgpu + + kw.setdefault("grid", (1, 3)) + kw.setdefault("mem_size", 16) + return run_tgpu(src, **kw) + + +def test_gpu_a_label_starting_with_r_is_parsed_as_a_register(): + # The emitter names labels L0, L1, ...; `relu_end` would crash tiny-gpu's assembler. + with pytest.raises(ValueError): + _run("SET R0, 1\nBRZ relu_end\nrelu_end:\nST 0, R0\n") + _run("SET R0, 1\nBRZ L0\nL0:\nST 0, R0\n") + + +def test_gpu_a_loop_of_k_iterations_costs_k_cycles(): + loop = "SET R0, 0\nL0:\nADD R0, R0, 1\nCMP R0, {k}\nBRLT L0\nST R7, R0\n" + for k in (1, 2, 5): + assert _run(loop.format(k=k)).cycles == k + + +def test_gpu_relu_idiom_costs_a_cycle_only_when_a_thread_does_not_clamp(): + relu = "LD R0, R7\nCMP R0, 0\nBRGT L0\nSET R0, 0\nL0:\nADD R2, R7, 8\nST R2, R0\n" + clamped = _run(relu, memory_init={0: [-1, -2, -3]}) + assert (clamped.cycles, clamped.op_counts["SET"]) == (1, 3) # every thread runs the SET + kept = _run(relu, memory_init={0: [1, 2, 3]}) + assert (kept.cycles, kept.op_counts["SET"]) == (2, 0) # every thread takes the skip branch + assert _run(relu, memory_init={0: [-1, 2, -3]}).read(8, 3) == [0, 2, 0] + + +def test_gpu_branching_to_the_end_of_the_program_costs_a_stray_cycle(): + src = "SET R0, 1\nSET R1, 0\nCMP R0, R1\nBRGT L9\nSET R2, 5\nL9:\n" + assert _run(src, grid=(1, 1)).cycles == 2 # why a kernel must not end with an `if` + + +def test_gpu_addresses_may_be_immediates_and_comments_are_allowed(): + res = _run("LD R0, 1 ; load\nADD R0, R0, 5\nST 2, R0\n", grid=(1, 1), memory_init={1: [10]}) + assert res.read(0, 4) == [0, 10, 15, 0] + + +def test_gpu_reserved_registers_are_reported_intact_after_a_normal_run(): + assert _run("ADD R0, R7, 1\nST R7, R0\n", grid=(2, 3)).reserved_intact + assert not _run("SET R7, 99\n", grid=(2, 3)).reserved_intact # the check can fail diff --git a/tinyc/cli.py b/tinyc/cli.py index 6c2ca56..741b7fa 100644 --- a/tinyc/cli.py +++ b/tinyc/cli.py @@ -1,4 +1,4 @@ -"""`tinyc FILE.tiny`: compile through the pass pipeline; print IR, a tile schedule or a program.""" +"""`tinyc FILE.tiny`: compile through the pass pipeline; print IR, a schedule, kernels or a program.""" from __future__ import annotations @@ -9,26 +9,31 @@ import numpy as np -from tinyc.driver import build_pipeline, compile_tpu +from tinyc.driver import build_pipeline, compile_gpu, 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.gpu.emit import dumps as dumps_launch +from tinyc.gpu.kernel_ir import print_kernel 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.passes import REGISTRY, PassContext, PassManager, PassVerifyError from tinyc.tpu.emit import dumps from tinyc.tpu.tile_ir import print_schedule +TPU_EMITS = ("tile", "ttpu") +GPU_EMITS = ("kernel", "tgpu", "launch") + 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", "tile", "ttpu"], default="ir", - help="tile and ttpu need --target tpu") # fmt: skip + ap.add_argument("--emit", choices=["ast", "ir", *TPU_EMITS, *GPU_EMITS], default="ir", + help="tile/ttpu need --target tpu; kernel/tgpu/launch need --target gpu") # 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("--target", choices=["gpu", "tpu"], help="also run that target's backend") ap.add_argument( "-O", dest="opt", type=int, choices=[0, 1], default=1, help="optimization level" ) @@ -37,29 +42,44 @@ def main(argv: list[str] | None = None) -> int: 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 + help="run the program after every pass/stage 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 + help="run on the real chip/simulator with seeded random inputs") # 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") + ap.add_argument("--gpu-regs", type=int, default=16, help="tiny-gpu registers per thread") 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)" - ) + wants_backend = args.emit in (*TPU_EMITS, *GPU_EMITS) or args.run or args.stats + if wants_backend and args.target is None: + ap.error("--emit tile/ttpu/kernel/tgpu/launch, --run and --stats need --target gpu|tpu") + if args.emit in TPU_EMITS and args.target != "tpu": + ap.error(f"--emit {args.emit} needs --target tpu") + if args.emit in GPU_EMITS and args.target != "gpu": + ap.error(f"--emit {args.emit} needs --target gpu") 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="") + ctx = PassContext(gpu_num_registers=args.gpu_regs) try: program = parse(source) if args.emit == "ast": print(dump(program.stmts)) return 0 + if wants_backend: + compile_fn = compile_tpu if args.target == "tpu" else compile_gpu + kw = {"ctx": ctx} if args.target == "gpu" else {} + comp = compile_fn( + source.text, source.name, opt_level=args.opt, disabled=args.disable, + verify_each=args.verify_each, dump=dump_ir, **kw, + ) # fmt: skip + if args.time_passes: + print(comp.report.timing_table(), file=sys.stderr) + return (_tpu_output if args.target == "tpu" else _gpu_output)(args, comp) func = lower(program, source) verify(func) manager = PassManager( @@ -67,14 +87,7 @@ def dump_ir(name: str, text: str) -> None: verify_each=args.verify_each, 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) + report = manager.run(func, ctx) except CompileError as err: print(err.render(), file=sys.stderr) return 1 @@ -96,18 +109,16 @@ def dump_ir(name: str, text: str) -> None: return 0 -def _tpu_output(args, comp, report) -> int: +def _tpu_output(args, comp) -> 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): + exact = all((got.outputs[n] == w).all() for n, w in zip(comp.schedule.outputs, want)) + for name in comp.schedule.outputs: 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}" @@ -123,5 +134,39 @@ def _tpu_output(args, comp, report) -> int: return 0 +def _gpu_output(args, comp) -> int: + prog = comp.program + est = prog.estimate + if args.run: + from runtime.gpu_runner import run_launch + + inputs = ref.random_inputs(comp.func, np.random.default_rng(args.seed)) + got = run_launch(comp.launch, inputs) + want = ref.evaluate(comp.func, inputs) + names = comp.launch["outputs"] + exact = all((got.outputs[n] == w).all() for n, w in zip(names, want)) + for name in names: + print(f"%{name} =\n{got.outputs[name]}") + print(f"bit-exact vs numpy oracle: {'yes' if exact else 'NO'}") + print(f"launches: {est['launches']} threads: {est['threads']} registers: {est['registers']}" + f" (R5-R7 untouched: {'yes' if got.reserved_intact else 'NO'})") # fmt: skip + c, i, o = est["cycles"], est["instructions"], got.op_counts + print(f"measured cycles: {got.cycles} (estimate {c['min']}..{c['max']})") + print(f"measured instructions: {got.instructions} (estimate {i['min']}..{i['max']})") + print(f"measured loads: {o['LD']} (estimate {est['loads']}) " + f"stores: {o['ST']} (estimate {est['stores']})") # fmt: skip + print("note: tiny-gpu runs straight-line code in one cycle, so compare instructions too") + return 0 if exact and got.reserved_intact else 1 + if args.stats: + print(json.dumps(est, indent=2)) + elif args.emit == "kernel": + print("\n".join(print_kernel(k) for k in prog.kernels), end="") + elif args.emit == "tgpu": + print(prog.asm_text(), end="") + else: + print(dumps_launch(comp.launch), end="") + return 0 + + if __name__ == "__main__": raise SystemExit(main()) diff --git a/tinyc/driver.py b/tinyc/driver.py index 473ed33..2c2bf03 100644 --- a/tinyc/driver.py +++ b/tinyc/driver.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from tinyc.frontend import compile_to_graph +from tinyc.gpu import GpuProgram, compile_gpu_program from tinyc.ir.graph import Func from tinyc.passes import REGISTRY, Fixpoint, PassContext, PassManager, PassReport from tinyc.tpu import compile_tile_ir @@ -98,3 +99,34 @@ def compile_tpu( ) # fmt: skip schedule = compile_tile_ir(func, opt_level, verify_each=verify_each, dump=dump) return TpuCompilation(func, schedule, to_json(schedule), report) + + +@dataclass +class GpuCompilation: + func: Func # the optimized, gpu-legalized Graph IR + program: GpuProgram # kernels, memory plan, assembly + report: PassReport # graph pass records + + @property + def launch(self) -> dict: + return self.program.launch + + +def compile_gpu( + 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, +) -> GpuCompilation: + """Source text -> kernels and `.tgpu` assembly for tiny-gpu (graph passes, then the backend).""" + ctx = ctx or PassContext() + func, report = compile_graph( + text, filename, target="gpu", opt_level=opt_level, disabled=disabled, + verify_each=verify_each, ctx=ctx, dump=dump, + ) # fmt: skip + program = compile_gpu_program(func, ctx, verify_each=verify_each, dump=dump) + return GpuCompilation(func, program, report) diff --git a/tinyc/gpu/__init__.py b/tinyc/gpu/__init__.py new file mode 100644 index 0000000..79721ad --- /dev/null +++ b/tinyc/gpu/__init__.py @@ -0,0 +1,115 @@ +"""tiny-gpu backend: Graph IR -> memory plan -> Kernel IR -> registers -> `.tgpu` + launch JSON.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field + +import numpy as np + +from tinyc.gpu import cost, interp +from tinyc.gpu.emit import emit_asm, to_launch_json +from tinyc.gpu.kernel_ir import Kernel, KernelError, flatten, print_kernel, verify_kernel +from tinyc.gpu.lower import lower +from tinyc.gpu.memplan import MemPlan, check_plan, plan_memory +from tinyc.gpu.regalloc import RegisterPressure, allocate +from tinyc.ir import ref +from tinyc.ir.errors import PassError, RangeError +from tinyc.ir.graph import Func +from tinyc.passes.manager import PassContext, PassVerifyError + + +@dataclass +class GpuProgram: + func: Func + plan: MemPlan + kernels: list[Kernel] + num_registers: int + consts: dict[str, np.ndarray] + asm: dict[str, str] = field(default_factory=dict) # kernel name -> .tgpu text + registers: dict[str, int] = field(default_factory=dict) # kernel name -> registers used + + @property + def estimate(self) -> dict: + regs = {k.name: self.registers.get(k.name, 0) for k in self.kernels} + return cost.estimate(self.kernels, regs) + + @property + def launch(self) -> dict: + return to_launch_json(self) + + def asm_text(self) -> str: + return "\n".join(self.asm[k.name] for k in self.kernels) + + +def _kernel_text(kernels: list[Kernel]) -> str: + return "\n".join(print_kernel(k) for k in kernels) + + +def _problem(want, got_outputs, names) -> str | None: + for name, w in zip(names, want): + if got_outputs[name].shape != w.shape or not (got_outputs[name] == w).all(): + return f"output '{name}' differs from the oracle on seeded random inputs" + return None + + +def compile_gpu_program( + func: Func, + ctx: PassContext | None = None, + *, + verify_each: bool = False, + dump: Callable[[str, str], None] | None = None, + seed: int = 42, + num_checks: int = 3, +) -> GpuProgram: + """Lower a gpu-legalized Graph IR function. Kernel code is the straightforward (-O0) form.""" + ctx = ctx or PassContext() + expectations = ref.seeded_expectations(func, num_checks, seed) if verify_each else [] + plan = plan_memory(func) + check_plan(func, plan) + if plan.mem_size > ctx.gpu_mem_words: + raise PassError( + f"program needs {plan.mem_size} words of tiny-gpu memory, budget is {ctx.gpu_mem_words}" + ) + consts = {op.result.name: op.data for op in func.ops if op.kind == "const"} + prog = GpuProgram(func, plan, lower(func, plan), ctx.gpu_num_registers, consts) + outputs = [v.name for v in func.outputs] + + text = _kernel_text(prog.kernels) + if dump: + dump("lower-to-kernel", text) + for inputs, want in expectations: # the structured kernels against the oracle + got = interp.run_program(prog, inputs)["outputs"] + if why := _problem(want, got, outputs): + raise PassVerifyError("lower-to-kernel", why, "", text) + + for k in prog.kernels: + try: + verify_kernel(k) + insts = flatten(k) + alloc = allocate(insts, prog.num_registers, k.name) + except RegisterPressure as e: + raise PassError(e.message, e.loc) from None + except KernelError as e: # pragma: no cover - lowering never builds an invalid kernel + raise PassVerifyError("lower-to-kernel", str(e), "", text) from None + prog.asm[k.name] = emit_asm(k, insts, alloc) + prog.registers[k.name] = alloc.registers_used + asm = prog.asm_text() + if dump: + dump("regalloc-emit", asm) + + if verify_each: # the emitted assembly on the real simulator against the oracle + from runtime.gpu_runner import run_launch + + for inputs, want in expectations: + try: + res = run_launch(prog.launch, inputs) + except (RangeError, OverflowError) as e: + raise PassVerifyError( + "regalloc-emit", f"simulator failed: {e}", text, asm + ) from None + if not res.reserved_intact: + raise PassVerifyError("regalloc-emit", "a kernel clobbered R5-R7", text, asm) + if why := _problem(want, res.outputs, outputs): + raise PassVerifyError("regalloc-emit", why, text, asm) + return prog diff --git a/tinyc/gpu/cost.py b/tinyc/gpu/cost.py new file mode 100644 index 0000000..c43e41b --- /dev/null +++ b/tinyc/gpu/cost.py @@ -0,0 +1,82 @@ +"""Static cost bounds for kernels, to be compared with what the simulator measures. + +Loads, stores and launch counts are exact. Instruction and cycle counts are data dependent +exactly where a kernel has an `if` (relu): a thread that clamps runs one more instruction, and a +thread that does not clamp takes a skip branch, which ends its simulator cycle. The bounds are +the two extremes; `gpu.interp` gives the exact numbers for a given input. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from tinyc.gpu.kernel_ir import Assign, For, Kernel, Load, Store + + +@dataclass +class _Cost: + imin: int = 0 + imax: int = 0 + tmin: int = 0 # taken branches per thread + tmax: int = 0 + loads: int = 0 + stores: int = 0 + + def seq(self, o: _Cost) -> _Cost: + return _Cost(self.imin + o.imin, self.imax + o.imax, self.tmin + o.tmin, + self.tmax + o.tmax, self.loads + o.loads, self.stores + o.stores) # fmt: skip + + +def _cost(body: list) -> _Cost: + total = _Cost() + for s in body: + if isinstance(s, (Assign, Load, Store)): + total = total.seq(_Cost(1, 1, 0, 0, isinstance(s, Load), isinstance(s, Store))) + elif isinstance(s, For): + n, b = s.stop - s.start, _cost(s.body) + total = total.seq(_Cost( + 1 + n * (b.imin + 3), 1 + n * (b.imax + 3), + (n - 1) + n * b.tmin, (n - 1) + n * b.tmax, n * b.loads, n * b.stores, + )) # fmt: skip + else: + b = _cost(s.body) + assert not ( + b.loads or b.stores + ), "memory ops inside `if` would make loads data dependent" + # skipped: 2 instructions + 1 taken branch; executed: 2 + body, body's own branches + total = total.seq(_Cost(2, 2 + b.imax, min(1, b.tmin), max(1, b.tmax), 0, 0)) + return total + + +def kernel_cost(k: Kernel, registers: int) -> dict: + c, t = _cost(k.body), k.threads + return { + "name": k.name, + "grid": list(k.grid), + "threads": t, + "instructions": {"min": t * c.imin, "max": t * c.imax}, + "cycles": {"min": c.tmin + 1, "max": c.tmax + 1}, + "loads": t * c.loads, + "stores": t * c.stores, + "registers": registers, + } + + +def estimate(kernels: list[Kernel], registers: dict[str, int]) -> dict: + per = [kernel_cost(k, registers[k.name]) for k in kernels] + return { + "launches": len(per), + "threads": sum(p["threads"] for p in per), + "instructions": { + "min": sum(p["instructions"]["min"] for p in per), + "max": sum(p["instructions"]["max"] for p in per), + }, + "cycles": { + "min": sum(p["cycles"]["min"] for p in per), + "max": sum(p["cycles"]["max"] for p in per), + }, + "loads": sum(p["loads"] for p in per), + "stores": sum(p["stores"] for p in per), + "registers": max((p["registers"] for p in per), default=0), + "per_kernel": per, + } diff --git a/tinyc/gpu/emit.py b/tinyc/gpu/emit.py new file mode 100644 index 0000000..9884bbe --- /dev/null +++ b/tinyc/gpu/emit.py @@ -0,0 +1,80 @@ +"""Kernel -> `.tgpu` assembly text, and the whole program -> launch JSON (docs/TGPU_FORMAT.md).""" + +from __future__ import annotations + +import json + +from tinyc.gpu.kernel_ir import Kernel, VInst +from tinyc.gpu.regalloc import Allocation, reg_name + +FORMAT = "tgpu-launch" +VERSION = 1 + + +def _inst_text(i: VInst, asg) -> str: + r = lambda x: reg_name(x, asg) + if i.op == "LABEL": + return f"{i.label}:" + if i.op in ("BRLT", "BRGT"): + return f"{i.op} {i.label}" + if i.op == "ST": + return f"ST {r(i.srcs[0])}, {r(i.srcs[1])}" + if i.op == "CMP": + return f"CMP {r(i.srcs[0])}, {r(i.srcs[1])}" + return f"{i.op} {r(i.dst)}, " + ", ".join(r(s) for s in i.srcs) + + +def emit_asm(kernel: Kernel, insts: list[VInst], alloc: Allocation) -> str: + """tiny-gpu assembly. Comments (`;`) carry the source line:col each instruction came from.""" + where = f"{kernel.loc.file}:{kernel.loc.line}:{kernel.loc.col}" if kernel.loc else "?" + lines = [ + f"; kernel {kernel.name}: %{kernel.value} = {kernel.kind} ({where})", + ( + f"; grid: {kernel.grid[0]} blocks x {kernel.grid[1]} threads;" + f" {alloc.registers_used} registers (R5-R7 reserved, read-only)" + ), + ] + texts = [_inst_text(i, alloc.assignment) for i in insts] + width = max((len(t) for t, i in zip(texts, insts) if i.op != "LABEL"), default=0) + 2 + for text, i in zip(texts, insts): + if i.op == "LABEL" or i.loc is None: + lines.append(text) + else: + lines.append(f"{text.ljust(width)}; {i.loc.line}:{i.loc.col}") + return "\n".join(lines) + "\n" + + +def to_launch_json(prog) -> dict: + """`prog` is a tinyc.gpu.GpuProgram.""" + buffers = {} + for name, r in prog.plan.buffers.items(): + entry = {"addr": r.addr, "dtype": r.dtype, "shape": list(r.shape), "kind": r.kind} + if r.kind == "const": + entry["data"] = prog.consts[name].tolist() + buffers[name] = entry + return { + "format": FORMAT, + "version": VERSION, + "name": prog.func.name, + "num_registers": prog.num_registers, + "mem_size": prog.plan.mem_size, + "buffers": buffers, + "inputs": [a.name for a in prog.func.args], + "outputs": [v.name for v in prog.func.outputs], + "kernels": [ + { + "name": k.name, + "value": k.value, + "op": k.kind, + "grid": list(k.grid), + "shared_size": 0, + "asm": prog.asm.get(k.name, ""), # empty before assembly (validation of the IR) + } + for k in prog.kernels + ], + "estimate": prog.estimate, + } + + +def dumps(launch: dict) -> str: + return json.dumps(launch, indent=2) + "\n" diff --git a/tinyc/gpu/interp.py b/tinyc/gpu/interp.py new file mode 100644 index 0000000..2961d28 --- /dev/null +++ b/tinyc/gpu/interp.py @@ -0,0 +1,144 @@ +"""Reference interpreter for the Kernel IR, with the same dynamic counts the simulator reports. + +It runs the *structured* kernels (before register allocation and assembly), thread by thread, +so it is independent of the lowering to tiny-gpu instructions and of the allocator. It also counts +what tiny-gpu would execute, so a test can require interpreter counts == simulator counts exactly: +that checks lowering, allocation and emission together. + +Counting follows the idioms in `kernel_ir.flatten` and tiny-gpu's `step()`: a thread's cycle ends +at a taken branch, so a thread needs (taken branches + 1) cycles and a launch needs the maximum +over its threads. +""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping +from dataclasses import dataclass, field + +import numpy as np + +from tinyc.gpu.kernel_ir import I32_HI, I32_LO, SPECIAL_REGS, Assign, For, Load, Special, Store +from tinyc.ir.errors import InputError, RangeError + + +class GpuFault(Exception): + """A thread touched memory outside [0, mem_size).""" + + +@dataclass +class KernelCounts: + name: str + threads: int + cycles: int # max over threads of (taken branches + 1) + ops: Counter = field(default_factory=Counter) # tiny-gpu op name -> executions, all threads + + @property + def instructions(self) -> int: + return sum(self.ops.values()) + + +def _run_thread(kernel, memory, env, ops: Counter) -> int: + """Execute one thread; returns its taken-branch count.""" + taken = 0 + + def val(x): + if isinstance(x, Special): + return env[x.name] + return env[x] if not isinstance(x, int) else x + + def fit(v: int, what: str) -> int: + if not I32_LO <= v <= I32_HI: + raise RangeError(f"{what} overflows i32 in kernel @{kernel.name}") + return v + + def check_addr(a: int) -> int: + if not 0 <= a < len(memory): + raise GpuFault(f"kernel @{kernel.name}: address {a} outside memory [0, {len(memory)})") + return a + + def walk(body): + nonlocal taken + for s in body: + if isinstance(s, Assign): + ops[s.op.upper()] += 1 + a = [val(x) for x in s.args] + env[s.dst] = fit( + a[0] if s.op == "set" else a[0] + a[1] if s.op == "add" else a[0] * a[1], s.op + ) + elif isinstance(s, Load): + ops["LD"] += 1 + env[s.dst] = memory[check_addr(val(s.addr))] + elif isinstance(s, Store): + ops["ST"] += 1 + memory[check_addr(val(s.addr))] = fit(val(s.src), "st") + elif isinstance(s, For): + ops["SET"] += 1 + env[s.var] = s.start + for it in range(s.stop - s.start): + walk(s.body) + env[s.var] = env[s.var] + 1 + ops["ADD"] += 1 + ops["CMP"] += 1 + ops["BRLT"] += 1 + if it < s.stop - s.start - 1: + taken += 1 # the back edge is taken on every iteration but the last + else: + ops["CMP"] += 1 + ops["BRGT"] += 1 + if val(s.a) <= val(s.b): + walk(s.body) + else: + taken += 1 # skip branch taken + + walk(kernel.body) + return taken + + +def initial_memory(launch: Mapping, inputs: Mapping[str, object]) -> list[int]: + """Memory image before the first kernel: constants and inputs placed, the rest zero.""" + memory = [0] * launch["mem_size"] + for name, b in launch["buffers"].items(): + if b["kind"] == "const": + flat = np.asarray(b["data"], dtype=np.int64).reshape(-1).tolist() + elif b["kind"] == "input": + if name not in inputs: + raise InputError(f"missing input '{name}'") + arr = np.asarray(inputs[name], dtype=np.int64) + if list(arr.shape) != b["shape"]: + raise InputError(f"input '{name}' must have shape {b['shape']}") + flat = arr.reshape(-1).tolist() + else: + continue + memory[b["addr"] : b["addr"] + len(flat)] = flat + return memory + + +def read_outputs(launch: Mapping, memory) -> dict[str, np.ndarray]: + out = {} + for name in launch["outputs"]: + b = launch["buffers"][name] + size = int(np.prod(b["shape"])) + out[name] = np.array(memory[b["addr"] : b["addr"] + size], dtype=np.int64).reshape( + b["shape"] + ) + return out + + +def run_program(prog, inputs: Mapping[str, object]) -> dict: + """Run a tinyc.gpu.GpuProgram's kernels. Returns {"outputs": {...}, "kernels": [KernelCounts]}.""" + memory = initial_memory(prog.launch, inputs) + results = [] + for k in prog.kernels: + ops: Counter = Counter() + worst = 0 + blocks, tpb = k.grid + for block in range(blocks): + for t in range(tpb): + env = {"block": block, "thread": t, "tid": block * tpb + t} + worst = max(worst, _run_thread(k, memory, env, ops)) + results.append(KernelCounts(k.name, k.threads, worst + 1, ops)) + return {"outputs": read_outputs(prog.launch, memory), "kernels": results} + + +assert set(SPECIAL_REGS) == {"block", "thread", "tid"} diff --git a/tinyc/gpu/kernel_ir.py b/tinyc/gpu/kernel_ir.py new file mode 100644 index 0000000..3821e86 --- /dev/null +++ b/tinyc/gpu/kernel_ir.py @@ -0,0 +1,263 @@ +"""Kernel IR: one tiny-gpu kernel as a structured program over virtual registers. + +Every thread runs the same body. The launch grid is (blocks, threads per block); the simulator +gives each thread three read-only registers, exposed here as `Special` operands: + block = R5 (block id) thread = R6 (thread in block) tid = R7 (global thread id) +Virtual registers (`VReg`) are mapped to R0-R4 and R8.. by the register allocator; R5-R7 are +never written. tiny-gpu has no SUB/DIV/shift/MAX, so the IR has only what the ISA can do: +`set`/`add`/`mul`, global `ld`/`st`, a counted `for` (do-while, trip count >= 1) and `if le`. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from tinyc.ir.loc import Loc +from tinyc.ir.types import RANGES + +SPECIAL_REGS = {"block": 5, "thread": 6, "tid": 7} +RESERVED_REGS = tuple(sorted(SPECIAL_REGS.values())) +I32_LO, I32_HI = RANGES["i32"] +MAX_THREADS = 1024 + + +@dataclass(frozen=True) +class VReg: + id: int + + def __str__(self) -> str: + return f"v{self.id}" + + +@dataclass(frozen=True) +class Special: + name: str # "block" | "thread" | "tid" + + def __str__(self) -> str: + return self.name + + +Operand = VReg | Special | int + + +@dataclass +class Assign: + op: str # "set" | "add" | "mul" + dst: VReg + args: tuple[Operand, ...] + loc: Loc | None = None + + +@dataclass +class Load: + dst: VReg + addr: Operand + loc: Loc | None = None + + +@dataclass +class Store: + addr: Operand + src: Operand + loc: Loc | None = None + + +@dataclass +class For: + """`var` counts start, start+1, ... while < stop. Always runs at least once (stop > start).""" + + var: VReg + start: int + stop: int + body: list = field(default_factory=list) + loc: Loc | None = None + + +@dataclass +class If: + """Run `body` when `a <= b` (the only condition tiny-gpu's CMP + BRGT gives us for free).""" + + a: Operand + b: Operand + body: list = field(default_factory=list) + loc: Loc | None = None + + +Stmt = Assign | Load | Store | For | If + + +@dataclass +class Kernel: + name: str + grid: tuple[int, int] # (blocks, threads per block) + body: list[Stmt] + value: str # name of the Graph IR value this kernel computes + kind: str # the Graph op: "matmul" | "add" | "relu" + loc: Loc | None = None + + @property + def threads(self) -> int: + return self.grid[0] * self.grid[1] + + +class KernelBuilder: + def __init__(self): + self._next = 0 + + def fresh(self) -> VReg: + v = VReg(self._next) + self._next += 1 + return v + + +class KernelError(Exception): + pass + + +# --- verifier --------------------------------------------------------------------------------- + + +def verify_kernel(k: Kernel) -> None: + problems: list[str] = [] + blocks, tpb = k.grid + if blocks < 1 or tpb < 1: + problems.append(f"grid {k.grid} must be positive") + if blocks * tpb > MAX_THREADS: + problems.append(f"{blocks * tpb} threads exceed the {MAX_THREADS}-thread launch limit") + + def operand(x: Operand, defined: set[VReg], where: str) -> None: + if isinstance(x, VReg): + if x not in defined: + problems.append(f"{where}: {x} used before it is assigned") + elif isinstance(x, Special): + if x.name not in SPECIAL_REGS: + problems.append(f"{where}: unknown special register '{x.name}'") + elif isinstance(x, int) and not isinstance(x, bool): + if not I32_LO <= x <= I32_HI: + problems.append(f"{where}: immediate {x} does not fit in i32") + else: + problems.append(f"{where}: bad operand {x!r}") + + def walk(body: list, defined: set[VReg]) -> None: + for s in body: + where = type(s).__name__.lower() + if isinstance(s, Assign): + arity = {"set": 1, "add": 2, "mul": 2}.get(s.op) + if arity is None or len(s.args) != arity: + problems.append(f"{where}: '{s.op}' with {len(s.args)} operands") + for a in s.args: + operand(a, defined, s.op) + defined.add(s.dst) + elif isinstance(s, Load): + operand(s.addr, defined, "ld") + defined.add(s.dst) + elif isinstance(s, Store): + operand(s.addr, defined, "st") + operand(s.src, defined, "st") + elif isinstance(s, For): + if s.stop <= s.start: + problems.append(f"for {s.var}: needs stop > start (got {s.start}..{s.stop})") + defined.add(s.var) + walk(s.body, defined) # runs at least once, so its assignments are definite + elif isinstance(s, If): + operand(s.a, defined, "if") + operand(s.b, defined, "if") + walk(s.body, set(defined)) # conditional: its assignments are not definite after + else: + problems.append(f"unknown statement {s!r}") + + walk(k.body, set()) + if k.body and isinstance(k.body[-1], If): + # tiny-gpu: branching to the end of the program costs a stray extra cycle (measured) + problems.append("a kernel must not end with an `if` (its skip label would be at the end)") + if problems: + raise KernelError(f"invalid kernel @{k.name}:\n " + "\n ".join(problems)) + + +# --- printing --------------------------------------------------------------------------------- + + +def _ops(args) -> str: + return ", ".join(str(a) for a in args) + + +def _print_body(body: list, indent: int, out: list[str]) -> None: + pad = " " * indent + for s in body: + if isinstance(s, Assign): + out.append(f"{pad}{s.dst} = {s.op} {_ops(s.args)}") + elif isinstance(s, Load): + out.append(f"{pad}{s.dst} = ld [{s.addr}]") + elif isinstance(s, Store): + out.append(f"{pad}st [{s.addr}], {s.src}") + elif isinstance(s, For): + out.append(f"{pad}for {s.var} in {s.start}..{s.stop} {{") + _print_body(s.body, indent + 1, out) + out.append(f"{pad}}}") + else: + out.append(f"{pad}if le {s.a}, {s.b} {{") + _print_body(s.body, indent + 1, out) + out.append(f"{pad}}}") + + +def print_kernel(k: Kernel) -> str: + where = f" at {k.loc.line}:{k.loc.col}" if k.loc else "" + head = ( + f"kernel @{k.name} grid=({k.grid[0]} blocks, {k.grid[1]} threads)" + f" // %{k.value} = {k.kind}{where}" + ) + out = [head + " {"] + _print_body(k.body, 1, out) + out.append("}") + return "\n".join(out) + "\n" + + +# --- flattening to virtual tiny-gpu instructions ----------------------------------------------- + + +@dataclass +class VInst: + """One tiny-gpu instruction over virtual registers (or a LABEL pseudo-instruction).""" + + op: str # SET ADD MUL LD ST CMP BRLT BRGT LABEL + dst: VReg | None = None + srcs: tuple[Operand, ...] = () + label: str | None = None # LABEL: the name defined; BR*: the target + loc: Loc | None = None + + +def flatten(k: Kernel) -> list[VInst]: + """Lower structured control flow to labels and branches, in tiny-gpu's own idiom: + + for: SET var, start / L: body / ADD var, var, 1 / CMP var, stop / BRLT L + if : CMP a, b / BRGT skip / body / skip: + """ + out: list[VInst] = [] + labels = iter(range(10**6)) + # Labels must not start with R/r: tiny-gpu's assembler would parse them as registers. + + def walk(body: list) -> None: + for s in body: + if isinstance(s, Assign): + out.append(VInst(s.op.upper(), s.dst, s.args, loc=s.loc)) + elif isinstance(s, Load): + out.append(VInst("LD", s.dst, (s.addr,), loc=s.loc)) + elif isinstance(s, Store): + out.append(VInst("ST", None, (s.addr, s.src), loc=s.loc)) + elif isinstance(s, For): + top = f"L{next(labels)}" + out.append(VInst("SET", s.var, (s.start,), loc=s.loc)) + out.append(VInst("LABEL", label=top, loc=s.loc)) + walk(s.body) + out.append(VInst("ADD", s.var, (s.var, 1), loc=s.loc)) + out.append(VInst("CMP", None, (s.var, s.stop), loc=s.loc)) + out.append(VInst("BRLT", label=top, loc=s.loc)) + else: + skip = f"L{next(labels)}" + out.append(VInst("CMP", None, (s.a, s.b), loc=s.loc)) + out.append(VInst("BRGT", label=skip, loc=s.loc)) + walk(s.body) + out.append(VInst("LABEL", label=skip, loc=s.loc)) + + walk(k.body) + return out diff --git a/tinyc/gpu/lower.py b/tinyc/gpu/lower.py new file mode 100644 index 0000000..7e7213d --- /dev/null +++ b/tinyc/gpu/lower.py @@ -0,0 +1,123 @@ +"""Graph IR (legalized for tiny-gpu) -> Kernel IR, one kernel per op (fused ops stay one kernel). + +Mapping: one thread per output element. For a rank-2 result [M, N] the grid is M blocks of N +threads, so block = row, thread = column, and the global thread id is the flat row-major index. +Rank-1 results use one block. Addresses are word addresses from the memory plan, baked in as +immediates (tiny-gpu has no base+offset addressing, so each access computes its address). + +This is the straightforward (-O0) code: loops are not unrolled and addresses are recomputed on +every iteration. +""" + +from __future__ import annotations + +from tinyc.gpu.kernel_ir import ( + Assign, + For, + If, + Kernel, + KernelBuilder, + Load, + Special, + Store, + verify_kernel, +) +from tinyc.gpu.memplan import MemPlan +from tinyc.ir.graph import Func, Op +from tinyc.ir.types import epilogue + +BLOCK, THREAD, TID = Special("block"), Special("thread"), Special("tid") + + +def _grid(shape: tuple[int, ...]) -> tuple[int, int]: + return (shape[0], shape[1]) if len(shape) == 2 else (1, shape[0]) + + +def _index(operand_shape, result_shape): + """Which special register indexes an elementwise operand: the flat id, or the column.""" + return TID if operand_shape == result_shape else THREAD + + +def _lower_matmul(op: Op, plan: MemPlan, name: str) -> Kernel: + b = KernelBuilder() + a_val, b_val = op.operands[0], op.operands[1] + stages = epilogue(op.attrs) + (_, k_dim), n_dim = a_val.type.shape, b_val.type.shape[1] + a_base, b_base, c_base = plan.addr(a_val.name), plan.addr(b_val.name), plan.addr(op.result.name) + loc = op.loc + acc, k = b.fresh(), b.fresh() + body: list = [Assign("set", acc, (0,), loc)] + + t_a, x, t_b, y, p = (b.fresh() for _ in range(5)) + loop = [ + Assign("mul", t_a, (BLOCK, k_dim), loc), # &A[row][k] = a_base + row*K + k + Assign("add", t_a, (t_a, k), loc), + Assign("add", t_a, (t_a, a_base), loc), + Load(x, t_a, loc), + Assign("mul", t_b, (k, n_dim), loc), # &B[k][col] = b_base + k*N + col + Assign("add", t_b, (t_b, THREAD), loc), + Assign("add", t_b, (t_b, b_base), loc), + Load(y, t_b, loc), + Assign("mul", p, (x, y), loc), + Assign("add", acc, (acc, p), loc), + ] + body.append(For(k, 0, k_dim, loop, loc)) + if "bias" in stages: + t_bias, z = b.fresh(), b.fresh() + bias_base = plan.addr(op.operands[2].name) + body += [ + Assign("add", t_bias, (THREAD, bias_base), loc), + Load(z, t_bias, loc), + Assign("add", acc, (acc, z), loc), + ] + if "relu" in stages: + body.append(If(acc, 0, [Assign("set", acc, (0,), loc)], loc)) + t_c = b.fresh() + body += [Assign("add", t_c, (TID, c_base), loc), Store(t_c, acc, loc)] + return Kernel(name, _grid(op.result.type.shape), body, op.result.name, "matmul", loc) + + +def _lower_add(op: Op, plan: MemPlan, name: str) -> Kernel: + b = KernelBuilder() + shape = op.result.type.shape + t_a, x, t_b, y, s, t_c = (b.fresh() for _ in range(6)) + lhs, rhs = op.operands + loc = op.loc + body = [ + Assign("add", t_a, (_index(lhs.type.shape, shape), plan.addr(lhs.name)), loc), + Load(x, t_a, loc), + Assign("add", t_b, (_index(rhs.type.shape, shape), plan.addr(rhs.name)), loc), + Load(y, t_b, loc), + Assign("add", s, (x, y), loc), + Assign("add", t_c, (TID, plan.addr(op.result.name)), loc), + Store(t_c, s, loc), + ] + return Kernel(name, _grid(shape), body, op.result.name, "add", loc) + + +def _lower_relu(op: Op, plan: MemPlan, name: str) -> Kernel: + b = KernelBuilder() + t_a, x, t_c = (b.fresh() for _ in range(3)) + loc = op.loc + body = [ + Assign("add", t_a, (TID, plan.addr(op.operands[0].name)), loc), + Load(x, t_a, loc), + If(x, 0, [Assign("set", x, (0,), loc)], loc), # relu: clamp x <= 0 to 0 + Assign("add", t_c, (TID, plan.addr(op.result.name)), loc), + Store(t_c, x, loc), + ] + return Kernel(name, _grid(op.result.type.shape), body, op.result.name, "relu", loc) + + +_LOWER = {"matmul": _lower_matmul, "add": _lower_add, "relu": _lower_relu} + + +def lower(func: Func, plan: MemPlan) -> list[Kernel]: + kernels = [] + for op in func.ops: + if op.kind == "const": + continue # constants are the memory's initial contents, not kernels + kernel = _LOWER[op.kind](op, plan, f"k{len(kernels)}_{op.kind}") + verify_kernel(kernel) + kernels.append(kernel) + return kernels diff --git a/tinyc/gpu/memplan.py b/tinyc/gpu/memplan.py new file mode 100644 index 0000000..31a7ca4 --- /dev/null +++ b/tinyc/gpu/memplan.py @@ -0,0 +1,120 @@ +"""Assign tiny-gpu global-memory word addresses to every Graph IR value. + +Inputs and constants are pinned (constants are the memory's initial contents); outputs are pinned +until the end. Every other intermediate is allocated when its op runs and its words are reused +once its last reader has finished. `i8` values take one int32 word each. + +Reuse is only ever from ops *strictly before* the one allocating: an op's result never overlaps +the buffers that same op reads (its threads read and write concurrently). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from tinyc.ir.graph import Func + + +@dataclass +class Region: + name: str + addr: int + size: int + dtype: str + shape: tuple[int, ...] + kind: str # "input" | "const" | "temp" + + +@dataclass +class MemPlan: + buffers: dict[str, Region] + mem_size: int + + def addr(self, name: str) -> int: + return self.buffers[name].addr + + +class _FreeList: + """First-fit free list over [0, top) with coalescing.""" + + def __init__(self): + self.free: list[tuple[int, int]] = [] # (addr, size), sorted, non-adjacent + self.top = 0 + + def alloc(self, size: int) -> int: + for i, (a, n) in enumerate(self.free): + if n >= size: + if n == size: + del self.free[i] + else: + self.free[i] = (a + size, n - size) + return a + addr = self.top + self.top += size + return addr + + def release(self, addr: int, size: int) -> None: + self.free.append((addr, size)) + self.free.sort() + merged: list[tuple[int, int]] = [] + for a, n in self.free: + if merged and merged[-1][0] + merged[-1][1] == a: + merged[-1] = (merged[-1][0], merged[-1][1] + n) + else: + merged.append((a, n)) + self.free = merged + + +def plan_memory(func: Func) -> MemPlan: + last_use: dict[str, int] = {} + for i, op in enumerate(func.ops): + for v in op.operands: + last_use[v.name] = i + pinned_until_end = {v.name for v in func.outputs} + + fl = _FreeList() + buffers: dict[str, Region] = {} + + def place(v, kind: str) -> Region: + r = Region(v.name, fl.alloc(v.type.size), v.type.size, v.type.dtype, v.type.shape, kind) + buffers[v.name] = r + return r + + for a in func.args: + place(a, "input") + for op in func.ops: + if op.kind == "const": + place(op.result, "const") + + live: dict[str, Region] = {} + for i, op in enumerate(func.ops): + if op.kind == "const": + continue + for name, r in list(live.items()): # readers that finished before op i + if last_use.get(name, -1) < i and name not in pinned_until_end: + fl.release(r.addr, r.size) + del live[name] + live[op.result.name] = place(op.result, "temp") + return MemPlan(buffers, fl.top) + + +def check_plan(func: Func, plan: MemPlan) -> None: + """No two buffers that are live at the same time may overlap (used by tests and verify).""" + n = len(func.ops) + first = {a.name: -1 for a in func.args} + last = {a.name: n for a in func.args} + for i, op in enumerate(func.ops): + first[op.result.name] = -1 if op.kind == "const" else i + last[op.result.name] = n if op.kind == "const" else i + for v in op.operands: + last[v.name] = max(last[v.name], i) + for v in func.outputs: + last[v.name] = n + names = list(plan.buffers) + for x in range(len(names)): + for y in range(x + 1, len(names)): + a, b = plan.buffers[names[x]], plan.buffers[names[y]] + alive_together = first[a.name] <= last[b.name] and first[b.name] <= last[a.name] + overlap = a.addr < b.addr + b.size and b.addr < a.addr + a.size + if alive_together and overlap: + raise AssertionError(f"buffers {a.name} and {b.name} overlap while both live") diff --git a/tinyc/gpu/regalloc.py b/tinyc/gpu/regalloc.py new file mode 100644 index 0000000..7d495fe --- /dev/null +++ b/tinyc/gpu/regalloc.py @@ -0,0 +1,112 @@ +"""Linear-scan register allocation for a flattened kernel. + +Pool: every register except the reserved R5 (block id), R6 (thread in block) and R7 (global tid), +which the simulator sets and the kernel only reads. There is no spilling: running out of +registers is a compile error that says how many were needed. + +Liveness is a [first mention, last mention] interval per virtual register, extended around loops: +a value that is live when a loop starts and used inside it must survive until the loop's back +edge, because the next iteration reads it again. A register is reusable from the instruction at +which its previous owner is last read (an instruction reads its sources before writing). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from tinyc.gpu.kernel_ir import RESERVED_REGS, SPECIAL_REGS, Special, VInst, VReg + + +class RegisterPressure(Exception): + def __init__(self, message: str, loc=None): + super().__init__(message) + self.message, self.loc = message, loc + + +@dataclass +class Allocation: + assignment: dict[VReg, int] + registers_used: int # distinct physical registers + max_live: int # peak number of simultaneously live values + + +def allocatable(num_registers: int) -> list[int]: + if num_registers < 8: + raise ValueError("tiny-gpu needs at least 8 registers (R7 holds the global thread id)") + return [r for r in range(num_registers) if r not in RESERVED_REGS] + + +def live_intervals(insts: list[VInst]) -> dict[VReg, list[int]]: + spans: dict[VReg, list[int]] = {} + + def touch(v, i): + if isinstance(v, VReg): + span = spans.setdefault(v, [i, i]) + span[0], span[1] = min(span[0], i), max(span[1], i) + + for i, inst in enumerate(insts): + for s in inst.srcs: + touch(s, i) + touch(inst.dst, i) + + label_at = {inst.label: i for i, inst in enumerate(insts) if inst.op == "LABEL"} + loops = [ + (label_at[inst.label], i) + for i, inst in enumerate(insts) + if inst.op in ("BRLT", "BRGT") and label_at[inst.label] <= i # a backward branch + ] + changed = True + while changed: # nested loops: extending for the inner one can matter for the outer one + changed = False + for top, bottom in loops: + for span in spans.values(): + if span[0] < top <= span[1] < bottom: + span[1] = bottom + changed = True + return spans + + +def allocate(insts: list[VInst], num_registers: int, kernel: str = "?") -> Allocation: + pool = allocatable(num_registers) + spans = live_intervals(insts) + order = sorted(spans, key=lambda v: (spans[v][0], spans[v][1], v.id)) + free = list(pool) + active: list[VReg] = [] + assignment: dict[VReg, int] = {} + max_live = 0 + for v in order: + start = spans[v][0] + for u in [u for u in active if spans[u][1] <= start]: # last read at or before `start` + active.remove(u) + free.append(assignment[u]) + free.sort() + if not free: + at = next((i.loc for i in insts if i.dst == v), None) + raise RegisterPressure( + f"kernel @{kernel} needs more than {len(pool)} registers: {len(active) + 1} " + f"values are live at once (R5-R7 are reserved by tiny-gpu); " + f"raise --gpu-regs (now {num_registers})", + at, + ) + assignment[v] = free.pop(0) + active.append(v) + max_live = max(max_live, len(active)) + assert_reserved_untouched(insts, assignment) + return Allocation(assignment, len(set(assignment.values())), max_live) + + +def assert_reserved_untouched(insts: list[VInst], assignment: dict[VReg, int]) -> None: + """R5-R7 are only ever read: no virtual register may land on them, and no instruction writes one.""" + bad = [v for v, r in assignment.items() if r in RESERVED_REGS] + assert not bad, f"allocator assigned reserved registers to {bad}" + for inst in insts: + if inst.dst is not None: + assert assignment[inst.dst] not in RESERVED_REGS, f"{inst} writes a reserved register" + + +def reg_name(x, assignment: dict[VReg, int]) -> str: + if isinstance(x, VReg): + return f"R{assignment[x]}" + if isinstance(x, Special): + return f"R{SPECIAL_REGS[x.name]}" + return str(x) diff --git a/tinyc/ir/ref.py b/tinyc/ir/ref.py index b66f8be..65297cc 100644 --- a/tinyc/ir/ref.py +++ b/tinyc/ir/ref.py @@ -76,3 +76,16 @@ def random_inputs( lo, hi = RANGES["i8"] if a.type.dtype == "i8" else i32_range out[a.name] = rng.integers(lo, hi + 1, size=a.type.shape, dtype=np.int64) return out + + +def seeded_expectations(func: Func, n: int, seed: int) -> list[tuple[dict, list[np.ndarray]]]: + """(inputs, oracle outputs) for n seeded random inputs, skipping inputs that overflow i32 in + the source program itself. Used by every `verify_each` translation-validation step.""" + out = [] + for i in range(n): + inputs = random_inputs(func, np.random.default_rng(seed + i)) + try: + out.append((inputs, evaluate(func, inputs))) + except RangeError: + continue + return out diff --git a/tinyc/passes/legalize_gpu.py b/tinyc/passes/legalize_gpu.py index 0cb5457..a4f5b78 100644 --- a/tinyc/passes/legalize_gpu.py +++ b/tinyc/passes/legalize_gpu.py @@ -11,6 +11,10 @@ class LegalizeGpu: name = "legalize-gpu" def run(self, func, ctx): + if ctx.gpu_num_registers < 8: + raise PassError( + f"tiny-gpu needs at least 8 registers (R7 holds the thread id); got {ctx.gpu_num_registers}" + ) for op in func.ops: if op.kind != "const" and op.result.type.size > ctx.gpu_max_threads: raise PassError( diff --git a/tinyc/passes/manager.py b/tinyc/passes/manager.py index de27534..cba65a3 100644 --- a/tinyc/passes/manager.py +++ b/tinyc/passes/manager.py @@ -13,8 +13,6 @@ 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 @@ -30,6 +28,7 @@ class PassContext: 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 + gpu_num_registers: int = 16 # tiny-gpu registers per thread (R5-R7 reserved by the simulator) analysis: dict = field(default_factory=dict) # e.g. analysis["ranges"] from range_check @@ -95,14 +94,7 @@ def __init__( 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 + return ref.seeded_expectations(func, self.num_checks, self.seed) def _semantic_problem(self, func: Func, expectations) -> str | None: for inputs, want in expectations: diff --git a/tinyc/tpu/__init__.py b/tinyc/tpu/__init__.py index 82ac427..e785149 100644 --- a/tinyc/tpu/__init__.py +++ b/tinyc/tpu/__init__.py @@ -4,8 +4,6 @@ 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 @@ -17,17 +15,6 @@ 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: @@ -51,7 +38,7 @@ def compile_tile_ir( 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 [] + expectations = ref.seeded_expectations(func, num_checks, seed) if verify_each else [] schedule = lower(func) stages = [("lower-to-tile", lambda: None)] if opt_level >= 1: