From 54071a6b98fd013f792936e0e91805d87136e7e5 Mon Sep 17 00:00:00 2001 From: Dineth Hettiarachchi Date: Sun, 25 Oct 2026 12:31:37 +0530 Subject: [PATCH] Phase 5: tiny-gpu kernel optimizations (-O1/-O2) Add kernel passes on the Kernel IR: loop-invariant code motion, strength reduction of counter-driven address chains into bumped pointers, liveness based dead-code elimination, shared-memory tiling (cooperative copy into SHST, SYNCB, then SHLD), full unrolling of loops up to 16 iterations, and constant/copy folding. -O1 enables licm and strength reduction; -O2 adds tiling, unrolling and folding. Every pass is switchable by name. The Kernel IR gains SharedLoad, SharedStore and Barrier. The interpreter models tiny-gpu's barrier from its source (waiting threads re-execute SYNCB every step) and a test requires it to equal the simulator exactly on hand-written kernels whose threads arrive at different times. The cost model becomes an abstract interpretation: loads, stores, shared accesses and a tiled copy's guards are exact; only a branch on loaded data (relu) yields min/max bounds. Measured on the simulator (mm_relu 8x8): -O2 cuts instructions by 55% and global loads by 47% versus -O0, and cycles from 10 to 3. -O1/-O2 trade registers for instructions. The table is generated by scripts/opt_table.py into docs/ARCHITECTURE.md and a test fails if it drifts. Tests: directed tests per pass, the exact load formula (2MNK -> MK + MNK) on every matmul example, every level on the real simulator, kernels validated by the interpreter after each pass, and Hypothesis properties over generated programs. A pass-generated instruction without a source position was caught by the existing source-mapping test and fixed. --- CLAUDE.md | 4 +- CONTEXT.md | 9 +- README.md | 5 +- docs/ARCHITECTURE.md | 107 ++++ docs/PASSES.md | 29 +- docs/TGPU_FORMAT.md | 19 +- docs/launch.schema.json | 13 +- scripts/opt_table.py | 36 ++ tests/golden/kernel/O0/big_mm.kernel | 27 + tests/golden/kernel/O0/const_mm.kernel | 27 + tests/golden/kernel/O0/mlp_gpu.kernel | 55 ++ tests/golden/kernel/O0/mm.kernel | 17 + tests/golden/kernel/O0/mm_bias_relu.kernel | 37 ++ tests/golden/kernel/O0/mm_relu.kernel | 27 + tests/golden/kernel/O1/big_mm.kernel | 19 + tests/golden/kernel/O1/const_mm.kernel | 19 + tests/golden/kernel/O1/mlp_gpu.kernel | 39 ++ tests/golden/kernel/O1/mm.kernel | 16 + tests/golden/kernel/O1/mm_bias_relu.kernel | 22 + tests/golden/kernel/O1/mm_relu.kernel | 19 + tests/golden/kernel/O2/big_mm.kernel | 74 +++ tests/golden/kernel/O2/const_mm.kernel | 24 + tests/golden/kernel/O2/mlp_gpu.kernel | 110 ++++ tests/golden/kernel/O2/mm.kernel | 31 ++ tests/golden/kernel/O2/mm_bias_relu.kernel | 61 +++ tests/golden/kernel/O2/mm_relu.kernel | 58 ++ 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/O0/big_mm.tgpu | 31 ++ tests/golden/tgpu/O0/const_mm.tgpu | 31 ++ tests/golden/tgpu/O0/mlp_gpu.tgpu | 62 +++ tests/golden/tgpu/{ => O0}/mm.tgpu | 2 +- tests/golden/tgpu/O0/mm_bias_relu.tgpu | 41 ++ tests/golden/tgpu/O0/mm_relu.tgpu | 31 ++ tests/golden/tgpu/O1/big_mm.tgpu | 23 + tests/golden/tgpu/O1/const_mm.tgpu | 23 + tests/golden/tgpu/O1/mlp_gpu.tgpu | 46 ++ tests/golden/tgpu/O1/mm.tgpu | 19 + tests/golden/tgpu/O1/mm_bias_relu.tgpu | 26 + tests/golden/tgpu/O1/mm_relu.tgpu | 23 + tests/golden/tgpu/O2/big_mm.tgpu | 75 +++ tests/golden/tgpu/O2/const_mm.tgpu | 25 + tests/golden/tgpu/O2/mlp_gpu.tgpu | 112 ++++ tests/golden/tgpu/O2/mm.tgpu | 31 ++ tests/golden/tgpu/O2/mm_bias_relu.tgpu | 62 +++ tests/golden/tgpu/O2/mm_relu.tgpu | 59 ++ 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_bias_relu.tgpu | 27 - tests/golden/tgpu/mm_relu.tgpu | 24 - tests/gpu/test_gpu_emit.py | 17 +- tests/gpu/test_gpu_perf.py | 33 +- tests/gpu/test_gpu_properties.py | 23 +- tests/gpu/test_gpu_sim.py | 45 +- tests/gpu/test_kernel_ir.py | 10 +- tests/gpu/test_kernel_passes.py | 504 +++++++++++++++++ tests/gpu/test_opt_table.py | 37 ++ tests/gpu/test_regalloc.py | 23 +- tests/perf_baseline.json | 164 +++++- tests/test_cli.py | 27 +- tinyc/cli.py | 15 +- tinyc/driver.py | 10 +- tinyc/gpu/__init__.py | 69 ++- tinyc/gpu/cost.py | 179 ++++-- tinyc/gpu/emit.py | 15 +- tinyc/gpu/interp.py | 105 ++-- tinyc/gpu/kernel_ir.py | 79 ++- tinyc/gpu/kernel_passes.py | 606 +++++++++++++++++++++ tinyc/gpu/lower.py | 22 +- tinyc/gpu/report.py | 77 +++ 75 files changed, 3500 insertions(+), 474 deletions(-) create mode 100644 docs/ARCHITECTURE.md create mode 100755 scripts/opt_table.py create mode 100644 tests/golden/kernel/O0/big_mm.kernel create mode 100644 tests/golden/kernel/O0/const_mm.kernel create mode 100644 tests/golden/kernel/O0/mlp_gpu.kernel create mode 100644 tests/golden/kernel/O0/mm.kernel create mode 100644 tests/golden/kernel/O0/mm_bias_relu.kernel create mode 100644 tests/golden/kernel/O0/mm_relu.kernel create mode 100644 tests/golden/kernel/O1/big_mm.kernel create mode 100644 tests/golden/kernel/O1/const_mm.kernel create mode 100644 tests/golden/kernel/O1/mlp_gpu.kernel create mode 100644 tests/golden/kernel/O1/mm.kernel create mode 100644 tests/golden/kernel/O1/mm_bias_relu.kernel create mode 100644 tests/golden/kernel/O1/mm_relu.kernel create mode 100644 tests/golden/kernel/O2/big_mm.kernel create mode 100644 tests/golden/kernel/O2/const_mm.kernel create mode 100644 tests/golden/kernel/O2/mlp_gpu.kernel create mode 100644 tests/golden/kernel/O2/mm.kernel create mode 100644 tests/golden/kernel/O2/mm_bias_relu.kernel create mode 100644 tests/golden/kernel/O2/mm_relu.kernel delete mode 100644 tests/golden/kernel/big_mm.kernel delete mode 100644 tests/golden/kernel/const_mm.kernel delete mode 100644 tests/golden/kernel/mlp_gpu.kernel delete mode 100644 tests/golden/kernel/mm.kernel delete mode 100644 tests/golden/kernel/mm_bias_relu.kernel delete mode 100644 tests/golden/kernel/mm_relu.kernel create mode 100644 tests/golden/tgpu/O0/big_mm.tgpu create mode 100644 tests/golden/tgpu/O0/const_mm.tgpu create mode 100644 tests/golden/tgpu/O0/mlp_gpu.tgpu rename tests/golden/tgpu/{ => O0}/mm.tgpu (100%) create mode 100644 tests/golden/tgpu/O0/mm_bias_relu.tgpu create mode 100644 tests/golden/tgpu/O0/mm_relu.tgpu create mode 100644 tests/golden/tgpu/O1/big_mm.tgpu create mode 100644 tests/golden/tgpu/O1/const_mm.tgpu create mode 100644 tests/golden/tgpu/O1/mlp_gpu.tgpu create mode 100644 tests/golden/tgpu/O1/mm.tgpu create mode 100644 tests/golden/tgpu/O1/mm_bias_relu.tgpu create mode 100644 tests/golden/tgpu/O1/mm_relu.tgpu create mode 100644 tests/golden/tgpu/O2/big_mm.tgpu create mode 100644 tests/golden/tgpu/O2/const_mm.tgpu create mode 100644 tests/golden/tgpu/O2/mlp_gpu.tgpu create mode 100644 tests/golden/tgpu/O2/mm.tgpu create mode 100644 tests/golden/tgpu/O2/mm_bias_relu.tgpu create mode 100644 tests/golden/tgpu/O2/mm_relu.tgpu delete mode 100644 tests/golden/tgpu/big_mm.tgpu delete mode 100644 tests/golden/tgpu/const_mm.tgpu delete mode 100644 tests/golden/tgpu/mlp_gpu.tgpu delete mode 100644 tests/golden/tgpu/mm_bias_relu.tgpu delete mode 100644 tests/golden/tgpu/mm_relu.tgpu create mode 100644 tests/gpu/test_kernel_passes.py create mode 100644 tests/gpu/test_opt_table.py create mode 100644 tinyc/gpu/kernel_passes.py create mode 100644 tinyc/gpu/report.py diff --git a/CLAUDE.md b/CLAUDE.md index 5737de9..82eb67b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -19,7 +19,9 @@ 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) +tinyc examples/mm_relu.tiny --target gpu --run -O2 # compile + run on the simulator (--emit kernel|tgpu|launch, --gpu-regs N, + # --disable licm|strength-reduce|shared-tile|unroll|fold|kernel-dce) +python scripts/opt_table.py --write # regenerate the measured table in docs/ARCHITECTURE.md 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 949d769..0d7b967 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -21,8 +21,10 @@ faked; host-side work is always labelled as host work. - 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. +- Phase 5 (GPU kernel optimizations): DONE. 557 tests green. `-O0/-O1/-O2`: licm, strength-reduce, kernel-dce | + shared-tile, unroll, fold. + Measured table in `docs/ARCHITECTURE.md` (generated by `scripts/opt_table.py`, drift-tested). `-O2` has strictly fewer + global loads than `-O0` on every matmul example (exact formula tested). Cost model is now an abstract interpretation. +- Next: Phase 6 (differential fuzzing + hardening; nightly 5000-example job; publishable CLI). ## Fixed decisions (do not re-litigate; change only via docs/ARCHITECTURE.md) | Decision | Choice | @@ -39,6 +41,9 @@ faked; host-side work is always labelled as host work. | 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 | +| `-O` levels | -O0 naive; -O1 graph passes + kernel licm/strength-reduce; -O2 + shared-tile, unroll, fold (gpu only; tpu -O2 == -O1). Kernel passes are switchable by name (`--disable`) | +| Kernel soundness | Lowering assigns every temporary once; only accumulators/pointers/counters are reassigned. Passes may move or rewrite single-assignment values freely | +| GPU barrier | `SYNCB` only at kernel top level; waiting threads re-execute it each step (counts modelled exactly in `gpu.interp`); out-of-range shared access is a fault in the interpreter | | 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 | diff --git a/README.md b/README.md index 3806031..4db2ae4 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ 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 4.** The frontend, optimizing pass pipeline and **both backends** work. One program, +**Status: Phase 5.** The frontend, optimizing pass pipeline and **both backends** work. One program, two chips, bit-exact against numpy on each: ```bash @@ -12,7 +12,8 @@ tinyc examples/mm_relu.tiny --target tpu --run # tile passes on the real tiny- tinyc examples/mm_relu.tiny --target gpu --run # SIMT kernels on the tiny-gpu simulator ``` -GPU kernel optimizations are next. +`-O2` tiles operands into tiny-gpu's shared memory and unrolls the loops (see +[docs/ARCHITECTURE.md](docs/ARCHITECTURE.md) for the measured effect of every optimization). 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) and [docs/TGPU_FORMAT.md](docs/TGPU_FORMAT.md). diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md new file mode 100644 index 0000000..c725a95 --- /dev/null +++ b/docs/ARCHITECTURE.md @@ -0,0 +1,107 @@ +# Architecture + +How the compiler is put together, and the decisions that are not obvious from the code. Change a +fixed decision here first (see `CONTEXT.md`). Per-pass details are in [PASSES.md](PASSES.md); the +formats are in [LANGUAGE.md](LANGUAGE.md), [IR.md](IR.md), [TTPU_FORMAT.md](TTPU_FORMAT.md) and +[TGPU_FORMAT.md](TGPU_FORMAT.md). + +## Pipeline + +``` + .tiny -> AST -> Graph IR -> graph passes -+-> Tile IR -> tile passes -> .ttpu.json (tiny-tpu RTL) + | + +-> memory plan -> Kernel IR -> kernel passes + -> registers -> .tgpu + launch (tiny-gpu) +``` + +Every level has a printer, a verifier and an interpreter; every pass can be validated against the +numpy oracle (`--verify-each`); and both backends run on the real chip/simulator, never on a +re-implementation. + +## Optimization levels (`-O`) + +| Level | Graph IR | Tile IR (tpu) | Kernel IR (gpu) | +|---|---|---|---| +| `-O0` | shape-infer, range-check | none | naive kernels | +| `-O1` (default) | + simplify fixpoint, fuse | + skip-zero-passes | + licm, strength-reduce, kernel-dce | +| `-O2` | same as `-O1` | same as `-O1` | + shared-tile, unroll, fold | + +Individual passes can be switched off by name (`--disable licm`), which is how the table below +isolates each optimization. + +## Kernel optimizations, measured + +Real simulator runs (`tinyc/gpu/report.py`; regenerate with `python scripts/opt_table.py --write`, +and a test fails if this table drifts). "Cycles" are tiny-gpu cycles: a thread's cycle ends at a taken +branch, so straight-line code is one cycle however long it is; read them together with instructions. +`relu` makes the exact counts data dependent, so these use a fixed seed. + + +**`mm_relu` 4x4** (`C = relu(A @ B)`, int8, seed 42) + +| configuration | instructions | cycles | global loads | shared loads | registers | +|---|---:|---:|---:|---:|---:| +| -O0 (naive) | 1000 | 6 | 144 | 0 | 4 | +| + licm | 840 (-16%) | 5 | 128 (-11%) | 0 | 5 | +| + strength-reduce (= -O1) | 728 (-27%) | 5 | 128 (-11%) | 0 | 6 | +| -O2 with shared-memory tiling only | 728 (-27%) | 6 | 80 (-44%) | 64 | 5 | +| -O2 with unrolling + folding only | 504 (-50%) | 2 | 128 (-11%) | 0 | 7 | +| -O2 (all passes) | 504 (-50%) | 3 | 80 (-44%) | 64 | 5 | + +**`mm_relu` 8x8** (`C = relu(A @ B)`, int8, seed 42) + +| configuration | instructions | cycles | global loads | shared loads | registers | +|---|---:|---:|---:|---:|---:| +| -O0 (naive) | 7320 | 10 | 1088 | 0 | 4 | +| + licm | 6168 (-16%) | 9 | 1024 (-6%) | 0 | 5 | +| + strength-reduce (= -O1) | 5208 (-29%) | 9 | 1024 (-6%) | 0 | 6 | +| -O2 with shared-memory tiling only | 4952 (-32%) | 10 | 576 (-47%) | 512 | 5 | +| -O2 with unrolling + folding only | 3544 (-52%) | 2 | 1024 (-6%) | 0 | 7 | +| -O2 (all passes) | 3288 (-55%) | 3 | 576 (-47%) | 512 | 5 | + + + +What each pass does, and what it costs: + +- **licm** hoists address arithmetic that does not depend on the loop counter (`row*K + base`). + It never hoists loads, so global loads are unchanged. It *raises register pressure*: hoisted values + stay live across the loop, which is why `-O0` is the code that fits the smallest register file. +- **strength-reduce** turns `k*N + col + base` (a multiply and two adds per iteration) into a pointer + initialised once and bumped by `N`. Chains of one statement are left alone: an add costs the same as + the increment that would replace it. +- **shared-tile** exploits that every thread of a block reads the same row of A. The block's threads + copy it into shared memory cooperatively, meet at a barrier (`SYNCB`), and read it with `SHLD`. + Global loads per matmul drop from `2*M*N*K` to `M*K + M*N*K` (about -44% to -47% above). It is roughly + instruction-neutral: the copy and the barrier cost about what the cheaper `shld` (no address + arithmetic) saves, and in loop form it adds a cycle. Skipped for one-thread blocks. +- **unroll** removes the loop for trip counts up to 16, so the kernel has no taken back edges and its + cycle count collapses. **fold** then folds the constant addresses and copies that unrolling exposes. + +### Why -O2 is not simply better + +- Unrolling alone is the biggest instruction and cycle win but the most register-hungry (7 registers + above: every iteration's temporaries are live at once before the allocator packs them); tiling + *reduces* pressure back to 5 because the A address chain disappears. The default register file (13 + allocatable) has room, but `--gpu-regs 8` only fits `-O0`. +- It needs a barrier, and tiny-gpu re-executes `SYNCB` on every step a thread spends waiting, so the + instruction count includes those repeats (modelled exactly in `gpu/interp.py`). +- Tiling only pays when a block has at least two threads; with a reduction as short as `K = 2` the + copy and barrier are not amortised, so loads and cycles improve more than instructions. +- A kernel tiles `K` shared words per operand per block; the simulator allocates exactly that. + +## Decisions worth knowing + +| Decision | Why | +|---|---| +| Temporaries are assigned once; only accumulators, pointers and counters are reassigned | It makes every kernel pass sound without dataflow analysis: a single-assignment value can be moved or rewritten freely. | +| Kernel passes are validated against the oracle after each pass | The structured kernels are interpreted (independently of the allocator and assembler); the final assembly is run on the simulator. | +| The estimate is an abstract interpretation | Values that depend only on block/thread ids are exact, so loads, stores, shared accesses and the tiled copy's guards are exact; only a branch on loaded data (relu) gives `min`/`max`. | +| One fresh simulator per kernel launch, memory carried over | tiny-gpu's `set_grid` resize leaves R7 = 0 for threads it adds. | +| `-O` levels are cumulative, kernel passes are switchable by name | Reproducibility, and so the table above can attribute each gain. | + +## Known limits + +- Kernel optimization is specific to what `lower.py` emits (matmul-shaped loops); the passes are + written generally but only exercised on those kernels. +- No spilling: a program that needs more registers than the budget is a compile error. +- The dead final pointer bump of an unrolled loop is not removed (the DCE is not flow-sensitive). diff --git a/docs/PASSES.md b/docs/PASSES.md index c88e41d..99e3549 100644 --- a/docs/PASSES.md +++ b/docs/PASSES.md @@ -66,9 +66,26 @@ validates the schedule after lowering and after each tile pass, by interpreting ## 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. +Kernel IR (`gpu/lower.py`, one kernel per op) -> kernel passes (`gpu/kernel_passes.py`) -> verify -> +flatten -> linear-scan register allocation (`gpu/regalloc.py`) -> `.tgpu` emission (`gpu/emit.py`). + +| `-O` | Kernel passes | +|---|---| +| 0 | none (the naive kernels) | +| 1 | `licm`, `strength-reduce`, `kernel-dce` | +| 2 | `licm`, `strength-reduce`, `shared-tile`, `unroll`, `fold`, `kernel-dce` | + +- `licm`: hoist single-assignment arithmetic that does not change inside a loop (never loads). +- `strength-reduce`: replace a `mul`/`add` chain on the loop counter by a pointer bumped each iteration + (chains of two or more statements only; not when the value is read after the loop). +- `kernel-dce`: remove values that never reach a store, `if` or loop control (liveness-based, so a dead + value that only updates itself goes too). +- `shared-tile`: a loop load whose address is `u + k`, with `u` the same for the whole block, is copied + cooperatively into shared memory (`SHST`), a barrier (`SYNCB`) follows, and the loop reads `SHLD`. +- `unroll`: full unroll of innermost loops of at most 16 iterations (the counter becomes a constant). +- `fold`: constant propagation, constant folding, add-chain reassociation and copy propagation. + +With `--verify-each` the structured kernels are interpreted against the numpy oracle after lowering +and after *every* kernel pass, and the emitted assembly is run on the real simulator against the +oracle. `--print-after-all` shows `lower-to-kernel`, each kernel pass, and `regalloc-emit`. Each pass +can be disabled by name. The measured effect of each is in [ARCHITECTURE.md](ARCHITECTURE.md). diff --git a/docs/TGPU_FORMAT.md b/docs/TGPU_FORMAT.md index 576ac7b..127e411 100644 --- a/docs/TGPU_FORMAT.md +++ b/docs/TGPU_FORMAT.md @@ -27,14 +27,29 @@ Assembler traps found by measurement (and enforced by the emitter): - **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. +## Shared memory and barriers (`-O2`) +`SHST saddr, Rs` / `SHLD Rd, saddr` access the block's own shared memory (`shared_size` words per +block, declared per kernel in the launch document); `SYNCB` is a block-wide barrier. Measured and +modelled behaviour: +- a thread that reaches `SYNCB` stops for that step, and a **waiting thread executes `SYNCB` again on + every following step** until all threads of its block are waiting; then the block resumes on the next + step. With arrival step `s_t = taken branches before the barrier + 1` and release step `R = max s_t`, + a thread executes `SYNCB` `R - s_t + 1` times and finishes at step `R + taken branches after + 1`. + `gpu/interp.py` implements this and a test requires it to equal the simulator exactly, including for + hand-written kernels whose threads arrive at different times; +- out-of-range `SHLD` silently reads 0 and out-of-range `SHST` is silently dropped in tiny-gpu; the + interpreter treats both as faults, so a compiler bug cannot hide; +- the Kernel IR only allows a barrier at the top level of a kernel, so every thread reaches 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 +cycle. The static `estimate` is an abstract interpretation: anything that depends only on block/thread ids is +exact (loop trip counts, a tiled kernel's copy guards), so loads, stores and shared accesses are exact; +only a branch on loaded data gives `min`/`max` for instructions and cycles. `gpu.interp` gives exact counts for given inputs, and tests require them to equal the simulator's. ## Launch document diff --git a/docs/launch.schema.json b/docs/launch.schema.json index c62e0ef..93d7f94 100644 --- a/docs/launch.schema.json +++ b/docs/launch.schema.json @@ -51,30 +51,35 @@ }, "estimate": { "type": "object", - "required": ["launches", "threads", "instructions", "cycles", "loads", "stores", "registers", "per_kernel"], + "required": ["launches", "threads", "instructions", "cycles", "loads", "stores", "shared_loads", "shared_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 }, + "loads": { "type": "integer", "minimum": 0, "description": "global loads, exact" }, + "stores": { "type": "integer", "minimum": 0, "description": "global stores, exact" }, + "shared_loads": { "type": "integer", "minimum": 0 }, + "shared_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"], + "required": ["name", "grid", "threads", "shared_size", "instructions", "cycles", "loads", "stores", "shared_loads", "shared_stores", "registers"], "additionalProperties": false, "properties": { "name": { "type": "string" }, "grid": { "type": "array", "items": { "type": "integer" }, "minItems": 2, "maxItems": 2 }, "threads": { "type": "integer", "minimum": 1 }, + "shared_size": { "type": "integer", "minimum": 0 }, "instructions": { "$ref": "#/$defs/range" }, "cycles": { "$ref": "#/$defs/range" }, "loads": { "type": "integer", "minimum": 0 }, "stores": { "type": "integer", "minimum": 0 }, + "shared_loads": { "type": "integer", "minimum": 0 }, + "shared_stores": { "type": "integer", "minimum": 0 }, "registers": { "type": "integer", "minimum": 0 } } } diff --git a/scripts/opt_table.py b/scripts/opt_table.py new file mode 100755 index 0000000..8ca48d2 --- /dev/null +++ b/scripts/opt_table.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +"""Print the measured optimization table, or rewrite it inside docs/ARCHITECTURE.md. + +python scripts/opt_table.py # print +python scripts/opt_table.py --write # update docs/ARCHITECTURE.md between its markers +""" + +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +from tinyc.gpu.report import opt_table_markdown + +START, END = "", "" + + +def main() -> int: + table = opt_table_markdown() + if "--write" not in sys.argv: + print(table) + return 0 + doc = ROOT / "docs" / "ARCHITECTURE.md" + text = doc.read_text() + new = re.sub( + f"{re.escape(START)}.*?{re.escape(END)}", f"{START}\n{table}\n{END}", text, flags=re.DOTALL + ) + doc.write_text(new) + print(f"updated {doc.relative_to(ROOT)}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/golden/kernel/O0/big_mm.kernel b/tests/golden/kernel/O0/big_mm.kernel new file mode 100644 index 0000000..849f088 --- /dev/null +++ b/tests/golden/kernel/O0/big_mm.kernel @@ -0,0 +1,27 @@ +kernel @k0_matmul grid=(12 blocks, 12 threads) // %0 = matmul at 4:10 { + v0 = set 0 + for v1 in 0..12 { + v2 = mul block, 12 + v3 = add v2, 0 + v4 = add v3, v1 + v5 = ld [v4] + v6 = mul v1, 12 + v7 = add v6, thread + v8 = add v7, 144 + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + } + v11 = add tid, 288 + st [v11], v0 +} + +kernel @k1_relu grid=(12 blocks, 12 threads) // %1 = relu at 4:5 { + v0 = add tid, 288 + v1 = ld [v0] + if le v1, 0 { + v1 = set 0 + } + v2 = add tid, 432 + st [v2], v1 +} diff --git a/tests/golden/kernel/O0/const_mm.kernel b/tests/golden/kernel/O0/const_mm.kernel new file mode 100644 index 0000000..f759c8b --- /dev/null +++ b/tests/golden/kernel/O0/const_mm.kernel @@ -0,0 +1,27 @@ +kernel @k0_matmul grid=(3 blocks, 2 threads) // %1 = matmul at 5:10 { + v0 = set 0 + for v1 in 0..2 { + v2 = mul block, 2 + v3 = add v2, 0 + v4 = add v3, v1 + v5 = ld [v4] + v6 = mul v1, 2 + v7 = add v6, thread + v8 = add v7, 6 + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + } + v11 = add tid, 10 + st [v11], v0 +} + +kernel @k1_relu grid=(3 blocks, 2 threads) // %2 = relu at 5:5 { + v0 = add tid, 10 + v1 = ld [v0] + if le v1, 0 { + v1 = set 0 + } + v2 = add tid, 16 + st [v2], v1 +} diff --git a/tests/golden/kernel/O0/mlp_gpu.kernel b/tests/golden/kernel/O0/mlp_gpu.kernel new file mode 100644 index 0000000..e94dcb5 --- /dev/null +++ b/tests/golden/kernel/O0/mlp_gpu.kernel @@ -0,0 +1,55 @@ +kernel @k0_matmul grid=(4 blocks, 6 threads) // %0 = matmul at 6:10 { + v0 = set 0 + for v1 in 0..8 { + v2 = mul block, 8 + v3 = add v2, 0 + v4 = add v3, v1 + v5 = ld [v4] + v6 = mul v1, 6 + v7 = add v6, thread + v8 = add v7, 32 + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + } + v11 = add tid, 104 + st [v11], v0 +} + +kernel @k1_add grid=(4 blocks, 6 threads) // %1 = add at 6:10 { + v0 = add tid, 104 + v1 = ld [v0] + v2 = add thread, 80 + v3 = ld [v2] + v4 = add v1, v3 + v5 = add tid, 128 + st [v5], v4 +} + +kernel @k2_relu grid=(4 blocks, 6 threads) // %2 = relu at 6:5 { + v0 = add tid, 128 + v1 = ld [v0] + if le v1, 0 { + v1 = set 0 + } + v2 = add tid, 104 + st [v2], v1 +} + +kernel @k3_matmul grid=(4 blocks, 3 threads) // %3 = matmul at 7:5 { + v0 = set 0 + for v1 in 0..6 { + v2 = mul block, 6 + v3 = add v2, 104 + v4 = add v3, v1 + v5 = ld [v4] + v6 = mul v1, 3 + v7 = add v6, thread + v8 = add v7, 86 + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + } + v11 = add tid, 128 + st [v11], v0 +} diff --git a/tests/golden/kernel/O0/mm.kernel b/tests/golden/kernel/O0/mm.kernel new file mode 100644 index 0000000..6df5ceb --- /dev/null +++ b/tests/golden/kernel/O0/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 + v3 = add v2, 0 + v4 = add v3, v1 + v5 = ld [v4] + v6 = mul v1, 4 + v7 = add v6, thread + v8 = add v7, 16 + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + } + v11 = add tid, 32 + st [v11], v0 +} diff --git a/tests/golden/kernel/O0/mm_bias_relu.kernel b/tests/golden/kernel/O0/mm_bias_relu.kernel new file mode 100644 index 0000000..97e7451 --- /dev/null +++ b/tests/golden/kernel/O0/mm_bias_relu.kernel @@ -0,0 +1,37 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %0 = matmul at 5:10 { + v0 = set 0 + for v1 in 0..8 { + v2 = mul block, 8 + v3 = add v2, 0 + v4 = add v3, v1 + v5 = ld [v4] + v6 = mul v1, 4 + v7 = add v6, thread + v8 = add v7, 32 + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + } + v11 = add tid, 68 + st [v11], v0 +} + +kernel @k1_add grid=(4 blocks, 4 threads) // %1 = add at 5:10 { + v0 = add tid, 68 + v1 = ld [v0] + v2 = add thread, 64 + v3 = ld [v2] + v4 = add v1, v3 + v5 = add tid, 84 + st [v5], v4 +} + +kernel @k2_relu grid=(4 blocks, 4 threads) // %2 = relu at 5:5 { + v0 = add tid, 84 + v1 = ld [v0] + if le v1, 0 { + v1 = set 0 + } + v2 = add tid, 68 + st [v2], v1 +} diff --git a/tests/golden/kernel/O0/mm_relu.kernel b/tests/golden/kernel/O0/mm_relu.kernel new file mode 100644 index 0000000..424a46a --- /dev/null +++ b/tests/golden/kernel/O0/mm_relu.kernel @@ -0,0 +1,27 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %0 = matmul at 4:10 { + v0 = set 0 + for v1 in 0..8 { + v2 = mul block, 8 + v3 = add v2, 0 + v4 = add v3, v1 + v5 = ld [v4] + v6 = mul v1, 4 + v7 = add v6, thread + v8 = add v7, 32 + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + } + v11 = add tid, 64 + st [v11], v0 +} + +kernel @k1_relu grid=(4 blocks, 4 threads) // %1 = relu at 4:5 { + v0 = add tid, 64 + v1 = ld [v0] + if le v1, 0 { + v1 = set 0 + } + v2 = add tid, 80 + st [v2], v1 +} diff --git a/tests/golden/kernel/O1/big_mm.kernel b/tests/golden/kernel/O1/big_mm.kernel new file mode 100644 index 0000000..32a1fe0 --- /dev/null +++ b/tests/golden/kernel/O1/big_mm.kernel @@ -0,0 +1,19 @@ +kernel @k0_matmul grid=(12 blocks, 12 threads) // %1 = matmul at 4:5 { + v0 = set 0 + v2 = mul block, 12 + v3 = add v2, 0 + v8 = add thread, 144 + for v1 in 0..12 { + v4 = add v3, v1 + v5 = ld [v4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + } + if le v0, 0 { + v0 = set 0 + } + v11 = add tid, 288 + st [v11], v0 +} diff --git a/tests/golden/kernel/O1/const_mm.kernel b/tests/golden/kernel/O1/const_mm.kernel new file mode 100644 index 0000000..75dee48 --- /dev/null +++ b/tests/golden/kernel/O1/const_mm.kernel @@ -0,0 +1,19 @@ +kernel @k0_matmul grid=(3 blocks, 2 threads) // %2 = matmul at 5:5 { + v0 = set 0 + v2 = mul block, 2 + v3 = add v2, 0 + v8 = add thread, 6 + for v1 in 0..2 { + v4 = add v3, v1 + v5 = ld [v4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 2 + } + if le v0, 0 { + v0 = set 0 + } + v11 = add tid, 10 + st [v11], v0 +} diff --git a/tests/golden/kernel/O1/mlp_gpu.kernel b/tests/golden/kernel/O1/mlp_gpu.kernel new file mode 100644 index 0000000..b42784e --- /dev/null +++ b/tests/golden/kernel/O1/mlp_gpu.kernel @@ -0,0 +1,39 @@ +kernel @k0_matmul grid=(4 blocks, 6 threads) // %2 = matmul at 6:5 { + v0 = set 0 + v2 = mul block, 8 + v3 = add v2, 0 + v8 = add thread, 32 + for v1 in 0..8 { + v4 = add v3, v1 + v5 = ld [v4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + } + v11 = add thread, 80 + v12 = ld [v11] + v0 = add v0, v12 + if le v0, 0 { + v0 = set 0 + } + v13 = add tid, 104 + st [v13], v0 +} + +kernel @k1_matmul grid=(4 blocks, 3 threads) // %3 = matmul at 7:5 { + v0 = set 0 + v2 = mul block, 6 + v3 = add v2, 104 + v8 = add thread, 86 + for v1 in 0..6 { + v4 = add v3, v1 + v5 = ld [v4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 3 + } + v11 = add tid, 128 + st [v11], v0 +} diff --git a/tests/golden/kernel/O1/mm.kernel b/tests/golden/kernel/O1/mm.kernel new file mode 100644 index 0000000..cd7b740 --- /dev/null +++ b/tests/golden/kernel/O1/mm.kernel @@ -0,0 +1,16 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %0 = matmul at 4:5 { + v0 = set 0 + v2 = mul block, 4 + v3 = add v2, 0 + v8 = add thread, 16 + for v1 in 0..4 { + v4 = add v3, v1 + v5 = ld [v4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + } + v11 = add tid, 32 + st [v11], v0 +} diff --git a/tests/golden/kernel/O1/mm_bias_relu.kernel b/tests/golden/kernel/O1/mm_bias_relu.kernel new file mode 100644 index 0000000..50f0c46 --- /dev/null +++ b/tests/golden/kernel/O1/mm_bias_relu.kernel @@ -0,0 +1,22 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %2 = matmul at 5:5 { + v0 = set 0 + v2 = mul block, 8 + v3 = add v2, 0 + v8 = add thread, 32 + for v1 in 0..8 { + v4 = add v3, v1 + v5 = ld [v4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + } + v11 = add thread, 64 + v12 = ld [v11] + v0 = add v0, v12 + if le v0, 0 { + v0 = set 0 + } + v13 = add tid, 68 + st [v13], v0 +} diff --git a/tests/golden/kernel/O1/mm_relu.kernel b/tests/golden/kernel/O1/mm_relu.kernel new file mode 100644 index 0000000..ba7b93b --- /dev/null +++ b/tests/golden/kernel/O1/mm_relu.kernel @@ -0,0 +1,19 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) // %1 = matmul at 4:5 { + v0 = set 0 + v2 = mul block, 8 + v3 = add v2, 0 + v8 = add thread, 32 + for v1 in 0..8 { + v4 = add v3, v1 + v5 = ld [v4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + } + if le v0, 0 { + v0 = set 0 + } + v11 = add tid, 64 + st [v11], v0 +} diff --git a/tests/golden/kernel/O2/big_mm.kernel b/tests/golden/kernel/O2/big_mm.kernel new file mode 100644 index 0000000..bf4b1c6 --- /dev/null +++ b/tests/golden/kernel/O2/big_mm.kernel @@ -0,0 +1,74 @@ +kernel @k0_matmul grid=(12 blocks, 12 threads) shared=12 words // %1 = matmul at 4:5 { + v0 = set 0 + v2 = mul block, 12 + v8 = add thread, 144 + v12 = add v2, thread + v13 = ld [v12] + shst [thread], v13 + syncb + v5 = shld [0] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [1] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [2] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [3] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [5] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [6] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [7] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [8] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [9] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [10] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + v5 = shld [11] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 12 + if le v0, 0 { + v0 = set 0 + } + v11 = add tid, 288 + st [v11], v0 +} diff --git a/tests/golden/kernel/O2/const_mm.kernel b/tests/golden/kernel/O2/const_mm.kernel new file mode 100644 index 0000000..43672c6 --- /dev/null +++ b/tests/golden/kernel/O2/const_mm.kernel @@ -0,0 +1,24 @@ +kernel @k0_matmul grid=(3 blocks, 2 threads) shared=2 words // %2 = matmul at 5:5 { + v0 = set 0 + v2 = mul block, 2 + v8 = add thread, 6 + v12 = add v2, thread + v13 = ld [v12] + shst [thread], v13 + syncb + v5 = shld [0] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 2 + v5 = shld [1] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 2 + if le v0, 0 { + v0 = set 0 + } + v11 = add tid, 10 + st [v11], v0 +} diff --git a/tests/golden/kernel/O2/mlp_gpu.kernel b/tests/golden/kernel/O2/mlp_gpu.kernel new file mode 100644 index 0000000..a890bfb --- /dev/null +++ b/tests/golden/kernel/O2/mlp_gpu.kernel @@ -0,0 +1,110 @@ +kernel @k0_matmul grid=(4 blocks, 6 threads) shared=8 words // %2 = matmul at 6:5 { + v0 = set 0 + v2 = mul block, 8 + v8 = add thread, 32 + v14 = add v2, thread + v15 = ld [v14] + shst [thread], v15 + v16 = add thread, 6 + if le v16, 7 { + v17 = add v2, v16 + v18 = ld [v17] + shst [v16], v18 + } + syncb + v5 = shld [0] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v5 = shld [1] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v5 = shld [2] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v5 = shld [3] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v5 = shld [4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v5 = shld [5] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v5 = shld [6] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v5 = shld [7] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 6 + v11 = add thread, 80 + v12 = ld [v11] + v0 = add v0, v12 + if le v0, 0 { + v0 = set 0 + } + v13 = add tid, 104 + st [v13], v0 +} + +kernel @k1_matmul grid=(4 blocks, 3 threads) shared=6 words // %3 = matmul at 7:5 { + v0 = set 0 + v2 = mul block, 6 + v3 = add v2, 104 + v8 = add thread, 86 + v12 = add v3, thread + v13 = ld [v12] + shst [thread], v13 + v14 = add thread, 3 + v15 = add v3, v14 + v16 = ld [v15] + shst [v14], v16 + syncb + v5 = shld [0] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 3 + v5 = shld [1] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 3 + v5 = shld [2] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 3 + v5 = shld [3] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 3 + v5 = shld [4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 3 + v5 = shld [5] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 3 + v11 = add tid, 128 + st [v11], v0 +} diff --git a/tests/golden/kernel/O2/mm.kernel b/tests/golden/kernel/O2/mm.kernel new file mode 100644 index 0000000..d079cb8 --- /dev/null +++ b/tests/golden/kernel/O2/mm.kernel @@ -0,0 +1,31 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) shared=4 words // %0 = matmul at 4:5 { + v0 = set 0 + v2 = mul block, 4 + v8 = add thread, 16 + v12 = add v2, thread + v13 = ld [v12] + shst [thread], v13 + syncb + v5 = shld [0] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [1] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [2] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [3] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v11 = add tid, 32 + st [v11], v0 +} diff --git a/tests/golden/kernel/O2/mm_bias_relu.kernel b/tests/golden/kernel/O2/mm_bias_relu.kernel new file mode 100644 index 0000000..f321103 --- /dev/null +++ b/tests/golden/kernel/O2/mm_bias_relu.kernel @@ -0,0 +1,61 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) shared=8 words // %2 = matmul at 5:5 { + v0 = set 0 + v2 = mul block, 8 + v8 = add thread, 32 + v14 = add v2, thread + v15 = ld [v14] + shst [thread], v15 + v16 = add thread, 4 + v17 = add v2, v16 + v18 = ld [v17] + shst [v16], v18 + syncb + v5 = shld [0] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [1] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [2] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [3] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [5] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [6] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [7] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v11 = add thread, 64 + v12 = ld [v11] + v0 = add v0, v12 + if le v0, 0 { + v0 = set 0 + } + v13 = add tid, 68 + st [v13], v0 +} diff --git a/tests/golden/kernel/O2/mm_relu.kernel b/tests/golden/kernel/O2/mm_relu.kernel new file mode 100644 index 0000000..1ef42e1 --- /dev/null +++ b/tests/golden/kernel/O2/mm_relu.kernel @@ -0,0 +1,58 @@ +kernel @k0_matmul grid=(4 blocks, 4 threads) shared=8 words // %1 = matmul at 4:5 { + v0 = set 0 + v2 = mul block, 8 + v8 = add thread, 32 + v12 = add v2, thread + v13 = ld [v12] + shst [thread], v13 + v14 = add thread, 4 + v15 = add v2, v14 + v16 = ld [v15] + shst [v14], v16 + syncb + v5 = shld [0] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [1] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [2] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [3] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [4] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [5] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [6] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + v5 = shld [7] + v9 = ld [v8] + v10 = mul v5, v9 + v0 = add v0, v10 + v8 = add v8, 4 + if le v0, 0 { + v0 = set 0 + } + v11 = add tid, 64 + st [v11], v0 +} diff --git a/tests/golden/kernel/big_mm.kernel b/tests/golden/kernel/big_mm.kernel deleted file mode 100644 index d79a058..0000000 --- a/tests/golden/kernel/big_mm.kernel +++ /dev/null @@ -1,20 +0,0 @@ -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 deleted file mode 100644 index 05393e8..0000000 --- a/tests/golden/kernel/const_mm.kernel +++ /dev/null @@ -1,20 +0,0 @@ -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 deleted file mode 100644 index 4f9ee57..0000000 --- a/tests/golden/kernel/mlp_gpu.kernel +++ /dev/null @@ -1,41 +0,0 @@ -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 deleted file mode 100644 index d9d7c24..0000000 --- a/tests/golden/kernel/mm.kernel +++ /dev/null @@ -1,17 +0,0 @@ -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 deleted file mode 100644 index 4e615f9..0000000 --- a/tests/golden/kernel/mm_bias_relu.kernel +++ /dev/null @@ -1,23 +0,0 @@ -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 deleted file mode 100644 index 478be65..0000000 --- a/tests/golden/kernel/mm_relu.kernel +++ /dev/null @@ -1,20 +0,0 @@ -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/O0/big_mm.tgpu b/tests/golden/tgpu/O0/big_mm.tgpu new file mode 100644 index 0000000..7b4ab57 --- /dev/null +++ b/tests/golden/tgpu/O0/big_mm.tgpu @@ -0,0 +1,31 @@ +; kernel k0_matmul: %0 = matmul (t.tiny:4:10) +; grid: 12 blocks x 12 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:10 +SET R1, 0 ; 4:10 +L0: +MUL R2, R5, 12 ; 4:10 +ADD R2, R2, 0 ; 4:10 +ADD R2, R2, R1 ; 4:10 +LD R2, R2 ; 4:10 +MUL R3, R1, 12 ; 4:10 +ADD R3, R3, R6 ; 4:10 +ADD R3, R3, 144 ; 4:10 +LD R3, R3 ; 4:10 +MUL R2, R2, R3 ; 4:10 +ADD R0, R0, R2 ; 4:10 +ADD R1, R1, 1 ; 4:10 +CMP R1, 12 ; 4:10 +BRLT L0 ; 4:10 +ADD R1, R7, 288 ; 4:10 +ST R1, R0 ; 4:10 + +; kernel k1_relu: %1 = relu (t.tiny:4:5) +; grid: 12 blocks x 12 threads; 2 registers (R5-R7 reserved, read-only) +ADD R0, R7, 288 ; 4:5 +LD R0, R0 ; 4:5 +CMP R0, 0 ; 4:5 +BRGT L0 ; 4:5 +SET R0, 0 ; 4:5 +L0: +ADD R1, R7, 432 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/O0/const_mm.tgpu b/tests/golden/tgpu/O0/const_mm.tgpu new file mode 100644 index 0000000..024839b --- /dev/null +++ b/tests/golden/tgpu/O0/const_mm.tgpu @@ -0,0 +1,31 @@ +; kernel k0_matmul: %1 = matmul (t.tiny:5:10) +; grid: 3 blocks x 2 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:10 +SET R1, 0 ; 5:10 +L0: +MUL R2, R5, 2 ; 5:10 +ADD R2, R2, 0 ; 5:10 +ADD R2, R2, R1 ; 5:10 +LD R2, R2 ; 5:10 +MUL R3, R1, 2 ; 5:10 +ADD R3, R3, R6 ; 5:10 +ADD R3, R3, 6 ; 5:10 +LD R3, R3 ; 5:10 +MUL R2, R2, R3 ; 5:10 +ADD R0, R0, R2 ; 5:10 +ADD R1, R1, 1 ; 5:10 +CMP R1, 2 ; 5:10 +BRLT L0 ; 5:10 +ADD R1, R7, 10 ; 5:10 +ST R1, R0 ; 5:10 + +; kernel k1_relu: %2 = relu (t.tiny:5:5) +; grid: 3 blocks x 2 threads; 2 registers (R5-R7 reserved, read-only) +ADD R0, R7, 10 ; 5:5 +LD R0, R0 ; 5:5 +CMP R0, 0 ; 5:5 +BRGT L0 ; 5:5 +SET R0, 0 ; 5:5 +L0: +ADD R1, R7, 16 ; 5:5 +ST R1, R0 ; 5:5 diff --git a/tests/golden/tgpu/O0/mlp_gpu.tgpu b/tests/golden/tgpu/O0/mlp_gpu.tgpu new file mode 100644 index 0000000..1d69348 --- /dev/null +++ b/tests/golden/tgpu/O0/mlp_gpu.tgpu @@ -0,0 +1,62 @@ +; kernel k0_matmul: %0 = matmul (t.tiny:6:10) +; grid: 4 blocks x 6 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 6:10 +SET R1, 0 ; 6:10 +L0: +MUL R2, R5, 8 ; 6:10 +ADD R2, R2, 0 ; 6:10 +ADD R2, R2, R1 ; 6:10 +LD R2, R2 ; 6:10 +MUL R3, R1, 6 ; 6:10 +ADD R3, R3, R6 ; 6:10 +ADD R3, R3, 32 ; 6:10 +LD R3, R3 ; 6:10 +MUL R2, R2, R3 ; 6:10 +ADD R0, R0, R2 ; 6:10 +ADD R1, R1, 1 ; 6:10 +CMP R1, 8 ; 6:10 +BRLT L0 ; 6:10 +ADD R1, R7, 104 ; 6:10 +ST R1, R0 ; 6:10 + +; kernel k1_add: %1 = add (t.tiny:6:10) +; grid: 4 blocks x 6 threads; 2 registers (R5-R7 reserved, read-only) +ADD R0, R7, 104 ; 6:10 +LD R0, R0 ; 6:10 +ADD R1, R6, 80 ; 6:10 +LD R1, R1 ; 6:10 +ADD R0, R0, R1 ; 6:10 +ADD R1, R7, 128 ; 6:10 +ST R1, R0 ; 6:10 + +; kernel k2_relu: %2 = relu (t.tiny:6:5) +; grid: 4 blocks x 6 threads; 2 registers (R5-R7 reserved, read-only) +ADD R0, R7, 128 ; 6:5 +LD R0, R0 ; 6:5 +CMP R0, 0 ; 6:5 +BRGT L0 ; 6:5 +SET R0, 0 ; 6:5 +L0: +ADD R1, R7, 104 ; 6:5 +ST R1, R0 ; 6:5 + +; kernel k3_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, 104 ; 7:5 +ADD R2, R2, R1 ; 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/O0/mm.tgpu similarity index 100% rename from tests/golden/tgpu/mm.tgpu rename to tests/golden/tgpu/O0/mm.tgpu index f08787a..10d421d 100644 --- a/tests/golden/tgpu/mm.tgpu +++ b/tests/golden/tgpu/O0/mm.tgpu @@ -4,8 +4,8 @@ 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 +ADD R2, R2, R1 ; 4:5 LD R2, R2 ; 4:5 MUL R3, R1, 4 ; 4:5 ADD R3, R3, R6 ; 4:5 diff --git a/tests/golden/tgpu/O0/mm_bias_relu.tgpu b/tests/golden/tgpu/O0/mm_bias_relu.tgpu new file mode 100644 index 0000000..df7163f --- /dev/null +++ b/tests/golden/tgpu/O0/mm_bias_relu.tgpu @@ -0,0 +1,41 @@ +; kernel k0_matmul: %0 = matmul (t.tiny:5:10) +; grid: 4 blocks x 4 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:10 +SET R1, 0 ; 5:10 +L0: +MUL R2, R5, 8 ; 5:10 +ADD R2, R2, 0 ; 5:10 +ADD R2, R2, R1 ; 5:10 +LD R2, R2 ; 5:10 +MUL R3, R1, 4 ; 5:10 +ADD R3, R3, R6 ; 5:10 +ADD R3, R3, 32 ; 5:10 +LD R3, R3 ; 5:10 +MUL R2, R2, R3 ; 5:10 +ADD R0, R0, R2 ; 5:10 +ADD R1, R1, 1 ; 5:10 +CMP R1, 8 ; 5:10 +BRLT L0 ; 5:10 +ADD R1, R7, 68 ; 5:10 +ST R1, R0 ; 5:10 + +; kernel k1_add: %1 = add (t.tiny:5:10) +; grid: 4 blocks x 4 threads; 2 registers (R5-R7 reserved, read-only) +ADD R0, R7, 68 ; 5:10 +LD R0, R0 ; 5:10 +ADD R1, R6, 64 ; 5:10 +LD R1, R1 ; 5:10 +ADD R0, R0, R1 ; 5:10 +ADD R1, R7, 84 ; 5:10 +ST R1, R0 ; 5:10 + +; kernel k2_relu: %2 = relu (t.tiny:5:5) +; grid: 4 blocks x 4 threads; 2 registers (R5-R7 reserved, read-only) +ADD R0, R7, 84 ; 5:5 +LD R0, R0 ; 5:5 +CMP R0, 0 ; 5:5 +BRGT L0 ; 5:5 +SET R0, 0 ; 5:5 +L0: +ADD R1, R7, 68 ; 5:5 +ST R1, R0 ; 5:5 diff --git a/tests/golden/tgpu/O0/mm_relu.tgpu b/tests/golden/tgpu/O0/mm_relu.tgpu new file mode 100644 index 0000000..10a666a --- /dev/null +++ b/tests/golden/tgpu/O0/mm_relu.tgpu @@ -0,0 +1,31 @@ +; kernel k0_matmul: %0 = matmul (t.tiny:4:10) +; grid: 4 blocks x 4 threads; 4 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:10 +SET R1, 0 ; 4:10 +L0: +MUL R2, R5, 8 ; 4:10 +ADD R2, R2, 0 ; 4:10 +ADD R2, R2, R1 ; 4:10 +LD R2, R2 ; 4:10 +MUL R3, R1, 4 ; 4:10 +ADD R3, R3, R6 ; 4:10 +ADD R3, R3, 32 ; 4:10 +LD R3, R3 ; 4:10 +MUL R2, R2, R3 ; 4:10 +ADD R0, R0, R2 ; 4:10 +ADD R1, R1, 1 ; 4:10 +CMP R1, 8 ; 4:10 +BRLT L0 ; 4:10 +ADD R1, R7, 64 ; 4:10 +ST R1, R0 ; 4:10 + +; kernel k1_relu: %1 = relu (t.tiny:4:5) +; grid: 4 blocks x 4 threads; 2 registers (R5-R7 reserved, read-only) +ADD R0, R7, 64 ; 4:5 +LD R0, R0 ; 4:5 +CMP R0, 0 ; 4:5 +BRGT L0 ; 4:5 +SET R0, 0 ; 4:5 +L0: +ADD R1, R7, 80 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/O1/big_mm.tgpu b/tests/golden/tgpu/O1/big_mm.tgpu new file mode 100644 index 0000000..eaffa0b --- /dev/null +++ b/tests/golden/tgpu/O1/big_mm.tgpu @@ -0,0 +1,23 @@ +; kernel k0_matmul: %1 = matmul (t.tiny:4:5) +; grid: 12 blocks x 12 threads; 6 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +MUL R1, R5, 12 ; 4:5 +ADD R1, R1, 0 ; 4:5 +ADD R2, R6, 144 ; 4:5 +SET R3, 0 ; 4:5 +L0: +ADD R4, R1, R3 ; 4:5 +LD R4, R4 ; 4:5 +LD R8, R2 ; 4:5 +MUL R4, R4, R8 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +ADD R3, R3, 1 ; 4:5 +CMP R3, 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/O1/const_mm.tgpu b/tests/golden/tgpu/O1/const_mm.tgpu new file mode 100644 index 0000000..78137bd --- /dev/null +++ b/tests/golden/tgpu/O1/const_mm.tgpu @@ -0,0 +1,23 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:5:5) +; grid: 3 blocks x 2 threads; 6 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:5 +MUL R1, R5, 2 ; 5:5 +ADD R1, R1, 0 ; 5:5 +ADD R2, R6, 6 ; 5:5 +SET R3, 0 ; 5:5 +L0: +ADD R4, R1, R3 ; 5:5 +LD R4, R4 ; 5:5 +LD R8, R2 ; 5:5 +MUL R4, R4, R8 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 2 ; 5:5 +ADD R3, R3, 1 ; 5:5 +CMP R3, 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/O1/mlp_gpu.tgpu b/tests/golden/tgpu/O1/mlp_gpu.tgpu new file mode 100644 index 0000000..459e598 --- /dev/null +++ b/tests/golden/tgpu/O1/mlp_gpu.tgpu @@ -0,0 +1,46 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:6:5) +; grid: 4 blocks x 6 threads; 6 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 6:5 +MUL R1, R5, 8 ; 6:5 +ADD R1, R1, 0 ; 6:5 +ADD R2, R6, 32 ; 6:5 +SET R3, 0 ; 6:5 +L0: +ADD R4, R1, R3 ; 6:5 +LD R4, R4 ; 6:5 +LD R8, R2 ; 6:5 +MUL R4, R4, R8 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +ADD R3, R3, 1 ; 6:5 +CMP R3, 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; 6 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 7:5 +MUL R1, R5, 6 ; 7:5 +ADD R1, R1, 104 ; 7:5 +ADD R2, R6, 86 ; 7:5 +SET R3, 0 ; 7:5 +L0: +ADD R4, R1, R3 ; 7:5 +LD R4, R4 ; 7:5 +LD R8, R2 ; 7:5 +MUL R4, R4, R8 ; 7:5 +ADD R0, R0, R4 ; 7:5 +ADD R2, R2, 3 ; 7:5 +ADD R3, R3, 1 ; 7:5 +CMP R3, 6 ; 7:5 +BRLT L0 ; 7:5 +ADD R1, R7, 128 ; 7:5 +ST R1, R0 ; 7:5 diff --git a/tests/golden/tgpu/O1/mm.tgpu b/tests/golden/tgpu/O1/mm.tgpu new file mode 100644 index 0000000..e5c7f01 --- /dev/null +++ b/tests/golden/tgpu/O1/mm.tgpu @@ -0,0 +1,19 @@ +; kernel k0_matmul: %0 = matmul (t.tiny:4:5) +; grid: 4 blocks x 4 threads; 6 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +MUL R1, R5, 4 ; 4:5 +ADD R1, R1, 0 ; 4:5 +ADD R2, R6, 16 ; 4:5 +SET R3, 0 ; 4:5 +L0: +ADD R4, R1, R3 ; 4:5 +LD R4, R4 ; 4:5 +LD R8, R2 ; 4:5 +MUL R4, R4, R8 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +ADD R3, R3, 1 ; 4:5 +CMP R3, 4 ; 4:5 +BRLT L0 ; 4:5 +ADD R1, R7, 32 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/O1/mm_bias_relu.tgpu b/tests/golden/tgpu/O1/mm_bias_relu.tgpu new file mode 100644 index 0000000..782a6cd --- /dev/null +++ b/tests/golden/tgpu/O1/mm_bias_relu.tgpu @@ -0,0 +1,26 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:5:5) +; grid: 4 blocks x 4 threads; 6 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:5 +MUL R1, R5, 8 ; 5:5 +ADD R1, R1, 0 ; 5:5 +ADD R2, R6, 32 ; 5:5 +SET R3, 0 ; 5:5 +L0: +ADD R4, R1, R3 ; 5:5 +LD R4, R4 ; 5:5 +LD R8, R2 ; 5:5 +MUL R4, R4, R8 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +ADD R3, R3, 1 ; 5:5 +CMP R3, 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/O1/mm_relu.tgpu b/tests/golden/tgpu/O1/mm_relu.tgpu new file mode 100644 index 0000000..e2a28d2 --- /dev/null +++ b/tests/golden/tgpu/O1/mm_relu.tgpu @@ -0,0 +1,23 @@ +; kernel k0_matmul: %1 = matmul (t.tiny:4:5) +; grid: 4 blocks x 4 threads; 6 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +MUL R1, R5, 8 ; 4:5 +ADD R1, R1, 0 ; 4:5 +ADD R2, R6, 32 ; 4:5 +SET R3, 0 ; 4:5 +L0: +ADD R4, R1, R3 ; 4:5 +LD R4, R4 ; 4:5 +LD R8, R2 ; 4:5 +MUL R4, R4, R8 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +ADD R3, R3, 1 ; 4:5 +CMP R3, 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/golden/tgpu/O2/big_mm.tgpu b/tests/golden/tgpu/O2/big_mm.tgpu new file mode 100644 index 0000000..efd5c40 --- /dev/null +++ b/tests/golden/tgpu/O2/big_mm.tgpu @@ -0,0 +1,75 @@ +; kernel k0_matmul: %1 = matmul (t.tiny:4:5) +; grid: 12 blocks x 12 threads; 12 shared words per block; 5 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +MUL R1, R5, 12 ; 4:5 +ADD R2, R6, 144 ; 4:5 +ADD R1, R1, R6 ; 4:5 +LD R1, R1 ; 4:5 +SHST R6, R1 ; 4:5 +SYNCB ; 4:5 +SHLD R1, 0 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 1 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 2 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 3 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 4 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 5 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 6 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 7 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 8 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 9 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 10 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +SHLD R1, 11 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 12 ; 4:5 +CMP R0, 0 ; 4:5 +BRGT L0 ; 4:5 +SET R0, 0 ; 4:5 +L0: +ADD R1, R7, 288 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/O2/const_mm.tgpu b/tests/golden/tgpu/O2/const_mm.tgpu new file mode 100644 index 0000000..3b12385 --- /dev/null +++ b/tests/golden/tgpu/O2/const_mm.tgpu @@ -0,0 +1,25 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:5:5) +; grid: 3 blocks x 2 threads; 2 shared words per block; 5 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:5 +MUL R1, R5, 2 ; 5:5 +ADD R2, R6, 6 ; 5:5 +ADD R1, R1, R6 ; 5:5 +LD R1, R1 ; 5:5 +SHST R6, R1 ; 5:5 +SYNCB ; 5:5 +SHLD R1, 0 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 2 ; 5:5 +SHLD R1, 1 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 2 ; 5:5 +CMP R0, 0 ; 5:5 +BRGT L0 ; 5:5 +SET R0, 0 ; 5:5 +L0: +ADD R1, R7, 10 ; 5:5 +ST R1, R0 ; 5:5 diff --git a/tests/golden/tgpu/O2/mlp_gpu.tgpu b/tests/golden/tgpu/O2/mlp_gpu.tgpu new file mode 100644 index 0000000..2471b33 --- /dev/null +++ b/tests/golden/tgpu/O2/mlp_gpu.tgpu @@ -0,0 +1,112 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:6:5) +; grid: 4 blocks x 6 threads; 8 shared words per block; 5 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 6:5 +MUL R1, R5, 8 ; 6:5 +ADD R2, R6, 32 ; 6:5 +ADD R3, R1, R6 ; 6:5 +LD R3, R3 ; 6:5 +SHST R6, R3 ; 6:5 +ADD R3, R6, 6 ; 6:5 +CMP R3, 7 ; 6:5 +BRGT L0 ; 6:5 +ADD R1, R1, R3 ; 6:5 +LD R1, R1 ; 6:5 +SHST R3, R1 ; 6:5 +L0: +SYNCB ; 6:5 +SHLD R1, 0 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +SHLD R1, 1 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +SHLD R1, 2 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +SHLD R1, 3 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +SHLD R1, 4 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +SHLD R1, 5 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +SHLD R1, 6 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 6:5 +SHLD R1, 7 ; 6:5 +LD R3, R2 ; 6:5 +MUL R4, R1, R3 ; 6:5 +ADD R0, R0, R4 ; 6:5 +ADD R2, R2, 6 ; 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; 6 shared words per block; 5 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 7:5 +MUL R1, R5, 6 ; 7:5 +ADD R1, R1, 104 ; 7:5 +ADD R2, R6, 86 ; 7:5 +ADD R3, R1, R6 ; 7:5 +LD R3, R3 ; 7:5 +SHST R6, R3 ; 7:5 +ADD R3, R6, 3 ; 7:5 +ADD R1, R1, R3 ; 7:5 +LD R1, R1 ; 7:5 +SHST R3, R1 ; 7:5 +SYNCB ; 7:5 +SHLD R1, 0 ; 7:5 +LD R3, R2 ; 7:5 +MUL R4, R1, R3 ; 7:5 +ADD R0, R0, R4 ; 7:5 +ADD R2, R2, 3 ; 7:5 +SHLD R1, 1 ; 7:5 +LD R3, R2 ; 7:5 +MUL R4, R1, R3 ; 7:5 +ADD R0, R0, R4 ; 7:5 +ADD R2, R2, 3 ; 7:5 +SHLD R1, 2 ; 7:5 +LD R3, R2 ; 7:5 +MUL R4, R1, R3 ; 7:5 +ADD R0, R0, R4 ; 7:5 +ADD R2, R2, 3 ; 7:5 +SHLD R1, 3 ; 7:5 +LD R3, R2 ; 7:5 +MUL R4, R1, R3 ; 7:5 +ADD R0, R0, R4 ; 7:5 +ADD R2, R2, 3 ; 7:5 +SHLD R1, 4 ; 7:5 +LD R3, R2 ; 7:5 +MUL R4, R1, R3 ; 7:5 +ADD R0, R0, R4 ; 7:5 +ADD R2, R2, 3 ; 7:5 +SHLD R1, 5 ; 7:5 +LD R3, R2 ; 7:5 +MUL R4, R1, R3 ; 7:5 +ADD R0, R0, R4 ; 7:5 +ADD R2, R2, 3 ; 7:5 +ADD R1, R7, 128 ; 7:5 +ST R1, R0 ; 7:5 diff --git a/tests/golden/tgpu/O2/mm.tgpu b/tests/golden/tgpu/O2/mm.tgpu new file mode 100644 index 0000000..1d742cb --- /dev/null +++ b/tests/golden/tgpu/O2/mm.tgpu @@ -0,0 +1,31 @@ +; kernel k0_matmul: %0 = matmul (t.tiny:4:5) +; grid: 4 blocks x 4 threads; 4 shared words per block; 5 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +MUL R1, R5, 4 ; 4:5 +ADD R2, R6, 16 ; 4:5 +ADD R1, R1, R6 ; 4:5 +LD R1, R1 ; 4:5 +SHST R6, R1 ; 4:5 +SYNCB ; 4:5 +SHLD R1, 0 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 1 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 2 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 3 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +ADD R1, R7, 32 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/O2/mm_bias_relu.tgpu b/tests/golden/tgpu/O2/mm_bias_relu.tgpu new file mode 100644 index 0000000..07fc335 --- /dev/null +++ b/tests/golden/tgpu/O2/mm_bias_relu.tgpu @@ -0,0 +1,62 @@ +; kernel k0_matmul: %2 = matmul (t.tiny:5:5) +; grid: 4 blocks x 4 threads; 8 shared words per block; 5 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 5:5 +MUL R1, R5, 8 ; 5:5 +ADD R2, R6, 32 ; 5:5 +ADD R3, R1, R6 ; 5:5 +LD R3, R3 ; 5:5 +SHST R6, R3 ; 5:5 +ADD R3, R6, 4 ; 5:5 +ADD R1, R1, R3 ; 5:5 +LD R1, R1 ; 5:5 +SHST R3, R1 ; 5:5 +SYNCB ; 5:5 +SHLD R1, 0 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +SHLD R1, 1 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +SHLD R1, 2 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +SHLD R1, 3 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +SHLD R1, 4 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +SHLD R1, 5 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +SHLD R1, 6 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 5:5 +SHLD R1, 7 ; 5:5 +LD R3, R2 ; 5:5 +MUL R4, R1, R3 ; 5:5 +ADD R0, R0, R4 ; 5:5 +ADD R2, R2, 4 ; 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 L0 ; 5:5 +SET R0, 0 ; 5:5 +L0: +ADD R1, R7, 68 ; 5:5 +ST R1, R0 ; 5:5 diff --git a/tests/golden/tgpu/O2/mm_relu.tgpu b/tests/golden/tgpu/O2/mm_relu.tgpu new file mode 100644 index 0000000..ccea473 --- /dev/null +++ b/tests/golden/tgpu/O2/mm_relu.tgpu @@ -0,0 +1,59 @@ +; kernel k0_matmul: %1 = matmul (t.tiny:4:5) +; grid: 4 blocks x 4 threads; 8 shared words per block; 5 registers (R5-R7 reserved, read-only) +SET R0, 0 ; 4:5 +MUL R1, R5, 8 ; 4:5 +ADD R2, R6, 32 ; 4:5 +ADD R3, R1, R6 ; 4:5 +LD R3, R3 ; 4:5 +SHST R6, R3 ; 4:5 +ADD R3, R6, 4 ; 4:5 +ADD R1, R1, R3 ; 4:5 +LD R1, R1 ; 4:5 +SHST R3, R1 ; 4:5 +SYNCB ; 4:5 +SHLD R1, 0 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 1 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 2 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 3 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 4 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 5 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 6 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +SHLD R1, 7 ; 4:5 +LD R3, R2 ; 4:5 +MUL R4, R1, R3 ; 4:5 +ADD R0, R0, R4 ; 4:5 +ADD R2, R2, 4 ; 4:5 +CMP R0, 0 ; 4:5 +BRGT L0 ; 4:5 +SET R0, 0 ; 4:5 +L0: +ADD R1, R7, 64 ; 4:5 +ST R1, R0 ; 4:5 diff --git a/tests/golden/tgpu/big_mm.tgpu b/tests/golden/tgpu/big_mm.tgpu deleted file mode 100644 index e652426..0000000 --- a/tests/golden/tgpu/big_mm.tgpu +++ /dev/null @@ -1,24 +0,0 @@ -; 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 deleted file mode 100644 index 363035e..0000000 --- a/tests/golden/tgpu/const_mm.tgpu +++ /dev/null @@ -1,24 +0,0 @@ -; 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 deleted file mode 100644 index 0f2369b..0000000 --- a/tests/golden/tgpu/mlp_gpu.tgpu +++ /dev/null @@ -1,48 +0,0 @@ -; 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_bias_relu.tgpu b/tests/golden/tgpu/mm_bias_relu.tgpu deleted file mode 100644 index 3a22e4e..0000000 --- a/tests/golden/tgpu/mm_bias_relu.tgpu +++ /dev/null @@ -1,27 +0,0 @@ -; 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 deleted file mode 100644 index 9b98046..0000000 --- a/tests/golden/tgpu/mm_relu.tgpu +++ /dev/null @@ -1,24 +0,0 @@ -; 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/test_gpu_emit.py b/tests/gpu/test_gpu_emit.py index 0123b85..96682ab 100644 --- a/tests/gpu/test_gpu_emit.py +++ b/tests/gpu/test_gpu_emit.py @@ -20,21 +20,30 @@ def test_the_schema_itself_is_valid(): @pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) -@pytest.mark.parametrize("opt", [0, 1]) +@pytest.mark.parametrize("opt", [0, 1, 2]) 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("opt", [0, 1, 2]) @pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) -def test_tgpu_assembly_golden(path, golden): +def test_tgpu_assembly_golden(path, opt, golden): golden( - ROOT / "tests" / "golden" / "tgpu" / f"{path.stem}.tgpu", - build(path.read_text()).program.asm_text(), + ROOT / "tests" / "golden" / "tgpu" / f"O{opt}" / f"{path.stem}.tgpu", + build(path.read_text(), opt=opt).program.asm_text(), ) +def test_tiled_kernels_declare_their_shared_memory_in_the_launch_document(): + comp = build("input A : i8[4, 8]\ninput B : i8[8, 4]\nC = A @ B\noutput C\n", opt=2) + (k,) = comp.launch["kernels"] + assert k["shared_size"] == 8 and "SYNCB" in k["asm"] and "SHLD" in k["asm"] + assert "8 shared words per block" in k["asm"] + assert comp.launch["estimate"]["per_kernel"][0]["shared_size"] == 8 + + def test_emission_is_deterministic(): src = (ROOT / "examples" / "mlp_gpu.tiny").read_text() assert dumps(build(src).launch) == dumps(build(src).launch) diff --git a/tests/gpu/test_gpu_perf.py b/tests/gpu/test_gpu_perf.py index 1a02700..f04e549 100644 --- a/tests/gpu/test_gpu_perf.py +++ b/tests/gpu/test_gpu_perf.py @@ -1,24 +1,29 @@ -"""Pinned static cost numbers: an unexpected change in launches/instructions/loads fails CI.""" +"""Pinned static cost numbers: an unexpected change in launches/instructions/loads fails CI. + +The optimization levels are pinned side by side, so a pass that stops helping (or starts hurting) +shows up as a diff here. Update deliberately with `pytest --update-golden`. +""" from tests.gpu.helpers import EXAMPLES, build +KEYS = ( + "launches", + "threads", + "instructions", + "cycles", + "loads", + "stores", + "shared_loads", + "shared_stores", + "registers", +) + def test_gpu_cost_numbers_are_pinned(perf_pin): pinned = {} - for opt in (0, 1): + for opt in (0, 1, 2): 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", - ) - } + pinned[f"O{opt}"][path.stem] = {k: est[k] for k in KEYS} perf_pin("gpu", pinned) diff --git a/tests/gpu/test_gpu_properties.py b/tests/gpu/test_gpu_properties.py index 3410d81..24879e9 100644 --- a/tests/gpu/test_gpu_properties.py +++ b/tests/gpu/test_gpu_properties.py @@ -15,16 +15,17 @@ def _compile(source, opt): try: - return build( - source, opt=opt, verify_each=True - ) # kernel interpreter AND simulator each stage - except PassError: + # every kernel pass is checked by the interpreter, and the assembly by the simulator + return build(source, opt=opt, verify_each=True) + except PassError as e: + if "registers" in str(e): # register pressure at the default budget would be a real finding + raise 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)) +@given(SOURCES, st.sampled_from([0, 1, 2]), 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: @@ -36,8 +37,10 @@ def test_kernel_interpreter_matches_oracle_on_random_programs(source, opt, seed) @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): +@given(SOURCES, st.sampled_from([0, 1, 2]), st.integers(0, 2**32 - 1)) +def test_simulator_matches_oracle_interpreter_counts_and_estimate_on_random_programs( + source, opt, seed +): comp = _compile(source, opt) if comp is None: return @@ -47,3 +50,9 @@ def test_simulator_matches_oracle_and_interpreter_counts_on_random_programs(sour 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 + est = comp.program.estimate + assert est["instructions"]["min"] <= res.instructions <= est["instructions"]["max"] + assert est["cycles"]["min"] <= res.cycles <= est["cycles"]["max"] + ops = res.op_counts + assert (ops["LD"], ops["ST"]) == (est["loads"], est["stores"]) + assert (ops["SHLD"], ops["SHST"]) == (est["shared_loads"], est["shared_stores"]) diff --git a/tests/gpu/test_gpu_sim.py b/tests/gpu/test_gpu_sim.py index ef585f7..82635f7 100644 --- a/tests/gpu/test_gpu_sim.py +++ b/tests/gpu/test_gpu_sim.py @@ -15,7 +15,7 @@ @pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) -@pytest.mark.parametrize("opt", [0, 1]) +@pytest.mark.parametrize("opt", [0, 1, 2]) 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): @@ -25,21 +25,23 @@ def test_examples_are_bit_exact_and_never_touch_the_reserved_registers(path, opt @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): +@pytest.mark.parametrize("opt", [0, 1, 2]) +def test_estimate_bounds_contain_the_measurement_and_memory_ops_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"] + assert res.op_counts["SHLD"] == est["shared_loads"] + assert res.op_counts["SHST"] == est["shared_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]) +@pytest.mark.parametrize("opt", [0, 1, 2]) 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.""" @@ -54,8 +56,9 @@ def test_interpreter_counts_equal_simulator_counts_exactly(path, opt): 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) +@pytest.mark.parametrize("opt", [0, 1, 2]) +def test_the_estimate_bounds_are_attained_at_the_data_extremes(opt): + comp = build(MM_RELU, opt=opt) 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 @@ -73,12 +76,38 @@ def test_a_straight_line_kernel_is_one_cycle_however_many_instructions(): assert res.cycles == 1 and res.instructions == 8 * 7 # 7 instructions per thread, 8 threads +@pytest.mark.parametrize("opt", [0, 1, 2]) +def test_without_data_dependent_branches_the_estimate_is_exact_at_every_level(opt): + """No relu -> nothing is uncertain: instructions and cycles must equal the measurement, + including the barrier model of the tiled -O2 kernel.""" + comp = build("input A : i8[4, 8]\ninput B : i8[8, 4]\nC = A @ B\noutput C\n", opt=opt) + est = comp.program.estimate + assert est["instructions"]["min"] == est["instructions"]["max"] + assert est["cycles"]["min"] == est["cycles"]["max"] + res, _, _ = run_on_simulator(comp, 5) + assert res.instructions == est["instructions"]["min"] + assert res.cycles == est["cycles"]["min"] + assert res.op_counts["SYNCB"] == interp_syncb(comp, res) + + +def interp_syncb(comp, res): + inputs = ref.random_inputs(comp.func, np.random.default_rng(5)) + got = interp.run_program(comp.program, inputs)["kernels"] + return sum(k.ops["SYNCB"] for k in got) + + def test_matmul_cycles_are_the_loop_trip_count_plus_the_relu_cycle(): - comp = build(MM_RELU) # K = 8 + comp = build(MM_RELU, opt=1) # K = 8, loop kept res, _, _ = run_on_simulator(comp, 1) assert res.cycles in (8, 9) # 8 iterations; +1 when some thread does not clamp +def test_unrolling_removes_the_loop_cycles(): + comp = build(MM_RELU, opt=2) # the K loop is unrolled: no taken back edges at all + res, _, _ = run_on_simulator(comp, 1) + assert res.cycles in (2, 3) # arrive at the barrier, resume, (+1 for a non-clamping thread) + + def test_global_memory_persists_across_launches(): comp = build((ROOT / "examples" / "mlp_gpu.tiny").read_text()) assert len(comp.launch["kernels"]) == 2 @@ -98,7 +127,7 @@ def test_unfused_o0_chain_matches_the_fused_o1_program(): def test_the_smallest_legal_register_budget_still_works(): - comp = build(MM_RELU, regs=8) + comp = build(MM_RELU, opt=0, regs=8) # -O0 is the register-frugal code res, _, want = run_on_simulator(comp) assert_bit_exact(comp, res, want) diff --git a/tests/gpu/test_kernel_ir.py b/tests/gpu/test_kernel_ir.py index 9eafc85..223bd59 100644 --- a/tests/gpu/test_kernel_ir.py +++ b/tests/gpu/test_kernel_ir.py @@ -55,15 +55,17 @@ def test_chained_matmuls_are_two_launches_over_shared_memory(): assert [(k.kind, k.grid) for k in ks] == [("matmul", (4, 6)), ("matmul", (4, 3))] +@pytest.mark.parametrize("opt", [0, 1, 2]) @pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) -def test_kernel_ir_golden(path, golden): - prog = build(path.read_text()).program +def test_kernel_ir_golden(path, opt, golden): + prog = build(path.read_text(), opt=opt).program text = "\n".join(print_kernel(k) for k in prog.kernels) - golden(ROOT / "tests" / "golden" / "kernel" / f"{path.stem}.kernel", text) + golden(ROOT / "tests" / "golden" / "kernel" / f"O{opt}" / 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]) + src = "input A : i8[4, 8]\ninput B : i8[8, 4]\nC = A @ B\noutput C\n" + text = print_kernel(kernels(src, opt=0)[0]) assert "for v1 in 0..8 {" in text assert ( text.count("= ld [") == 2 and text.count("= mul ") == 3 diff --git a/tests/gpu/test_kernel_passes.py b/tests/gpu/test_kernel_passes.py new file mode 100644 index 0000000..7493127 --- /dev/null +++ b/tests/gpu/test_kernel_passes.py @@ -0,0 +1,504 @@ +"""The kernel optimizations: directed tests per pass, then what they buy, measured on the simulator.""" + +import copy +import dataclasses +from collections import Counter + +import numpy as np +import pytest + +from runtime.gpu_runner import run_tgpu +from tests.gpu.helpers import EXAMPLES, ROOT, build, run_on_simulator +from tinyc.gpu import cost, interp +from tinyc.gpu.emit import emit_asm +from tinyc.gpu.kernel_ir import ( + Assign, + Barrier, + For, + If, + Kernel, + Load, + SharedLoad, + SharedStore, + Special, + Store, + VReg, + flatten, + print_kernel, + verify_kernel, +) +from tinyc.gpu.kernel_passes import ( + KERNEL_PASSES, + LICM, + Fold, + KernelDCE, + SharedTile, + StrengthReduce, + Unroll, + kernel_pass_names, +) +from tinyc.gpu.regalloc import allocate +from tinyc.ir import ref +from tinyc.ir.types import epilogue + +v = [VReg(i) for i in range(16)] +BLOCK, THREAD, TID = Special("block"), Special("thread"), Special("tid") + + +def matmul_src(m, k, n, body="C = A @ B"): + return f"input A : i8[{m}, {k}]\ninput B : i8[{k}, {n}]\n{body}\noutput C\n" + + +def lowered(source, opt=0): + """The kernels of a program, copied so a pass can be applied by hand.""" + prog = build(source, opt=opt).program + return prog, copy.deepcopy(prog.kernels) + + +def kernel(body, grid=(1, 4), shared=0): + return Kernel("k", grid, body, "0", "matmul", None, shared) + + +def run_kernel(k: Kernel, memory: list[int]) -> tuple[list[int], Counter]: + memory, ops = list(memory), Counter() + for block in range(k.grid[0]): + interp.run_block(k, memory, block, ops) + return memory, ops + + +def outputs_of(prog, kernels, seed=1): + inputs = ref.random_inputs(prog.func, np.random.default_rng(seed)) + return interp.run_program(dataclasses.replace(prog, kernels=kernels), inputs)["outputs"], inputs + + +def same_outputs(prog, kernels, seed=1): + got, inputs = outputs_of(prog, kernels, seed) + for name, want in zip([x.name for x in prog.func.outputs], ref.evaluate(prog.func, inputs)): + np.testing.assert_array_equal(got[name], want) + + +def top_loop(k: Kernel) -> For: + return next(s for s in k.body if isinstance(s, For)) + + +def kinds(body) -> list[str]: + return [s.op if isinstance(s, Assign) else type(s).__name__.lower() for s in body] + + +# --- levels and names ------------------------------------------------------------------------- + + +def test_levels_select_the_documented_passes(): + assert kernel_pass_names(0) == [] + assert kernel_pass_names(1) == ["licm", "strength-reduce", "kernel-dce"] + assert kernel_pass_names(2) == [ + "licm", "strength-reduce", "shared-tile", "unroll", "fold", "kernel-dce", + ] # fmt: skip + assert set(kernel_pass_names(2)) <= set(KERNEL_PASSES) + + +# --- licm -------------------------------------------------------------------------------------- + + +def test_licm_hoists_the_address_part_that_does_not_depend_on_k(): + prog, ks = lowered(matmul_src(4, 8, 4)) + (k,) = ks + LICM().run(k) + verify_kernel(k) + loop = top_loop(k) + mul_by_block = lambda s: isinstance(s, Assign) and s.op == "mul" and BLOCK in s.args + assert not any(mul_by_block(s) for s in loop.body) and any(mul_by_block(s) for s in k.body) + assert sum(isinstance(s, Load) for s in loop.body) == 2 # loads are never hoisted + assert any(s.op == "mul" and loop.var in s.args for s in loop.body if isinstance(s, Assign)) + same_outputs(prog, ks) + + +def test_licm_leaves_multiply_assigned_values_and_loads_in_place(): + k = kernel([ + Assign("set", v[0], (0,)), + Assign("set", v[1], (5,)), + For(v[2], 0, 3, [ + Assign("add", v[1], (v[1], 1)), # v1 changes every iteration: not invariant + Assign("mul", v[3], (BLOCK, 7)), # invariant + Load(v[4], v[3]), + Assign("add", v[0], (v[0], v[4])), + ]), + Store(TID, v[0]), + ]) # fmt: skip + LICM().run(k) + loop = top_loop(k) + assert kinds(loop.body) == ["add", "load", "add"] # only the invariant multiply left the loop + assert any(isinstance(s, Assign) and s.op == "mul" for s in k.body) + + +def test_licm_hoists_through_nested_loops(): + k = kernel([ + Assign("set", v[0], (0,)), + For(v[1], 0, 2, [For(v[2], 0, 2, [ + Assign("mul", v[3], (BLOCK, 3)), + Assign("add", v[0], (v[0], v[3])), + ])]), + Store(TID, v[0]), + ]) # fmt: skip + LICM().run(k) + assert kinds(k.body)[:2] == ["set", "mul"] # the multiply climbed out of both loops + + +# --- strength reduction ------------------------------------------------------------------------ + + +def test_strength_reduce_turns_the_b_address_chain_into_a_pointer(): + prog, ks = lowered(matmul_src(4, 8, 4)) + (k,) = ks + for p in (LICM(), StrengthReduce(), KernelDCE()): + p.run(k) + verify_kernel(k) + loop = top_loop(k) + muls = [s for s in loop.body if isinstance(s, Assign) and s.op == "mul"] + assert len(muls) == 1 # only the product x*y is left: no `mul k, N` + bumps = [s for s in loop.body if isinstance(s, Assign) and s.op == "add" and s.dst == s.args[0]] + assert any(s.args[1] == 4 for s in bumps) # the B pointer advances by N per iteration + same_outputs(prog, ks) + + +def test_strength_reduce_handles_a_loop_that_does_not_start_at_zero(): + def build_kernel(): + return kernel( + [ + Assign("set", v[0], (0,)), + For(v[1], 2, 6, [ + Assign("mul", v[2], (v[1], 3)), + Assign("add", v[3], (v[2], THREAD)), + Assign("add", v[4], (v[3], 100)), + Load(v[5], v[4]), + Assign("add", v[0], (v[0], v[5])), + ]), + Assign("add", v[6], (TID, 200)), + Store(v[6], v[0]), + ], + grid=(1, 3), + ) # fmt: skip + + memory = list(range(1000, 1300)) + before, ops_before = run_kernel(build_kernel(), memory) + k = build_kernel() + for p in (StrengthReduce(), KernelDCE()): + p.run(k) + verify_kernel(k) + after, ops_after = run_kernel(k, memory) + assert after == before + assert ops_after["MUL"] < ops_before["MUL"] + + +def test_strength_reduce_skips_chains_that_are_read_after_the_loop_or_are_too_short(): + k = kernel([ + Assign("set", v[0], (0,)), + For(v[1], 0, 3, [ + Assign("mul", v[2], (v[1], 3)), + Assign("add", v[3], (v[2], THREAD)), + Assign("add", v[0], (v[0], v[3])), + ]), + Store(TID, v[3]), # v3 is read after the loop: it must keep its per-iteration meaning + ]) # fmt: skip + before = print_kernel(k) + StrengthReduce().run(k) + assert print_kernel(k) == before + short = kernel([ + Assign("set", v[0], (0,)), + For(v[1], 0, 3, [Assign("add", v[2], (v[1], 9)), Assign("add", v[0], (v[0], v[2]))]), + Store(TID, v[0]), + ]) # fmt: skip + before = print_kernel(short) + StrengthReduce().run(short) + assert print_kernel(short) == before # one add costs the same as the increment: no gain + + +# --- dead code elimination --------------------------------------------------------------------- + + +def test_kernel_dce_removes_dead_chains_and_keeps_what_reaches_a_store(): + k = kernel([ + Assign("set", v[0], (1,)), + Assign("add", v[1], (v[0], 2)), + Assign("add", v[2], (v[1], 3)), # dead chain + Assign("set", v[3], (9,)), + Load(v[4], TID), # dead load + Store(TID, v[3]), + ]) # fmt: skip + KernelDCE().run(k) + assert kinds(k.body) == ["set", "store"] + + +def test_kernel_dce_removes_a_dead_value_that_only_updates_itself(): + k = kernel([ + Assign("set", v[0], (0,)), + For(v[1], 0, 3, [Assign("add", v[0], (v[0], 1))]), # v0 is never read by anything live + Store(TID, 5), + ]) # fmt: skip + KernelDCE().run(k) + assert isinstance(k.body[0], For) and k.body[0].body == [] + + +# --- shared-memory tiling ---------------------------------------------------------------------- + + +def tiled(m, k, n): + prog, ks = lowered(matmul_src(m, k, n), opt=1) + (kern,) = ks + SharedTile().run(kern) + verify_kernel(kern) + return prog, ks, kern + + +def test_shared_tile_copies_the_a_row_once_per_block_and_reads_it_with_shld(): + prog, ks, k = tiled(4, 8, 4) + assert k.shared_size == 8 + assert isinstance( + k.body[[type(s) for s in k.body].index(Barrier) + 1], For + ) # barrier, then loop + loop = top_loop(k) + assert sum(isinstance(s, SharedLoad) for s in loop.body) == 1 + assert sum(isinstance(s, Load) for s in loop.body) == 1 # only B is still loaded from global + same_outputs(prog, ks) + + +@pytest.mark.parametrize( + "m,k,n,guards", [(4, 8, 4, 0), (4, 6, 4, 1), (3, 3, 4, 1), (2, 5, 2, 1), (2, 8, 2, 0)] +) +def test_the_cooperative_copy_guards_only_the_rounds_that_can_run_past_the_row(m, k, n, guards): + prog, ks, kern = tiled(m, k, n) + assert sum(isinstance(s, If) for s in kern.body) == guards + stores = [s for s in kern.body if isinstance(s, SharedStore)] + stores += [ + t for s in kern.body if isinstance(s, If) for t in s.body if isinstance(t, SharedStore) + ] + assert len(stores) == -(-k // n) # ceil(K / N) rounds of one shared store each + same_outputs(prog, ks) + + +def test_shared_tile_skips_one_thread_blocks_and_kernels_without_a_uniform_load(): + _, ks = lowered(matmul_src(4, 8, 1), opt=1) # N = 1: a block is a single thread + SharedTile().run(ks[0]) + assert ks[0].shared_size == 0 and not any(isinstance(s, Barrier) for s in ks[0].body) + _prog, ks = lowered("input a : i8[8]\ninput b : i8[8]\nC = a + b\noutput C\n", opt=1) + SharedTile().run(ks[0]) + assert ks[0].shared_size == 0 + + +def test_a_tiled_kernel_does_k_global_loads_per_block_instead_of_n_times_k(): + _prog, _ks, k = tiled(4, 8, 4) + _, ops = run_kernel(k, [1] * 400) # data does not matter for counts here + blocks, n, kk = 4, 4, 8 + assert ops["LD"] == blocks * kk + blocks * n * kk # A row once per block + B per thread + assert ops["SHST"] == blocks * kk and ops["SHLD"] == blocks * n * kk + + +# --- unrolling --------------------------------------------------------------------------------- + + +def test_unroll_replaces_a_short_loop_by_straight_line_copies_with_constant_counters(): + prog, ks = lowered(matmul_src(2, 4, 2), opt=1) + (k,) = ks + Unroll().run(k) + verify_kernel(k) + assert not any(isinstance(s, For) for s in k.body) + assert sum(isinstance(s, Load) for s in k.body) == 2 * 4 # two loads per unrolled iteration + same_outputs(prog, ks) + + +def test_unroll_leaves_long_loops_and_loops_whose_counter_is_read_later(): + _prog, ks = lowered(matmul_src(2, 17, 2), opt=1) # 17 iterations > 16 + Unroll().run(ks[0]) + assert any(isinstance(s, For) for s in ks[0].body) + after_use = kernel([ + Assign("set", v[0], (0,)), + For(v[1], 0, 3, [Assign("add", v[0], (v[0], 1))]), + Store(TID, v[1]), # the counter's final value is read + ]) # fmt: skip + Unroll().run(after_use) + assert isinstance(after_use.body[1], For) + + +# --- constant folding -------------------------------------------------------------------------- + + +def test_fold_propagates_constants_and_reassociates_add_chains(): + k = kernel([ + Assign("set", v[0], (3,)), + Assign("add", v[1], (v[0], 4)), + Assign("add", v[2], (v[1], THREAD)), + Assign("add", v[3], (v[2], 10)), + Store(v[3], 1), + ]) # fmt: skip + Fold().run(k) + KernelDCE().run(k) + verify_kernel(k) + assert [type(s).__name__ for s in k.body] == ["Assign", "Store"] + assert k.body[0].args == (THREAD, 17) # (3 + 4 + thread) + 10 = thread + 17 + + +def test_fold_copy_propagates_adds_of_zero(): + k = kernel([ + Assign("add", v[0], (THREAD, 0)), # a copy of the thread id + Assign("add", v[1], (v[0], 0)), # a copy of the copy + Assign("add", v[2], (v[1], 5)), + Store(v[2], 1), + ]) # fmt: skip + Fold().run(k) + KernelDCE().run(k) + assert [type(s).__name__ for s in k.body] == ["Assign", "Store"] + assert k.body[0].args == (THREAD, 5) + + +def test_fold_does_not_copy_propagate_a_value_that_changes(): + k = kernel([ + Assign("set", v[0], (1,)), + Assign("add", v[1], (v[0], 0)), # copy of v0 as it is *now* + Assign("add", v[0], (v[0], 10)), # v0 changes afterwards + Assign("add", v[2], (v[1], TID)), + Store(v[2], v[1]), + ], grid=(1, 4)) # fmt: skip + memory = [0] * 32 + before, _ = run_kernel(k, memory) + Fold().run(k) + after, _ = run_kernel(k, memory) + assert after == before + + +def test_fold_does_not_reassociate_through_a_value_that_changes(): + k = kernel([ + Assign("set", v[0], (0,)), + Assign("add", v[1], (v[0], 1)), + Assign("add", v[0], (v[0], 5)), # v0 changes between v1's definition and v2's + Assign("add", v[2], (v[1], 2)), + Store(TID, v[2]), + ]) # fmt: skip + memory = [0] * 8 + before, _ = run_kernel(k, memory) + Fold().run(k) + after, _ = run_kernel(k, memory) + assert after == before and k.body[3].args == (v[1], 2) + + +# --- the barrier model, against the real simulator ---------------------------------------------- + + +def barrier_kernel(divergent_after: bool) -> Kernel: + body = [ + If(THREAD, 1, [Load(v[0], TID), SharedStore(THREAD, v[0])]), # threads 0,1 copy; 2,3 skip + Barrier(), + SharedLoad(v[1], 0), + SharedLoad(v[2], 1), + Assign("add", v[3], (v[1], v[2])), + ] + if divergent_after: + body.append(If(v[3], 100, [Assign("add", v[3], (v[3], 1000))])) # sum <= 100: some threads + body += [Assign("add", v[4], (TID, 8)), Store(v[4], v[3])] + return Kernel("k", (2, 4), body, "0", "matmul", None, 4) + + +@pytest.mark.parametrize("divergent_after", [False, True]) +def test_the_barrier_model_matches_the_simulator_exactly_when_arrival_times_differ(divergent_after): + k = barrier_kernel(divergent_after) + verify_kernel(k) + memory = [10, 20, 0, 0, 30, 130, 0, 0] + [0] * 8 + expected_mem, ops = run_kernel(k, memory) + + insts = flatten(k) + asm = emit_asm(k, insts, allocate(insts, 16, "k")) + sim = run_tgpu( + asm, grid=k.grid, mem_size=16, memory_init={0: memory}, shared_size=k.shared_size + ) + assert sim.read(0, 16) == expected_mem + assert dict(sim.op_counts) == dict(ops) # includes the SYNCB re-executions of waiting threads + assert sim.reserved_intact + + # threads 0,1 arrive at step 1, threads 2,3 (taken skip branch) at step 2; release at step 2 + assert ops["SYNCB"] == 2 * (2 + 2 + 1 + 1) + interp_cycles = max(interp.run_block(k, list(memory), b, Counter()) for b in range(2)) + assert sim.cycles == interp_cycles + + est = cost.kernel_cost(k, 0) # the abstract model must agree too (nothing is data dependent + if not divergent_after: # except the second `if`, which depends on loaded data) + assert est["instructions"]["min"] == est["instructions"]["max"] == sim.instructions + assert est["cycles"]["min"] == est["cycles"]["max"] == sim.cycles + else: + assert est["instructions"]["min"] <= sim.instructions <= est["instructions"]["max"] + assert est["cycles"]["min"] <= sim.cycles <= est["cycles"]["max"] + + +def test_a_barrier_inside_a_loop_or_if_is_rejected(): + from tinyc.gpu.kernel_ir import KernelError + + with pytest.raises(KernelError, match="barrier inside"): + verify_kernel(kernel([For(v[0], 0, 2, [Barrier()]), Store(TID, 1)], shared=1)) + with pytest.raises(KernelError, match="shared memory used but shared_size is 0"): + verify_kernel(kernel([SharedStore(0, 1), Store(TID, 1)])) + with pytest.raises(KernelError, match="shared address 5 outside"): + verify_kernel(kernel([SharedStore(5, 1), Store(TID, 1)], shared=2)) + + +# --- what the optimizations buy, measured on the simulator --------------------------------------- + + +def matmul_loads(comp): + """Expected global loads per matmul kernel at -O0/-O1, and with shared tiling at -O2.""" + out = {} + for op in comp.func.ops: + if op.kind != "matmul": + continue + (m, k), n = op.operands[0].type.shape, op.operands[1].type.shape[1] + bias = m * n if "bias" in epilogue(op.attrs) else 0 + out[op.result.name] = (2 * m * n * k + bias, m * k + m * n * k + bias, m, n, k) + return out + + +def non_matmul_loads(comp): + """Loads of add/relu kernels, which no level changes: 2 per add thread, 1 per relu thread.""" + per_kernel = comp.program.estimate["per_kernel"] + return sum(p["loads"] for p, k in zip(per_kernel, comp.program.kernels) if k.kind != "matmul") + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_o2_has_strictly_fewer_global_loads_than_o0_on_every_matmul(path): + """The Phase 5 exit criterion, with the exact formula: 2*M*N*K -> M*K + M*N*K per matmul.""" + comps = {o: build(path.read_text(), opt=o) for o in (0, 1, 2)} + measured = {o: run_on_simulator(c, 4)[0].op_counts["LD"] for o, c in comps.items()} + for o, comp in comps.items(): + assert measured[o] == comp.program.estimate["loads"] # the estimate is exact for loads + for level, column in ((0, 0), (1, 0), (2, 1)): # naive formula at -O0/-O1, tiled at -O2 + comp = comps[level] + expected = non_matmul_loads(comp) + sum(e[column] for e in matmul_loads(comp).values()) + assert measured[level] == expected + # hoisting and strength reduction never remove or add a load (same graph, passes switched off) + plain = build(path.read_text(), opt=1, disabled=["licm", "strength-reduce"]) + assert run_on_simulator(plain, 4)[0].op_counts["LD"] == measured[1] + assert measured[2] < measured[0] + + +@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) +def test_each_level_does_less_work_than_the_one_before(path): + r = {o: run_on_simulator(build(path.read_text(), opt=o), 6)[0] for o in (0, 1, 2)} + assert r[1].instructions < r[0].instructions # hoisting and pointer bumping remove work + assert r[1].cycles <= r[0].cycles # same trip count, plus fusion's saved launches + assert r[2].cycles <= r[1].cycles # no taken back edges at all + + +@pytest.mark.parametrize("opt", [1, 2]) +def test_optimized_and_naive_kernels_compute_the_same_thing(opt): + for path in EXAMPLES: + naive, fast = build(path.read_text(), opt=0), build(path.read_text(), opt=opt) + inputs = ref.random_inputs(naive.func, np.random.default_rng(13)) + a = interp.run_program(naive.program, inputs)["outputs"] + b = interp.run_program(fast.program, inputs)["outputs"] + for n0, n1 in zip(naive.launch["outputs"], fast.launch["outputs"]): + np.testing.assert_array_equal(a[n0], b[n1]) + + +def test_the_big_matmul_is_unrolled_only_when_the_loop_is_short_enough(): + big = build((ROOT / "examples" / "big_mm.tiny").read_text(), opt=2) # K = 12 <= 16 + assert "BRLT" not in big.program.asm_text() + long_k = build(matmul_src(2, 20, 3), opt=2) # K = 20 > 16: loop stays, A is still tiled + assert "BRLT" in long_k.program.asm_text() and "SYNCB" in long_k.program.asm_text() + res, _, want = run_on_simulator(long_k, 2) + np.testing.assert_array_equal(res.outputs["0"], want[0]) diff --git a/tests/gpu/test_opt_table.py b/tests/gpu/test_opt_table.py new file mode 100644 index 0000000..62ee9de --- /dev/null +++ b/tests/gpu/test_opt_table.py @@ -0,0 +1,37 @@ +"""docs/ARCHITECTURE.md quotes measured numbers; they must stay equal to what the tools measure.""" + +import re + +from tests.gpu.helpers import ROOT +from tinyc.gpu.report import CONFIGS, PROGRAMS, opt_table_markdown + +DOC = ROOT / "docs" / "ARCHITECTURE.md" + + +def test_the_documented_optimization_table_matches_a_fresh_measurement(): + text = DOC.read_text() + m = re.search(r"\n(.*?)\n", text, re.DOTALL) + assert m, "markers missing from docs/ARCHITECTURE.md" + stale = "ARCHITECTURE.md is stale: run `python scripts/opt_table.py --write`" + assert m.group(1).strip() == opt_table_markdown().strip(), stale + + +def test_every_configuration_in_the_table_is_correct_and_costs_are_ordered(): + """Measuring asserts bit-exactness; here, the claims the prose makes about the numbers.""" + from tinyc.gpu.report import measure + + src = PROGRAMS["8x8"] + m = {label: measure(src, opt, off) for label, opt, off in CONFIGS} + o0, licm, o1 = m["-O0 (naive)"], m["+ licm"], m["+ strength-reduce (= -O1)"] + tile, unroll, o2 = ( + m["-O2 with shared-memory tiling only"], + m["-O2 with unrolling + folding only"], + m["-O2 (all passes)"], + ) + assert o0["instructions"] > licm["instructions"] > o1["instructions"] + assert licm["registers"] > o0["registers"] and o1["registers"] >= licm["registers"] + assert o1["loads"] == licm["loads"] < o0["loads"] + 1 # loads are not removed by licm/sr + assert tile["loads"] < o1["loads"] and unroll["loads"] == o1["loads"] + assert unroll["cycles"] < o1["cycles"] and o2["cycles"] < o1["cycles"] + assert o2["instructions"] < o1["instructions"] + assert o2["loads"] == tile["loads"] and o2["instructions"] <= unroll["instructions"] diff --git a/tests/gpu/test_regalloc.py b/tests/gpu/test_regalloc.py index 3787773..6ea78db 100644 --- a/tests/gpu/test_regalloc.py +++ b/tests/gpu/test_regalloc.py @@ -101,17 +101,28 @@ def test_assertion_catches_writes_to_reserved_registers(): assert_reserved_untouched(insts, {v[0]: 4}) +@pytest.mark.parametrize("opt", [0, 1, 2]) @pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem) -def test_real_kernels_never_use_r5_to_r7_as_destinations(path): +def test_real_kernels_never_use_r5_to_r7_as_destinations(path, opt): import re - for asm in build(path.read_text()).program.asm.values(): + for asm in build(path.read_text(), opt=opt).program.asm.values(): for line in asm.splitlines(): - m = re.match(r"(SET|ADD|MUL|LD) R(\d+),", line) + m = re.match(r"(SET|ADD|MUL|LD|SHLD) 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} +MM_RELU = "input A : i8[4, 8]\ninput B : i8[8, 4]\nC = relu(A @ B)\noutput C\n" + + +def test_naive_matmul_needs_only_four_registers_and_fits_the_smallest_legal_budget(): + comp = build(MM_RELU, opt=0, regs=8) + assert comp.program.registers["k0_matmul"] == 4 and max(comp.program.registers.values()) == 4 + + +def test_optimizing_trades_registers_for_instructions(): + """Hoisted addresses and pointers stay live across the loop, so -O1/-O2 need more registers.""" + used = {o: max(build(MM_RELU, opt=o).program.registers.values()) for o in (0, 1, 2)} + assert used[0] == 4 and used[1] >= used[0] and used[2] >= used[0] + assert max(used.values()) <= 13, "must fit the default register file" diff --git a/tests/perf_baseline.json b/tests/perf_baseline.json index bea457c..8e02f4b 100644 --- a/tests/perf_baseline.json +++ b/tests/perf_baseline.json @@ -13,6 +13,8 @@ "launches": 2, "loads": 3600, "registers": 4, + "shared_loads": 0, + "shared_stores": 0, "stores": 288, "threads": 288 }, @@ -28,6 +30,8 @@ "launches": 2, "loads": 30, "registers": 4, + "shared_loads": 0, + "shared_stores": 0, "stores": 12, "threads": 12 }, @@ -43,6 +47,8 @@ "launches": 4, "loads": 600, "registers": 4, + "shared_loads": 0, + "shared_stores": 0, "stores": 84, "threads": 84 }, @@ -58,6 +64,8 @@ "launches": 1, "loads": 128, "registers": 4, + "shared_loads": 0, + "shared_stores": 0, "stores": 16, "threads": 16 }, @@ -73,6 +81,8 @@ "launches": 3, "loads": 304, "registers": 4, + "shared_loads": 0, + "shared_stores": 0, "stores": 48, "threads": 48 }, @@ -88,6 +98,8 @@ "launches": 2, "loads": 272, "registers": 4, + "shared_loads": 0, + "shared_stores": 0, "stores": 32, "threads": 32 } @@ -99,12 +111,14 @@ "min": 12 }, "instructions": { - "max": 23472, - "min": 23328 + "max": 16992, + "min": 16848 }, "launches": 1, "loads": 3456, - "registers": 4, + "registers": 6, + "shared_loads": 0, + "shared_stores": 0, "stores": 144, "threads": 144 }, @@ -114,12 +128,14 @@ "min": 2 }, "instructions": { - "max": 198, - "min": 192 + "max": 168, + "min": 162 }, "launches": 1, "loads": 24, - "registers": 4, + "registers": 6, + "shared_loads": 0, + "shared_stores": 0, "stores": 6, "threads": 6 }, @@ -129,12 +145,14 @@ "min": 14 }, "instructions": { - "max": 3720, - "min": 3696 + "max": 2772, + "min": 2748 }, "launches": 2, "loads": 552, - "registers": 4, + "registers": 6, + "shared_loads": 0, + "shared_stores": 0, "stores": 36, "threads": 36 }, @@ -144,12 +162,14 @@ "min": 4 }, "instructions": { - "max": 896, - "min": 896 + "max": 688, + "min": 688 }, "launches": 1, "loads": 128, - "registers": 4, + "registers": 6, + "shared_loads": 0, + "shared_stores": 0, "stores": 16, "threads": 16 }, @@ -159,12 +179,14 @@ "min": 8 }, "instructions": { - "max": 1824, - "min": 1808 + "max": 1360, + "min": 1344 }, "launches": 1, "loads": 272, - "registers": 4, + "registers": 6, + "shared_loads": 0, + "shared_stores": 0, "stores": 16, "threads": 16 }, @@ -174,12 +196,118 @@ "min": 8 }, "instructions": { - "max": 1776, - "min": 1760 + "max": 1312, + "min": 1296 }, "launches": 1, "loads": 256, - "registers": 4, + "registers": 6, + "shared_loads": 0, + "shared_stores": 0, + "stores": 16, + "threads": 16 + } + }, + "O2": { + "big_mm": { + "cycles": { + "max": 3, + "min": 2 + }, + "instructions": { + "max": 10368, + "min": 10224 + }, + "launches": 1, + "loads": 1872, + "registers": 5, + "shared_loads": 1728, + "shared_stores": 144, + "stores": 144, + "threads": 144 + }, + "const_mm": { + "cycles": { + "max": 3, + "min": 2 + }, + "instructions": { + "max": 132, + "min": 126 + }, + "launches": 1, + "loads": 18, + "registers": 5, + "shared_loads": 12, + "shared_stores": 6, + "stores": 6, + "threads": 6 + }, + "mlp_gpu": { + "cycles": { + "max": 6, + "min": 5 + }, + "instructions": { + "max": 1952, + "min": 1928 + }, + "launches": 2, + "loads": 344, + "registers": 5, + "shared_loads": 264, + "shared_stores": 56, + "stores": 36, + "threads": 36 + }, + "mm": { + "cycles": { + "max": 2, + "min": 2 + }, + "instructions": { + "max": 464, + "min": 464 + }, + "launches": 1, + "loads": 80, + "registers": 5, + "shared_loads": 64, + "shared_stores": 16, + "stores": 16, + "threads": 16 + }, + "mm_bias_relu": { + "cycles": { + "max": 3, + "min": 2 + }, + "instructions": { + "max": 944, + "min": 928 + }, + "launches": 1, + "loads": 176, + "registers": 5, + "shared_loads": 128, + "shared_stores": 32, + "stores": 16, + "threads": 16 + }, + "mm_relu": { + "cycles": { + "max": 3, + "min": 2 + }, + "instructions": { + "max": 896, + "min": 880 + }, + "launches": 1, + "loads": 160, + "registers": 5, + "shared_loads": 128, + "shared_stores": 32, "stores": 16, "threads": 16 } diff --git a/tests/test_cli.py b/tests/test_cli.py index 8070825..8b143e9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -164,11 +164,34 @@ def test_the_same_program_runs_on_both_chips_from_the_same_seed(capsys): 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") + # -O0 is the register-frugal code, so it fits the smallest legal register file + base = ["--target", "gpu", "--run", "-O", "0"] + assert run(capsys, EX / "mm_relu.tiny", *base, "--gpu-regs", "8")[0] == 0 + code, _, err = run(capsys, EX / "mm_relu.tiny", *base, "--gpu-regs", "7") assert code == 1 and "at least 8 registers" in err +def test_o2_tiles_into_shared_memory_and_o0_does_not(capsys): + o2 = run(capsys, EX / "mm_relu.tiny", "--target", "gpu", "--emit", "kernel", "-O", "2")[1] + assert "shared=8 words" in o2 and "syncb" in o2 and "shld" in o2 and "for " not in o2 + o0 = run(capsys, EX / "mm_relu.tiny", "--target", "gpu", "--emit", "kernel", "-O", "0")[1] + assert "shared" not in o0 and "syncb" not in o0 and "for v1 in 0..8" in o0 + + +def test_kernel_passes_can_be_disabled_by_name(capsys): + args = [EX / "mm_relu.tiny", "--target", "gpu", "--emit", "kernel", "-O", "2"] + no_tile = run(capsys, *args, "--disable", "shared-tile")[1] + assert "shared" not in no_tile and "for " not in no_tile # still unrolled + no_unroll = run(capsys, *args, "--disable", "unroll")[1] + assert "shared=8 words" in no_unroll and "for v1 in 0..8" in no_unroll # still tiled + + +def test_o2_run_is_bit_exact_and_reports_shared_memory_use(capsys): + code, out, _ = run(capsys, EX / "mlp_gpu.tiny", "--target", "gpu", "--run", "-O", "2") + assert code == 0 and "bit-exact vs numpy oracle: yes" in out + assert "shared memory: loads" in out and "barrier executions (SYNCB)" in out + + 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: diff --git a/tinyc/cli.py b/tinyc/cli.py index 741b7fa..fa66b25 100644 --- a/tinyc/cli.py +++ b/tinyc/cli.py @@ -15,6 +15,7 @@ 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.gpu.kernel_passes import KERNEL_PASSES from tinyc.ir import ref from tinyc.ir.errors import PassError from tinyc.ir.printer import print_func @@ -35,10 +36,12 @@ def main(argv: list[str] | None = None) -> int: 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 backend") ap.add_argument( - "-O", dest="opt", type=int, choices=[0, 1], default=1, help="optimization level" - ) + "-O", dest="opt", type=int, choices=[0, 1, 2], default=1, + help="0: none; 1: graph passes + kernel licm/strength-reduce; " + "2: also shared-memory tiling and unrolling (gpu)", + ) # fmt: skip ap.add_argument("--disable", action="append", default=[], metavar="PASS", - help=f"skip a pass ({', '.join(REGISTRY)})") # fmt: skip + help=f"skip a pass ({', '.join([*REGISTRY, *KERNEL_PASSES])})") # fmt: skip ap.add_argument("--print-after-all", action="store_true", help="dump the IR after every pass") ap.add_argument("--print-after", action="append", default=[], metavar="PASS") ap.add_argument("--verify-each", action="store_true", @@ -155,6 +158,12 @@ def _gpu_output(args, comp) -> int: 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 + if est["shared_loads"] or est["shared_stores"]: + print( + f"shared memory: loads {o['SHLD']} (estimate {est['shared_loads']}) " + f"stores {o['SHST']} (estimate {est['shared_stores']}) " + f"barrier executions (SYNCB): {o['SYNCB']}" + ) 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: diff --git a/tinyc/driver.py b/tinyc/driver.py index 2c2bf03..4e1c5fe 100644 --- a/tinyc/driver.py +++ b/tinyc/driver.py @@ -7,6 +7,7 @@ from tinyc.frontend import compile_to_graph from tinyc.gpu import GpuProgram, compile_gpu_program +from tinyc.gpu.kernel_passes import KERNEL_PASSES from tinyc.ir.graph import Func from tinyc.passes import REGISTRY, Fixpoint, PassContext, PassManager, PassReport from tinyc.tpu import compile_tile_ir @@ -35,9 +36,10 @@ def build_pipeline( target: str | None = None, opt_level: int = 1, disabled: Iterable[str] = () ) -> list: off = set(disabled) - unknown = off - set(REGISTRY) + unknown = off - set(REGISTRY) - set(KERNEL_PASSES) # kernel passes are disabled the same way if unknown: - raise ValueError(f"unknown passes: {sorted(unknown)} (known: {sorted(REGISTRY)})") + known = sorted(set(REGISTRY) | set(KERNEL_PASSES)) + raise ValueError(f"unknown passes: {sorted(unknown)} (known: {known})") out: list = [] group: list = [] for n in pipeline_names(target, opt_level): @@ -128,5 +130,7 @@ def compile_gpu( 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) + program = compile_gpu_program( + func, ctx, opt_level=opt_level, disabled=disabled, verify_each=verify_each, dump=dump + ) return GpuCompilation(func, program, report) diff --git a/tinyc/gpu/__init__.py b/tinyc/gpu/__init__.py index 79721ad..5b8760d 100644 --- a/tinyc/gpu/__init__.py +++ b/tinyc/gpu/__init__.py @@ -1,8 +1,8 @@ -"""tiny-gpu backend: Graph IR -> memory plan -> Kernel IR -> registers -> `.tgpu` + launch JSON.""" +"""tiny-gpu backend: Graph IR -> memory plan -> Kernel IR -> kernel passes -> registers -> `.tgpu`.""" from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass, field import numpy as np @@ -10,6 +10,7 @@ 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.kernel_passes import KERNEL_PASSES, kernel_pass_names from tinyc.gpu.lower import lower from tinyc.gpu.memplan import MemPlan, check_plan, plan_memory from tinyc.gpu.regalloc import RegisterPressure, allocate @@ -28,16 +29,26 @@ class GpuProgram: 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 + _estimate: dict | None = field(default=None, repr=False) @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) + """Static cost bounds. Cached once the program is fully compiled (registers known).""" + if self._estimate is not None: + return self._estimate + return cost.estimate( + self.kernels, {k.name: self.registers.get(k.name, 0) for k in self.kernels} + ) @property def launch(self) -> dict: return to_launch_json(self) + @property + def launch_core(self) -> dict: + """The launch document without the (costly to compute) estimate: what runners need.""" + return to_launch_json(self, with_estimate=False) + def asm_text(self) -> str: return "\n".join(self.asm[k.name] for k in self.kernels) @@ -53,17 +64,36 @@ def _problem(want, got_outputs, names) -> str | None: return None +def _kernel_problem(prog: GpuProgram, expectations, outputs) -> str | None: + """Interpret the structured kernels on seeded inputs and compare with the oracle.""" + for inputs, want in expectations: + try: + got = interp.run_program(prog, inputs)["outputs"] + except (RangeError, interp.GpuFault) as e: + return f"the kernels fail where the source program does not: {e}" + if why := _problem(want, got, outputs): + return why + return None + + def compile_gpu_program( func: Func, ctx: PassContext | None = None, *, + opt_level: int = 0, + disabled: Iterable[str] = (), 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.""" + """Lower a gpu-legalized Graph IR function, optimize its kernels, allocate registers, emit. + + Kernel optimization levels (see `kernel_passes`): -O0 none, -O1 licm + strength-reduce, + -O2 adds shared-memory tiling, full unrolling and constant folding. + """ ctx = ctx or PassContext() + off = set(disabled) expectations = ref.seeded_expectations(func, num_checks, seed) if verify_each else [] plan = plan_memory(func) check_plan(func, plan) @@ -78,22 +108,35 @@ def compile_gpu_program( 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) + if why := _kernel_problem(prog, expectations, outputs): + raise PassVerifyError("lower-to-kernel", why, "", text) + + for name in kernel_pass_names(opt_level): + if name in off: + continue + before = text + for k in prog.kernels: + KERNEL_PASSES[name]().run(k) + text = _kernel_text(prog.kernels) + try: + for k in prog.kernels: + verify_kernel(k) + except KernelError as e: + raise PassVerifyError(name, str(e), before, text) from None + if dump: + dump(name, text) + if why := _kernel_problem(prog, expectations, outputs): + raise PassVerifyError(name, why, before, 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 + prog._estimate = cost.estimate(prog.kernels, prog.registers) asm = prog.asm_text() if dump: dump("regalloc-emit", asm) @@ -103,7 +146,7 @@ def compile_gpu_program( for inputs, want in expectations: try: - res = run_launch(prog.launch, inputs) + res = run_launch(prog.launch_core, inputs) except (RangeError, OverflowError) as e: raise PassVerifyError( "regalloc-emit", f"simulator failed: {e}", text, asm diff --git a/tinyc/gpu/cost.py b/tinyc/gpu/cost.py index c43e41b..51907a8 100644 --- a/tinyc/gpu/cost.py +++ b/tinyc/gpu/cost.py @@ -1,63 +1,172 @@ -"""Static cost bounds for kernels, to be compared with what the simulator measures. +"""Static cost 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. +Each thread is evaluated *abstractly*: values that depend only on the block id, thread id and +constants are computed exactly, and anything loaded from memory is UNKNOWN. So loop trip counts, +the cooperative-copy guards of a tiled kernel and every load, store and shared access are exact. +The only uncertainty is an `if` on loaded data (relu): a thread that clamps runs one more +instruction, a thread that does not takes a skip branch, which adds a cycle. Instructions and +cycles are therefore bounds `min`/`max`; `gpu.interp` gives the exact numbers for given inputs. + +A barrier is handled per block with the arrival/release model described in `gpu.interp`. """ from __future__ import annotations from dataclasses import dataclass -from tinyc.gpu.kernel_ir import Assign, For, Kernel, Load, Store +from tinyc.gpu.kernel_ir import ( + Assign, + Barrier, + For, + If, + Kernel, + Load, + SharedLoad, + SharedStore, + Special, + Store, +) + +UNKNOWN = None # a value that depends on loaded data @dataclass -class _Cost: +class _Tally: imin: int = 0 imax: int = 0 - tmin: int = 0 # taken branches per thread + tmin: int = 0 # taken branches tmax: int = 0 loads: int = 0 stores: int = 0 + shld: int = 0 + shst: 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 _value(x, env): + if isinstance(x, Special): + return env[x.name] + return x if isinstance(x, int) else env.get(x) -def _cost(body: list) -> _Cost: - total = _Cost() +def _dsts(body): for s in body: - if isinstance(s, (Assign, Load, Store)): - total = total.seq(_Cost(1, 1, 0, 0, isinstance(s, Load), isinstance(s, Store))) + if isinstance(s, (Assign, Load, SharedLoad)): + yield s.dst + if isinstance(s, (For, If)): + yield from _dsts(s.body) + + +def _eval(body: list, env: dict, t: _Tally) -> None: + for s in body: + if isinstance(s, Assign): + t.imin += 1 + t.imax += 1 + a = [_value(x, env) for x in s.args] + if any(v is UNKNOWN for v in a): + env[s.dst] = UNKNOWN + else: + env[s.dst] = ( + a[0] if s.op == "set" else a[0] + a[1] if s.op == "add" else a[0] * a[1] + ) + elif isinstance(s, (Load, Store, SharedLoad, SharedStore)): + t.imin += 1 + t.imax += 1 + if isinstance(s, Load): + t.loads += 1 + elif isinstance(s, Store): + t.stores += 1 + elif isinstance(s, SharedLoad): + t.shld += 1 + else: + t.shst += 1 + if isinstance(s, (Load, SharedLoad)): + env[s.dst] = UNKNOWN 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 + t.imin += 1 # SET var, start + t.imax += 1 + env[s.var] = s.start + trips = s.stop - s.start + for it in range(trips): + _eval(s.body, env, t) + env[s.var] += 1 + t.imin += 3 # ADD, CMP, BRLT + t.imax += 3 + if it < trips - 1: + t.tmin += 1 + t.tmax += 1 + elif isinstance(s, If): + t.imin += 2 # CMP, BRGT + t.imax += 2 + a, b = _value(s.a, env), _value(s.b, env) + if a is not UNKNOWN and b is not UNKNOWN: + if a <= b: + _eval(s.body, env, t) + else: + t.tmin += 1 + t.tmax += 1 + else: # decided by loaded data: bound both outcomes + sub = _Tally() + _eval(s.body, dict(env), sub) + if sub.loads or sub.stores or sub.shld or sub.shst: + raise NotImplementedError("memory access under a data-dependent `if`") + t.imax += sub.imax + t.tmin += min(1, sub.tmin) + t.tmax += max(1, sub.tmax) + for d in _dsts(s.body): + env[d] = UNKNOWN + elif isinstance(s, Barrier): + raise TypeError("barriers are handled by kernel_cost, at the top level only") def kernel_cost(k: Kernel, registers: int) -> dict: - c, t = _cost(k.body), k.threads + blocks, tpb = k.grid + split = next((i for i, s in enumerate(k.body) if isinstance(s, Barrier)), None) + before = k.body if split is None else k.body[:split] + after = [] if split is None else k.body[split + 1 :] + + imin = imax = loads = stores = shld = shst = 0 + cyc_lo = cyc_hi = 0 + for block in range(blocks): + p1s, p2s = [], [] + for thread in range(tpb): + env = {"block": block, "thread": thread, "tid": block * tpb + thread} + p1, p2 = _Tally(), _Tally() + _eval(before, env, p1) + _eval(after, env, p2) + p1s.append(p1) + p2s.append(p2) + if split is None: + lo = max(p.tmin for p in p1s) + 1 + hi = max(p.tmax for p in p1s) + 1 + for p in p1s: + imin += p.imin + imax += p.imax + else: + smin = [p.tmin + 1 for p in p1s] # arrival step bounds + smax = [p.tmax + 1 for p in p1s] + rmin, rmax = max(smin), max(smax) # barrier release step bounds + lo = rmin + max(p.tmin for p in p2s) + 1 + hi = rmax + max(p.tmax for p in p2s) + 1 + for p, q, a_lo, a_hi in zip(p1s, p2s, smin, smax): + imin += p.imin + q.imin + max(1, rmin - a_hi + 1) # SYNCB re-executions + imax += p.imax + q.imax + (rmax - a_lo + 1) + cyc_lo, cyc_hi = max(cyc_lo, lo), max(cyc_hi, hi) + for p in (*p1s, *p2s): + loads += p.loads + stores += p.stores + shld += p.shld + shst += p.shst 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, + "threads": k.threads, + "shared_size": k.shared_size, + "instructions": {"min": imin, "max": imax}, + "cycles": {"min": cyc_lo, "max": cyc_hi}, + "loads": loads, + "stores": stores, + "shared_loads": shld, + "shared_stores": shst, "registers": registers, } @@ -77,6 +186,8 @@ def estimate(kernels: list[Kernel], registers: dict[str, int]) -> dict: }, "loads": sum(p["loads"] for p in per), "stores": sum(p["stores"] for p in per), + "shared_loads": sum(p["shared_loads"] for p in per), + "shared_stores": sum(p["shared_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 index 9884bbe..a370b22 100644 --- a/tinyc/gpu/emit.py +++ b/tinyc/gpu/emit.py @@ -17,8 +17,10 @@ def _inst_text(i: VInst, asg) -> str: 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 == "SYNCB": + return "SYNCB" + if i.op in ("ST", "SHST"): + return f"{i.op} {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) @@ -27,10 +29,11 @@ def _inst_text(i: VInst, asg) -> str: 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 "?" + shared = f"; {kernel.shared_size} shared words per block" if kernel.shared_size else "" lines = [ f"; kernel {kernel.name}: %{kernel.value} = {kernel.kind} ({where})", ( - f"; grid: {kernel.grid[0]} blocks x {kernel.grid[1]} threads;" + f"; grid: {kernel.grid[0]} blocks x {kernel.grid[1]} threads{shared};" f" {alloc.registers_used} registers (R5-R7 reserved, read-only)" ), ] @@ -44,7 +47,7 @@ def emit_asm(kernel: Kernel, insts: list[VInst], alloc: Allocation) -> str: return "\n".join(lines) + "\n" -def to_launch_json(prog) -> dict: +def to_launch_json(prog, with_estimate: bool = True) -> dict: """`prog` is a tinyc.gpu.GpuProgram.""" buffers = {} for name, r in prog.plan.buffers.items(): @@ -67,12 +70,12 @@ def to_launch_json(prog) -> dict: "value": k.value, "op": k.kind, "grid": list(k.grid), - "shared_size": 0, + "shared_size": k.shared_size, "asm": prog.asm.get(k.name, ""), # empty before assembly (validation of the IR) } for k in prog.kernels ], - "estimate": prog.estimate, + "estimate": prog.estimate if with_estimate else {}, } diff --git a/tinyc/gpu/interp.py b/tinyc/gpu/interp.py index 2961d28..c8f0cd2 100644 --- a/tinyc/gpu/interp.py +++ b/tinyc/gpu/interp.py @@ -8,6 +8,13 @@ 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. + +Barriers (`SYNCB`) are modelled from tiny-gpu's source: a thread that reaches the barrier stops +for that step, and a waiting thread executes the `SYNCB` instruction *again on every following +step* until every active thread of its block is waiting; then the block is released and resumes +on the next step. So with arrival step s_t = (taken branches before the barrier) + 1 and release +step R = max s_t over the block, a thread executes `SYNCB` (R - s_t + 1) times and finishes at +step R + (taken branches after the barrier) + 1. """ from __future__ import annotations @@ -18,19 +25,31 @@ import numpy as np -from tinyc.gpu.kernel_ir import I32_HI, I32_LO, SPECIAL_REGS, Assign, For, Load, Special, Store +from tinyc.gpu.kernel_ir import ( + I32_HI, + I32_LO, + Assign, + Barrier, + For, + Kernel, + Load, + SharedLoad, + SharedStore, + Special, + Store, +) from tinyc.ir.errors import InputError, RangeError class GpuFault(Exception): - """A thread touched memory outside [0, mem_size).""" + """A thread touched memory outside its bounds (the simulator would silently misbehave).""" @dataclass class KernelCounts: name: str threads: int - cycles: int # max over threads of (taken branches + 1) + cycles: int # max over threads of the step on which the thread finishes ops: Counter = field(default_factory=Counter) # tiny-gpu op name -> executions, all threads @property @@ -38,50 +57,60 @@ 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.""" +def _run_thread( + kernel: Kernel, body: list, memory: list, shared: list, env: dict, ops: Counter +) -> int: + """Execute `body` for 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 + return x if isinstance(x, int) else env[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)})") + def addr(a: int, size: int, space: str) -> int: + if not 0 <= a < size: + raise GpuFault(f"kernel @{kernel.name}: {space} address {a} outside [0, {size})") return a - def walk(body): + def walk(stmts): nonlocal taken - for s in body: + for s in stmts: 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 - ) + v = a[0] if s.op == "set" else a[0] + a[1] if s.op == "add" else a[0] * a[1] + env[s.dst] = fit(v, s.op) elif isinstance(s, Load): ops["LD"] += 1 - env[s.dst] = memory[check_addr(val(s.addr))] + env[s.dst] = memory[addr(val(s.addr), len(memory), "global")] elif isinstance(s, Store): ops["ST"] += 1 - memory[check_addr(val(s.addr))] = fit(val(s.src), "st") + memory[addr(val(s.addr), len(memory), "global")] = fit(val(s.src), "st") + elif isinstance(s, SharedLoad): + ops["SHLD"] += 1 + env[s.dst] = shared[addr(val(s.saddr), len(shared), "shared")] + elif isinstance(s, SharedStore): + ops["SHST"] += 1 + shared[addr(val(s.saddr), len(shared), "shared")] = fit(val(s.src), "shst") + elif isinstance(s, Barrier): + raise TypeError("barriers are handled by run_block, at the top level only") elif isinstance(s, For): ops["SET"] += 1 env[s.var] = s.start - for it in range(s.stop - s.start): + trips = s.stop - s.start + for it in range(trips): 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: + if it < trips - 1: taken += 1 # the back edge is taken on every iteration but the last else: ops["CMP"] += 1 @@ -91,10 +120,31 @@ def walk(body): else: taken += 1 # skip branch taken - walk(kernel.body) + walk(body) return taken +def run_block(kernel: Kernel, memory: list, block: int, ops: Counter) -> int: + """Run every thread of one block; returns the step on which the block's last thread finishes.""" + tpb = kernel.grid[1] + shared = [0] * kernel.shared_size + split = next((i for i, s in enumerate(kernel.body) if isinstance(s, Barrier)), None) + before = kernel.body if split is None else kernel.body[:split] + after = [] if split is None else kernel.body[split + 1 :] + assert not any(isinstance(s, Barrier) for s in after), "at most one barrier per kernel" + + envs = [{"block": block, "thread": t, "tid": block * tpb + t} for t in range(tpb)] + taken1 = [_run_thread(kernel, before, memory, shared, env, ops) for env in envs] + if split is None: + return max(taken1) + 1 + arrival = [t + 1 for t in taken1] + release = max(arrival) + for a in arrival: + ops["SYNCB"] += release - a + 1 # waiting threads re-execute the barrier every step + taken2 = [_run_thread(kernel, after, memory, shared, env, ops) for env in envs] + return max(release + t + 1 for t in taken2) + + 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"] @@ -127,18 +177,11 @@ def read_outputs(launch: Mapping, memory) -> dict[str, np.ndarray]: 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) + launch = prog.launch_core # buffers and memory size only: no need for the cost estimate + memory = initial_memory(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"} + worst = max(run_block(k, memory, block, ops) for block in range(k.grid[0])) + results.append(KernelCounts(k.name, k.threads, worst, ops)) + return {"outputs": read_outputs(launch, memory), "kernels": results} diff --git a/tinyc/gpu/kernel_ir.py b/tinyc/gpu/kernel_ir.py index 3821e86..6770f65 100644 --- a/tinyc/gpu/kernel_ir.py +++ b/tinyc/gpu/kernel_ir.py @@ -62,6 +62,34 @@ class Store: loc: Loc | None = None +@dataclass +class SharedLoad: + """dst = shared[block][saddr]: the thread's own block's shared memory (tiny-gpu SHLD).""" + + dst: VReg + saddr: Operand + loc: Loc | None = None + + +@dataclass +class SharedStore: + """shared[block][saddr] = src (tiny-gpu SHST).""" + + saddr: Operand + src: Operand + loc: Loc | None = None + + +@dataclass +class Barrier: + """Block-wide barrier (tiny-gpu SYNCB): every thread of the block waits until all arrive. + + Only allowed at the top level of a kernel, so that every thread reaches it. + """ + + loc: Loc | None = None + + @dataclass class For: """`var` counts start, start+1, ... while < stop. Always runs at least once (stop > start).""" @@ -83,7 +111,7 @@ class If: loc: Loc | None = None -Stmt = Assign | Load | Store | For | If +Stmt = Assign | Load | Store | SharedLoad | SharedStore | Barrier | For | If @dataclass @@ -94,6 +122,7 @@ class Kernel: value: str # name of the Graph IR value this kernel computes kind: str # the Graph op: "matmul" | "add" | "relu" loc: Loc | None = None + shared_size: int = 0 # words of shared memory per block (SHLD/SHST addresses 0..size-1) @property def threads(self) -> int: @@ -138,10 +167,31 @@ def operand(x: Operand, defined: set[VReg], where: str) -> None: else: problems.append(f"{where}: bad operand {x!r}") - def walk(body: list, defined: set[VReg]) -> None: + uses_shared = False + + def saddr(x: Operand, defined: set[VReg], where: str) -> None: + operand(x, defined, where) + if isinstance(x, int) and not 0 <= x < max(k.shared_size, 0): + problems.append(f"{where}: shared address {x} outside 0..{k.shared_size - 1}") + + def walk(body: list, defined: set[VReg], top: bool = False) -> None: + nonlocal uses_shared for s in body: where = type(s).__name__.lower() - if isinstance(s, Assign): + if isinstance(s, SharedLoad): + uses_shared = True + saddr(s.saddr, defined, "shld") + defined.add(s.dst) + elif isinstance(s, SharedStore): + uses_shared = True + saddr(s.saddr, defined, "shst") + operand(s.src, defined, "shst") + elif isinstance(s, Barrier): + if not top: + problems.append( + "barrier inside a loop or `if`: threads of a block could miss it" + ) + elif 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") @@ -166,7 +216,11 @@ def walk(body: list, defined: set[VReg]) -> None: else: problems.append(f"unknown statement {s!r}") - walk(k.body, set()) + walk(k.body, set(), top=True) + if uses_shared and k.shared_size < 1: + problems.append("shared memory used but shared_size is 0") + if k.shared_size < 0: + problems.append(f"negative shared_size {k.shared_size}") 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)") @@ -190,6 +244,12 @@ def _print_body(body: list, indent: int, out: list[str]) -> None: 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, SharedLoad): + out.append(f"{pad}{s.dst} = shld [{s.saddr}]") + elif isinstance(s, SharedStore): + out.append(f"{pad}shst [{s.saddr}], {s.src}") + elif isinstance(s, Barrier): + out.append(f"{pad}syncb") elif isinstance(s, For): out.append(f"{pad}for {s.var} in {s.start}..{s.stop} {{") _print_body(s.body, indent + 1, out) @@ -202,8 +262,9 @@ def _print_body(body: list, indent: int, out: list[str]) -> None: def print_kernel(k: Kernel) -> str: where = f" at {k.loc.line}:{k.loc.col}" if k.loc else "" + shared = f" shared={k.shared_size} words" if k.shared_size else "" head = ( - f"kernel @{k.name} grid=({k.grid[0]} blocks, {k.grid[1]} threads)" + f"kernel @{k.name} grid=({k.grid[0]} blocks, {k.grid[1]} threads){shared}" f" // %{k.value} = {k.kind}{where}" ) out = [head + " {"] @@ -219,7 +280,7 @@ def print_kernel(k: Kernel) -> str: 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 + op: str # SET ADD MUL LD ST SHLD SHST SYNCB CMP BRLT BRGT LABEL dst: VReg | None = None srcs: tuple[Operand, ...] = () label: str | None = None # LABEL: the name defined; BR*: the target @@ -244,6 +305,12 @@ def walk(body: list) -> None: 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, SharedLoad): + out.append(VInst("SHLD", s.dst, (s.saddr,), loc=s.loc)) + elif isinstance(s, SharedStore): + out.append(VInst("SHST", None, (s.saddr, s.src), loc=s.loc)) + elif isinstance(s, Barrier): + out.append(VInst("SYNCB", loc=s.loc)) elif isinstance(s, For): top = f"L{next(labels)}" out.append(VInst("SET", s.var, (s.start,), loc=s.loc)) diff --git a/tinyc/gpu/kernel_passes.py b/tinyc/gpu/kernel_passes.py new file mode 100644 index 0000000..aa2f57b --- /dev/null +++ b/tinyc/gpu/kernel_passes.py @@ -0,0 +1,606 @@ +"""Optimizations on the Kernel IR (the straightforward -O0 code is the baseline they improve). + + -O0 nothing + -O1 licm, strength-reduce, kernel-dce + -O2 licm, strength-reduce, shared-tile, unroll, fold, kernel-dce + +Soundness rests on one convention: the kernels `lower.py` builds assign every temporary exactly +once. Only accumulators, induction variables and loop counters are assigned repeatedly, so a pass +may move or rewrite a single-assignment value freely, and must treat a multi-assignment one as +changing over time. Every pass keeps `verify_kernel` happy, and `--verify-each` re-runs the +kernels against the numpy oracle after each one. +""" + +from __future__ import annotations + +import copy +from collections import Counter +from collections.abc import Callable + +from tinyc.gpu.kernel_ir import ( + Assign, + Barrier, + For, + If, + Kernel, + Load, + SharedLoad, + SharedStore, + Special, + Store, + VReg, +) + +I32_LO, I32_HI = -(2**31), 2**31 - 1 +MAX_UNROLL = 16 # full unroll when the trip count is at most this (plan §4.3) + + +# --- helpers ---------------------------------------------------------------------------------- + + +def _walk(body: list): + """Every statement, depth first.""" + for s in body: + yield s + if isinstance(s, (For, If)): + yield from _walk(s.body) + + +def _dst(s): + return s.dst if isinstance(s, (Assign, Load, SharedLoad)) else None + + +def _reads(s) -> list: + """Operands a statement reads (not its destination).""" + if isinstance(s, Assign): + return list(s.args) + if isinstance(s, Load): + return [s.addr] + if isinstance(s, Store): + return [s.addr, s.src] + if isinstance(s, SharedLoad): + return [s.saddr] + if isinstance(s, SharedStore): + return [s.saddr, s.src] + if isinstance(s, If): + return [s.a, s.b] + return [] + + +def _def_counts(kernel: Kernel) -> Counter: + """How many statements assign each virtual register. A For counts twice (var is updated).""" + counts: Counter = Counter() + for s in _walk(kernel.body): + if _dst(s) is not None: + counts[_dst(s)] += 1 + elif isinstance(s, For): + counts[s.var] += 2 + return counts + + +def _max_vreg(kernel: Kernel) -> int: + top = -1 + for s in _walk(kernel.body): + for x in [*_reads(s), _dst(s), getattr(s, "var", None)]: + if isinstance(x, VReg): + top = max(top, x.id) + return top + + +class _Fresh: + def __init__(self, kernel: Kernel): + self.next = _max_vreg(kernel) + 1 + + def __call__(self) -> VReg: + v = VReg(self.next) + self.next += 1 + return v + + +def _inside_defs(loop: For) -> set[VReg]: + out = {loop.var} + for s in _walk(loop.body): + if _dst(s) is not None: + out.add(_dst(s)) + elif isinstance(s, For): + out.add(s.var) + return out + + +def _invariant(x, inside: set[VReg]) -> bool: + return not isinstance(x, VReg) or x not in inside + + +def _uses_outside(kernel: Kernel, loop: For, v: VReg) -> int: + """Reads of `v` in statements that are not part of `loop`.""" + + def count(body) -> int: + n = 0 + for s in body: + if s is loop: + continue + n += sum(1 for x in _reads(s) if x == v) + if isinstance(s, (For, If)): + n += count(s.body) + return n + + return count(kernel.body) + + +def _rewrite_reads(s, fn: Callable) -> None: + """Replace every operand a statement reads with fn(operand), in place (recursing).""" + if isinstance(s, Assign): + s.args = tuple(fn(a) for a in s.args) + elif isinstance(s, Load): + s.addr = fn(s.addr) + elif isinstance(s, Store): + s.addr, s.src = fn(s.addr), fn(s.src) + elif isinstance(s, SharedLoad): + s.saddr = fn(s.saddr) + elif isinstance(s, SharedStore): + s.saddr, s.src = fn(s.saddr), fn(s.src) + elif isinstance(s, If): + s.a, s.b = fn(s.a), fn(s.b) + if isinstance(s, (For, If)): + for c in s.body: + _rewrite_reads(c, fn) + + +def _apply(op: str, a: int, b: int | None = None) -> int: + return a if op == "set" else a + b if op == "add" else a * b + + +# --- loop-invariant code motion --------------------------------------------------------------- + + +class LICM: + """Hoist single-assignment arithmetic whose operands do not change inside a loop. + + Loads are never hoisted (no alias analysis). Loops run at least once, so hoisting an + assignment that would have run anyway never adds an overflow or a fault. + """ + + name = "licm" + + def run(self, kernel: Kernel) -> None: + defc = _def_counts(kernel) + + def go(body: list) -> list: + out: list = [] + for s in body: + if isinstance(s, For): + s.body = go(s.body) # inner loops first: their hoists can climb further + inside = _inside_defs(s) + hoisted, kept = [], [] + for t in s.body: + if ( + isinstance(t, Assign) + and defc[t.dst] == 1 + and all(_invariant(a, inside) for a in t.args) + ): + hoisted.append(t) + inside.discard(t.dst) + else: + kept.append(t) + s.body = kept + out.extend(hoisted) + elif isinstance(s, If): + s.body = go(s.body) + out.append(s) + return out + + kernel.body = go(kernel.body) + + +# --- strength reduction ------------------------------------------------------------------------ + + +def _fold_add(a, b, pre: list, fresh: _Fresh, loc=None): + if isinstance(a, int) and isinstance(b, int): + return a + b + if isinstance(a, int) and a == 0: + return b + if isinstance(b, int) and b == 0: + return a + v = fresh() + pre.append(Assign("add", v, (a, b), loc)) + return v + + +def _fold_mul(a, c: int, pre: list, fresh: _Fresh, loc=None): + if isinstance(a, int): + return a * c + if c == 0: + return 0 + if c == 1: + return a + v = fresh() + pre.append(Assign("mul", v, (a, c), loc)) + return v + + +class StrengthReduce: + """Replace a multiply/add chain on a loop counter by one pointer bumped each iteration. + + t4 = mul k, N ; t5 = add t4, thread ; t6 = add t5, base ; x = ld [t6] + becomes, with `t6 = add thread, base` computed once before the loop, + x = ld [t6] ; ... ; t6 = add t6, N + + The chain's value is linear in k, so it can be advanced by its step instead of recomputed. + Only chains of two or more statements are reduced (a single add costs the same as the + increment that would replace it), only the last value of a chain is turned into a pointer, + and only when it is not read outside the loop. The dead rest of the chain is removed by + `kernel-dce`. + """ + + name = "strength-reduce" + + def run(self, kernel: Kernel) -> None: + defc = _def_counts(kernel) + fresh = _Fresh(kernel) + + def go(body: list) -> list: + out: list = [] + for s in body: + if isinstance(s, For): + s.body = go(s.body) + out.extend(self._reduce(kernel, s, defc, fresh)) + elif isinstance(s, If): + s.body = go(s.body) + out.append(s) + return out + + kernel.body = go(kernel.body) + + def _reduce(self, kernel: Kernel, loop: For, defc: Counter, fresh: _Fresh) -> list: + inside = _inside_defs(loop) + # vreg -> (step per iteration, parent IV, defining statement, the other operand) + iv: dict[VReg, tuple] = {loop.var: (1, None, None, None)} + for s in loop.body: + if not isinstance(s, Assign) or defc[s.dst] != 1 or s.op == "set": + continue + a, b = s.args + for x, y in ((a, b), (b, a)): + if not (isinstance(x, VReg) and x in iv): + continue + if s.op == "mul" and isinstance(y, int): + iv[s.dst] = (iv[x][0] * y, x, s, y) + break + if s.op == "add" and _invariant(y, inside): + iv[s.dst] = (iv[x][0], x, s, y) + break + + def chain_len(v: VReg) -> int: + return 0 if v == loop.var else 1 + chain_len(iv[v][1]) + + parents = {p for (_, p, _, _) in iv.values() if p is not None} + pre: list = [] + for d in [v for v in iv if v != loop.var]: + if chain_len(d) < 2 or d in parents or _uses_outside(kernel, loop, d): + continue + if not any(d in _reads(t) for t in _walk(loop.body)): + continue + step = iv[d][0] + + def init(v: VReg): + if v == loop.var: + return loop.start + _, parent, stmt, other = iv[v] + p = init(parent) + if stmt.op == "mul": + return _fold_mul(p, other, pre, fresh, stmt.loc) + return _fold_add(p, other, pre, fresh, stmt.loc) + + before = len(pre) + value = init(d) + if len(pre) > before and pre[-1].dst == value: + pre[-1].dst = d # the last computation of the init chain defines the pointer + elif isinstance(value, int): + pre.append(Assign("set", d, (value,), iv[d][2].loc)) + else: + pre.append(Assign("add", d, (value, 0), iv[d][2].loc)) + loop.body = [t for t in loop.body if t is not iv[d][2]] + loop.body.append(Assign("add", d, (d, step), iv[d][2].loc)) + return pre + + +# --- dead code elimination --------------------------------------------------------------------- + + +class KernelDCE: + """Remove assignments and loads whose value never reaches a store, `if` or loop control.""" + + name = "kernel-dce" + + def run(self, kernel: Kernel) -> None: + live: set[VReg] = set() + for s in _walk(kernel.body): + if isinstance(s, (Store, SharedStore, If)): + live.update(x for x in _reads(s) if isinstance(x, VReg)) + elif isinstance(s, For): + live.add(s.var) + changed = True + while changed: + changed = False + for s in _walk(kernel.body): + d = _dst(s) + if d is not None and d in live: + for x in _reads(s): + if isinstance(x, VReg) and x not in live: + live.add(x) + changed = True + + def prune(body: list) -> list: + out = [] + for s in body: + if _dst(s) is not None and _dst(s) not in live: + continue + if isinstance(s, (For, If)): + s.body = prune(s.body) + out.append(s) + return out + + kernel.body = prune(kernel.body) + + +# --- shared-memory tiling ---------------------------------------------------------------------- + + +def _thread_deps(kernel: Kernel) -> dict[VReg, set[str]]: + """What each vreg may depend on besides `block` and constants: 'thread' or loaded 'mem'.""" + deps: dict[VReg, set[str]] = {} + changed = True + while changed: + changed = False + for s in _walk(kernel.body): + d = _dst(s) + if d is None: + continue + new: set[str] = set() + if isinstance(s, (Load, SharedLoad)): + new.add("mem") + for x in _reads(s): + if isinstance(x, Special) and x.name in ("thread", "tid"): + new.add("thread") + elif isinstance(x, VReg): + new |= deps.get(x, set()) + if not new <= deps.get(d, set()): + deps.setdefault(d, set()).update(new) + changed = True + return deps + + +class SharedTile: + """Load block-uniform data once per block into shared memory instead of once per thread. + + In a loop `for k in 0..K`, a load whose address is `u + k` with `u` the same for every + thread of a block (it depends only on the block id and constants) reads the same K words in + every thread. The block's threads copy those words cooperatively (thread j copies words + j, j+N, ...), meet at a barrier, and the loop then reads `shld [k]`. Global loads drop from + N*K to K per block. Skipped for one-thread blocks, where there is nothing to share. + """ + + name = "shared-tile" + + def run(self, kernel: Kernel) -> None: + n = kernel.grid[1] + if n < 2 or kernel.shared_size: + return + defc = _def_counts(kernel) + deps = _thread_deps(kernel) + fresh = _Fresh(kernel) + uses = Counter() + for s in _walk(kernel.body): + uses.update(x for x in _reads(s) if isinstance(x, VReg)) + + for pos, loop in enumerate(kernel.body): + if not (isinstance(loop, For) and loop.start == 0): + continue + inside = _inside_defs(loop) + addr_def = { + t.dst: t + for t in loop.body + if isinstance(t, Assign) and t.op == "add" and defc[t.dst] == 1 + } + tiles = [] # (load, address statement, u) + for t in loop.body: + if not (isinstance(t, Load) and isinstance(t.addr, VReg)): + continue + a = addr_def.get(t.addr) + if a is None or uses[t.addr] != 1: + continue + for x, u in ((a.args[0], a.args[1]), (a.args[1], a.args[0])): + if ( + x == loop.var + and isinstance(u, VReg) + and u not in inside + and defc[u] == 1 + and not deps.get(u) + ): + tiles.append((t, a, u)) + break + if not tiles: + continue + + k_words, pre, offset = loop.stop, [], 0 + for load, addr_stmt, u in tiles: + for r in range((k_words + n - 1) // n): + idx = Special("thread") if r == 0 else fresh() + if r: + pre.append(Assign("add", idx, (Special("thread"), r * n), load.loc)) + addr, tmp, sidx = fresh(), fresh(), idx + body = [Assign("add", addr, (u, idx), load.loc), Load(tmp, addr, load.loc)] + if offset: + sidx = fresh() + body.append(Assign("add", sidx, (idx, offset), load.loc)) + body.append(SharedStore(sidx, tmp, load.loc)) + if (r + 1) * n <= k_words: # every thread's index is in range: no guard + pre.extend(body) + else: + pre.append(If(idx, k_words - 1, body, load.loc)) + # in the loop: the copied word for iteration k is at shared[offset + k] + replacement = [SharedLoad(load.dst, loop.var, load.loc)] + if offset: + s_k = fresh() + replacement = [ + Assign("add", s_k, (loop.var, offset), load.loc), + SharedLoad(load.dst, s_k, load.loc), + ] + new_body: list = [] + for t in loop.body: # identity, not equality: equal-looking statements may differ + if t is load: + new_body.extend(replacement) + elif t is not addr_stmt: + new_body.append(t) + loop.body = new_body + offset += k_words + pre.append(Barrier(loop.loc)) + kernel.shared_size = offset + kernel.body[pos:pos] = pre + return # one loop per kernel is enough for matmul + + +# --- full unrolling ---------------------------------------------------------------------------- + + +class Unroll: + """Fully unroll innermost loops with at most MAX_UNROLL iterations. + + The counter becomes a constant in every copy, so address arithmetic folds and a loop with + no taken branches runs in a single simulator cycle. Skipped when the counter is read after + the loop. + """ + + name = "unroll" + + def run(self, kernel: Kernel) -> None: + def go(body: list) -> list: + out: list = [] + for s in body: + if isinstance(s, For): + s.body = go(s.body) + trip = s.stop - s.start + nested = any(isinstance(t, For) for t in _walk(s.body)) + if trip <= MAX_UNROLL and not nested and not _uses_outside(kernel, s, s.var): + for i in range(trip): + for t in copy.deepcopy(s.body): + _rewrite_reads( + t, lambda x, v=s.var, c=s.start + i: c if x == v else x + ) + out.append(t) + continue + elif isinstance(s, If): + s.body = go(s.body) + out.append(s) + return out + + kernel.body = go(kernel.body) + + +# --- constant folding -------------------------------------------------------------------------- + + +class Fold: + """Constant propagation and folding on straight-line code (mostly after unrolling). + + * a single-assignment `set t, c` is propagated into its uses; + * `add`/`mul` of two constants becomes a `set`; + * constants are moved to the right-hand operand; + * `add (add x, c1), c2` becomes `add x, c1 + c2` when x cannot change in between + (a constant, a special register, or a single-assignment vreg); + * `d = add a, 0` is a copy: reads of d read `a` instead (same condition on `a`). + Whatever becomes dead is removed by `kernel-dce`. + """ + + name = "fold" + + def run(self, kernel: Kernel) -> None: + for _ in range(20): # each round can enable the next; stops at a fixpoint + if not self._round(kernel): + break + + def _round(self, kernel: Kernel) -> bool: + defc = _def_counts(kernel) + single = { + _dst(s): s for s in _walk(kernel.body) if _dst(s) is not None and defc[_dst(s)] == 1 + } + changed = False + + def const_of(x): + d = single.get(x) if isinstance(x, VReg) else None + return ( + d.args[0] + if isinstance(d, Assign) and d.op == "set" and isinstance(d.args[0], int) + else None + ) + + # `d = add a, 0` is a copy of `a`; reads of d can read a instead (when a cannot change) + copies = { + d: st.args[0] + for d, st in single.items() + if isinstance(st, Assign) + and st.op == "add" + and isinstance(st.args[1], int) + and st.args[1] == 0 + and not isinstance(st.args[0], int) + and (not isinstance(st.args[0], VReg) or st.args[0] in single) + } + + def prop(x): + nonlocal changed + while isinstance(x, VReg) and x in copies: # follow chains of copies + x = copies[x] + changed = True + c = const_of(x) + if c is not None: + changed = True + return c + return x + + for s in _walk(kernel.body): + _rewrite_reads_shallow(s, prop) + if not isinstance(s, Assign) or s.op == "set": + continue + a, b = s.args + if isinstance(a, int) and isinstance(b, int): + v = _apply(s.op, a, b) + if I32_LO <= v <= I32_HI: + s.op, s.args, changed = "set", (v,), True + continue + if isinstance(a, int): # constants go on the right + s.args, changed = (b, a), True + a, b = b, a + if s.op == "add" and isinstance(b, int) and isinstance(a, VReg): + inner = single.get(a) + if ( + isinstance(inner, Assign) + and inner.op == "add" + and isinstance(inner.args[1], int) + and (not isinstance(inner.args[0], VReg) or inner.args[0] in single) + ): + total = b + inner.args[1] + if I32_LO <= total <= I32_HI: + s.args, changed = (inner.args[0], total), True + return changed + + +def _rewrite_reads_shallow(s, fn: Callable) -> None: + """Like _rewrite_reads but without recursing (the caller already walks every statement).""" + if isinstance(s, (For, If)): + if isinstance(s, If): + s.a, s.b = fn(s.a), fn(s.b) + return + _rewrite_reads(s, fn) + + +KERNEL_PASSES = { + p.name: p for p in (LICM, StrengthReduce, KernelDCE, SharedTile, Unroll, Fold) +} # fmt: skip + + +def kernel_pass_names(opt_level: int) -> list[str]: + if opt_level <= 0: + return [] + if opt_level == 1: + return ["licm", "strength-reduce", "kernel-dce"] + return ["licm", "strength-reduce", "shared-tile", "unroll", "fold", "kernel-dce"] diff --git a/tinyc/gpu/lower.py b/tinyc/gpu/lower.py index 7e7213d..64b0167 100644 --- a/tinyc/gpu/lower.py +++ b/tinyc/gpu/lower.py @@ -48,16 +48,20 @@ def _lower_matmul(op: Op, plan: MemPlan, name: str) -> Kernel: 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)) + # Every temporary is assigned exactly once (only acc and k change over time), which is what + # lets the kernel passes move and rewrite them. The adds are ordered so that the part of + # each address that does not depend on k comes first: &A[row][k] = (row*K + a_base) + k and + # &B[k][col] = ((k*N) + col) + b_base. + t_a1, t_a2, t_a3, x, t_b1, t_b2, t_b3, y, p = (b.fresh() for _ in range(9)) 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", t_a1, (BLOCK, k_dim), loc), + Assign("add", t_a2, (t_a1, a_base), loc), + Assign("add", t_a3, (t_a2, k), loc), + Load(x, t_a3, loc), + Assign("mul", t_b1, (k, n_dim), loc), + Assign("add", t_b2, (t_b1, THREAD), loc), + Assign("add", t_b3, (t_b2, b_base), loc), + Load(y, t_b3, loc), Assign("mul", p, (x, y), loc), Assign("add", acc, (acc, p), loc), ] diff --git a/tinyc/gpu/report.py b/tinyc/gpu/report.py new file mode 100644 index 0000000..aff7e32 --- /dev/null +++ b/tinyc/gpu/report.py @@ -0,0 +1,77 @@ +"""The measured optimization table in docs/ARCHITECTURE.md, generated from real simulator runs. + +`python scripts/opt_table.py --write` regenerates it; a test fails if the document drifts. +""" + +from __future__ import annotations + +import numpy as np + +from tinyc.driver import compile_gpu +from tinyc.ir import ref + +PROGRAMS = { + "4x4": "input A : i8[4, 4]\ninput B : i8[4, 4]\nC = relu(A @ B)\noutput C\n", + "8x8": "input A : i8[8, 8]\ninput B : i8[8, 8]\nC = relu(A @ B)\noutput C\n", +} + +# (label, -O level, kernel passes switched off) +CONFIGS = [ + ("-O0 (naive)", 0, ()), + ("+ licm", 1, ("strength-reduce",)), + ("+ strength-reduce (= -O1)", 1, ()), + ("-O2 with shared-memory tiling only", 2, ("unroll", "fold")), + ("-O2 with unrolling + folding only", 2, ("shared-tile",)), + ("-O2 (all passes)", 2, ()), +] +SEED = 42 + + +def measure(source: str, opt: int, disabled=()) -> dict: + from runtime.gpu_runner import run_launch + + comp = compile_gpu(source, "mm_relu.tiny", opt_level=opt, disabled=disabled) + inputs = ref.random_inputs(comp.func, np.random.default_rng(SEED)) + res = run_launch(comp.launch, inputs) + want = ref.evaluate(comp.func, inputs) + exact = all((res.outputs[n] == w).all() for n, w in zip(comp.launch["outputs"], want)) + assert exact and res.reserved_intact, "a measured configuration must be correct" + ops = res.op_counts + return { + "instructions": res.instructions, + "cycles": res.cycles, + "loads": ops["LD"], + "shared": ops["SHLD"], + "registers": comp.program.estimate["registers"], + } + + +def opt_table_markdown() -> str: + out: list[str] = [] + for name, source in PROGRAMS.items(): + rows = [(label, measure(source, opt, off)) for label, opt, off in CONFIGS] + base = rows[0][1] + out.append(f"**`mm_relu` {name}** (`C = relu(A @ B)`, int8, seed {SEED})\n") + out.append( + "| configuration | instructions | cycles | global loads | shared loads | registers |" + ) + out.append("|---|---:|---:|---:|---:|---:|") + for label, m in rows: + cells = [ + ( + f"{m['instructions']} ({m['instructions'] / base['instructions'] - 1:+.0%})" + if m is not base + else str(m["instructions"]) + ), + str(m["cycles"]), + ( + f"{m['loads']} ({m['loads'] / base['loads'] - 1:+.0%})" + if m is not base + else str(m["loads"]) + ), + str(m["shared"]), + str(m["registers"]), + ] + out.append(f"| {label} | " + " | ".join(cells) + " |") + out.append("") + return "\n".join(out)