fuse

A small ML graph compiler: load real ONNX models into an IR, run optimizer passes over them (identity elimination, constant folding, conv+batchnorm fusion, activation fusion, dead-code elimination, liveness-based memory planning), execute both graphs on a from-scratch NumPy runtime, and prove the answers did not move. The victim is resnet18 as PyTorch exports it with the exporter's own cleanup disabled: 141 nodes of ceremony around 20 convolutions.

141 nodes to 32, outputs identical to 1.19e-06. Every Identity gone (all 72 of them), every BatchNorm folded into its convolution's weights (all 20), every eligible Relu fused into its producer (all 17), and the survivors are almost pure compute.

Correct before clever

The runtime (fuse/runtime.py, im2col convolutions, ~120 lines) is verified against onnxruntime before any pass touches anything: max deviation 2.6e-06 on resnet18, 8.2e-08 on an MLP, argmax identical (tests/test_runtime.py). Then the optimizer's contract is tested the same way: the optimized graph must reproduce the raw graph's outputs, and does, to 1.19e-06 (tests/test_passes.py).

The passes, and what each is worth

Identity elimination does the loudest work (141 to 69). Conv+BN fusion does the most valuable: the algebra folds a whole normalization layer into the convolution that feeds it, exactly:

(Animation source: assets/manim_fold.py, video in assets/fold.mp4.)

Where the milliseconds and megabytes went

End to end on my runtime: 201ms to 177ms per inference (1.14x). The honest context is right there in the chart: onnxruntime runs the same graph in 22ms, because its advantage is hand-written kernels, not graph surgery. The BatchNorm and Relu rows collapse to zero after fusion, which is precisely the work the passes deleted; the Conv row barely moves, because the passes cannot delete the actual mathematics.

The memory planner's liveness analysis cuts peak activation memory from 13.8MB to 4.0MB (3.4x) with free-after-last-use, and the live-memory sawtooth shows where the budget goes: the early 112x112 feature maps dwarf everything after the first downsampling stages. On models that do not fit, this pass is the one that decides whether you run at all.

What I learned building it

• PyTorch's exporter had already done my homework once. The default export folds BN and still leaves 16 dead Identity nodes; I re-exported with do_constant_folding=False so my passes had the real, messy graph compilers actually receive.
• Pass order is a correctness feature. Identity elimination must run first or the fusion patterns cannot see through the aliases; dead-code runs last to sweep what fusion orphaned (8 dead initializer tensors here).
• A compiler is only as trustworthy as its differential tests. Every pass is validated by running both graphs and comparing, the same way the serious ones (XLA, TensorRT) are.

Running it

pip install numpy onnx onnxruntime torch torchvision   # torch only to export
python -c "import benchmarks"                          # or: sh data/export.sh
python tests/test_runtime.py     # parity vs onnxruntime first
python tests/test_passes.py      # the optimizer's contract
python benchmarks/bench.py       # timings, attribution, memory plan
python assets/make_visuals.py

Where this sits

Eighth in the from-scratch ML systems series, after ember, forge-lm, tinyserve-engine, quantlab, gradmesh, hive-index, and relay-serve. The compiler layer: the part of the stack (TensorRT, XLA, torch.compile) where graphs go to get faster, rebuilt small enough to read in an afternoon.