# anneal

> A Go tensor compiler. Graph-rewrite-based, with autodiff and optimization as first-class compiler passes over one immutable UOp IR. The scheduler fuses across the forward/backward boundary.

anneal is an open-source project by georgebuilds. It is NOT a simulated-annealing optimizer. The name refers to the metallurgical process. Controlled heating and cooling that fuses grain structure and relaxes a material into a lower-energy, stable form, which is a near-exact metaphor for what this compiler does: it takes the whole computation (forward and backward), applies rewrite passes, and relaxes the graph into a fused, lower-cost form.

## Key claims

- Autodiff is a compiler pass, not a runtime hook. Per-op gradient rules are dispatched through a ruleset (`map[uop.Op]→GradRule`) over the same UOp IR as the forward graph, drift-checked against the curated `explain` documentation. The output is a uniform, fusible backward graph, not a runtime tape. Reverse-mode (`tensor.Backward`) is the training default; forward-mode autodiff (`tensor.JVP`, a Jacobian-vector product) is also shipped as the same kind of compiler pass (the tangent graph is just more UOps on the same arena) and powers the exact-MeanFlow objective's directional time-derivative.
- The scheduler (rangeify) sees forward and backward together and can fuse across that boundary.
- UOps are immutable, interned, and arena-allocated. Structural equality is identity equality.
- No reflection in the rewrite hot path. This is a hard invariant.
- `anneal viz` runs the actual compiler frontend compiled to WebAssembly and renders the real UOp graph in-browser. Not a mock. JSON output identifies which gradient rule produced each backward node.
- `anneal web` is a single-binary local browser studio with eight deep-linkable views. Zero telemetry, zero accounts, model bytes never leave the machine. Every view that compiles runs as WASM in a Web Worker; every view that executes streams over SSE from a native handler.
- ONNX import is shipped via `onnx.Import(bytes, arena, device)`. Zero-CGO (pure-Go protobuf bindings, committed in-tree). About 45 op handlers cover Stage-1 CNN and Stage-2 transformer cores. Symbolic `dim_param` axes ride through as anneal `Variable`s. `WithStructureOnly()` opt skips initializer payloads for visualization.

## Architecture

- **IR:** UOp DAG, immutable, arena-allocated, hash-consed
- **Rewrite engine:** PatternMatcher + graph_rewrite driver
- **Scheduler:** rangeify - fuses across the forward/backward boundary; memoized on a structural key
- **Codegen:** WGSL lowerer + renderer; Opt seam (`codegen/opt.go`) for composable kernel transforms (OptLocal, OptTile, OptUpcast, OptVectorize, OptVec4Load, where OptVec4Load rebinds matmul inputs as `array<vec4<f32>>` for 128-bit Metal tile loads); BEAM autotuning (`codegen/beam.go`) - beam-of-k search over opt sequences, env-gated (ANNEAL_BEAM=1), persistent disk cache at `~/.cache/anneal/beam_cache.json`
- **Backend:** WebGPU is the production target. `backend/` defines the five-interface contract (`Renderer`, `Compiler`, `Allocator`, `Program`, `DeviceBuffer`); `backend/webgpu/` implements them with `executor.go` as orchestrator. Pipeline cache lives on the `Compiler` (static + symbolic sub-caches, both keyed by normalized WGSL). A pure-Go CPU interpreter at `backend/cpu/` is a sibling executor for the zero-GPU case: it implements `backend.Executor` directly by walking the linearized instruction stream, skips Renderer / Compiler / Program entirely (no codegen to Go, no toolchain at runtime, zero-CGO preserved), and acts as the value oracle for the WebGPU path. Slice 1 covers the op set for the MLP forward+backward (forward max-abs-diff vs WebGPU 1.19e-07, gradient 2.98e-08); the remaining UOp coverage is the slice 2 backlog. Select via `--device=cpu` on `anneal train`.
- **Language:** Go (latest stable)
- **Symbolic shapes:** the batch dimension can be symbolic - a kernel compiles once and runs at any batch size via `NewSymbolicBatchInput` / `RealizeWithBinding`. The general-purpose tensor-surface constructor is `tensor.NewSymbolicShape(arena, []shape.Sint, dtype, device)`, fed from `tensor.NewVariable(arena, name, min, max).Sint()` for each symbolic axis; symbolic dims may appear at any position with multiple distinct Variables per shape, and `.Bind(val)` returns a `map[string]int64` entry (combine with `tensor.MergeBindings`). General axis movement is also supported: split/merge a symbolic axis (`[n,4]↔[n*4]`), pad/shrink amounts that are themselves symbolic, multi-dim symbolic dispatch with the symbolic axis in any position on both kernel-output and input buffers, an unbounded number of distinct symbolic vars per kernel (the `ParamsN` uniform grows dynamically), and cross-arena structural-key portability so symbolic graphs survive arena churn. Kernel opts: LOCAL applies on multi-dim symbolic kernels; TILE stays unavailable on symbolic axes (WGSL forbids non-const workgroup sizes, a hard platform ceiling); UPCAST/VECTORIZE/VEC4LOAD are matmul-only by lowerer design (only emitTiledReduce handles their per-lane positions); UPCAST and VECTORIZE are fail-loud at opt-application time when composed without OptTile (panic with diagnostic), and BEAM's ActionSpace pre-filters them on non-tiled kernels so the search never reaches either assertion. OptVec4Load is the throughput lever: composed after OptTile it binds the f32 matmul inputs as `array<vec4<f32>>` and turns tile fills into 128-bit Metal loads (best-known stack `OptLocal²+OptTile+OptUpcast²+OptVec4Load` ≈ 371/420 GFLOP/s @1024³/2048³ on M3). Correctness does not depend on autotuning.
- **JIT:** `tensor.JIT` captures the scheduling plan on the first step and replays subsequent steps without re-scheduling.
- **Migration I/O:** `tensor/npy` loads numpy `.npy`/`.npz` files; `tensor/safetensors` reads and writes HuggingFace `.safetensors` checkpoints, bidirectionally compatible with the real Python safetensors library. Pure Go, no cgo, no runtime Python. Tensor data carries as `float32`; source dtypes map with documented float32-range semantics.
- **Low-precision dtypes:** `tensor.Tensor` and `nn.Parameter` accept `Dtypes.Float16` (storage + arithmetic, narrowing uses IEEE 754 round-to-nearest-even; requires `shader-f16` WebGPU extension; fails closed on unavailability with clear error) and `Dtypes.BFloat16` (storage-only; arithmetic runs in f32 via bitcast/shift conversion at buffer boundaries; narrowing uses round-to-nearest-even per the PyTorch / Eigen reference, emitted as the `_bf16_rtne_bits` prelude helper; available on any WebGPU adapter). Reduction accumulators stay f32 even with f16 operands. `Parameter.Value` is f32 master; `Load` quantizes to device dtype each step (implicit mixed precision). FP16/FP16 fast path is deferred.
- **ONNX importer:** `onnx/` package. Pure-Go protobuf bindings under `onnx/onnxpb/` (committed). Strategy A bit-exact gate: every E2E test builds the model twice on the same arena (direct Tensor API + via importer) and asserts `[]float32` slice equality (max-abs-diff 0). Strategy B onnxruntime cross-check: committed ResNet-9 golden lands at max-abs-diff 8.2e-08. Phase 4 conformance harness over ONNX 1.17.0 backend node corpus: 234 cases committed under `onnx/testdata/node/`, 174 pass, 0 fail, 60 documented skips, max-abs-diff best=0, median=0, worst=7.324e-04 against a 1e-3 tolerance. Two new UOps shipped to close coverage: OpErf (Abramowitz-Stegun 7.1.26 WGSL polynomial, max abs error 1.68e-07) and OpMin. Symbolic `dim_param` ONNX axes ride through as anneal `Variable`s on the Option B path. `WithStructureOnly()` skips initializer payload materialisation for visualization use; `Runner.Run` fails closed in that mode. Documented v1.1 punts include Conv group>1, Resize, control flow, quantization, FLOAT8/STRING/COMPLEX dtypes, MaxPool ceil_mode/dilations/auto_pad, Slice |step|>1, Pow int-typed, Reduce empty-set semantics. The skip list in `onnx/conformance_skip.go` is the documented exclusion contract.
- **`anneal web` local studio:** single-binary local browser surface served by `anneal web [addr]` (default `:3001`, `:0` auto-allocates). Eight deep-linkable views: studio (`/`), visualize (`/v/<model>`), kernels (`/k/<model>`), explain (`/x/<op>`), train (`/t/<model>`), generate (`/g/<model>`), history (`/r/<id>`), doctor (`/d`). Load-bearing axis: WASM/native split. Every view that *compiles* runs as WASM in a Web Worker (visualize, kernels, explain, ONNX dropzone, tensor inspect); every view that *executes* streams over SSE from a native handler (train, generate). The doctor view shows both: native device card + browser `navigator.gpu` probe side by side. ONNX dropzone runs entirely in the browser via `onnx.WithStructureOnly()`; model bytes never reach the server. Bundle persistence under `~/.cache/anneal/runs/<ts>-<model>-<6hex>/` with `bundle_version=1` (manifest.json, schedule.json, kernels/*.wgsl, graph.json, loss.csv, generation.ndjson, events.ndjson, config.json). CLI default OFF (`--bundle` or `ANNEAL_BUNDLE=1` opts in); web default ON (`?bundle=0` disables). GPT-2-small verified on M3: 47s cold-start for 3-token completion, 3.82 GB peak RSS for 10 tokens.
- **Accessibility is binding, not aspirational.** WCAG 2.x AA across every view. Brand tokens carry verified contrast (forward teal #00ADD8 7.14:1 on dark, ember and gold tabulated). matchMedia listener tracks OS theme live. Skip link, semantic landmarks, single h1 per active view, focus-visible global rule, keyboard reachable, chord help via `?`, polite live region, aria-current="page", touch target 44x44, reduced-motion respected, forced-colors fallback maps brand tokens to system color keywords. Twenty foundation a11y tests in `cmd/anneal/cmd_web_test.go` plus a binding per-view checklist in `web/A11Y.md` are merge blockers for new views. Design invariants: DD1 colour is never alone, DD2 real compiler only, DD3 the library is the product, DD4 restraint (no notebook, no model hub, no telemetry, no auth, no third-party JS).

## Color semantics (visualizer and TUI)

- teal (#00ADD8) - forward pass
- ember (#FF7A45) - backward pass
- gold (#F2C57C) - fused kernel

Color is never the sole carrier of meaning: forward = solid line + teal, backward = dashed line + ember, fused = filled box + gold.

## CLI

- `anneal run mlp` - realize and execute a registered example graph
- `anneal train mlp`: training loop with the live TUI (models: mlp, conv, dynmlp, nanogpt, llama, moe, bert, vit, resnet9, gpt2, diffusion, dit, meanflow; use `--batch=N` with dynmlp for dynamic batch). ResNet-9 also trains end to end via `anneal train resnet9` (forward + backward + Adam on CIFAR-10) at a small default batch, since the scalar-WGSL codegen ceiling caps batch size; `anneal run resnet9` realizes the forward graph
- `anneal train nanogpt`: char-level transformer trained end to end on tinyshakespeare; first run fetches the corpus, prints a generated sample at the end
- `anneal train llama`: char-level Llama-style decoder trained from scratch on tinyshakespeare; the modern-decoder level-up over nanoGPT (RMSNorm instead of LayerNorm, grouped-query attention with RoPE instead of vanilla multi-head attention plus learned absolute positions, SwiGLU feed-forward instead of a GELU MLP, tied embeddings); first run fetches the corpus, prints a generated sample at the end. A small architecture demo, not a trained chat model
- `anneal train moe`: a small GPT whose per-block feed-forward is replaced by a router plus several expert FFNs, combined by soft (dense) gating with a load-balance auxiliary loss; a teaching version of sparse Mixture-of-Experts. The only change from nanoGPT is the block feed-forward; trained from scratch on tinyshakespeare, first run fetches the corpus and prints a generated sample at the end
- `anneal train bert`: a char-level BERT encoder trained from scratch on tinyshakespeare with masked language modeling (predict masked tokens from both-side context); encoder-only, the bidirectional contrast to the causal nanoGPT/Llama decoders; first run fetches the corpus and prints a masked-token reconstruction at the end
- `anneal train vit`: Vision Transformer (ViT-tiny: 32x32 RGB, patch 4, 64-dim embed, 2 encoder blocks, 4 heads, mean-pool classification head) trained on a synthetic 10-class image dataset; no asset downloads
- `anneal gpt2 sample <prompt>`: load HuggingFace GPT-2-small weights (~548 MB SHA-pinned safetensors fetched on first run), encode with pure-Go BPE, run forward through the same `nn.GPT` module, decode the sampled continuation
- `anneal train gpt2`: fine-tune GPT-2-small end to end on tinyshakespeare from the same HuggingFace weights - tied embedding/LM-head gradient, numerically stable cross-entropy, global-norm clip, AdamW + LR warmup, JIT-replayed train step (~40s/step on M3); the eval-batch loss converges
- `anneal viz` - launch the graph visualizer (WASM)
- `anneal web [addr]` - serve the local studio (eight deep-linkable views, no telemetry)
- `anneal graph mlp` - inspect the UOp DAG
- `anneal kernels mlp` - show generated WGSL with fusion boundaries
- `anneal explain add` - trace the rewrite rules that fire for one op
- `anneal doctor` - WebGPU / backend environment check

## Tutorial: how to train ML models from scratch in Go

The step-by-step tutorial (https://georgebuilds.github.io/anneal/tutorial/) is the place to learn how to train neural networks with anneal on a real GPU, in pure Go, with no Python and no CUDA. It walks through, in order:

- Install (#install): install the Go toolchain and build the `anneal` CLI.
- Verify your GPU (#sanity): run `anneal doctor` to confirm a working WebGPU adapter (Metal, Vulkan, or D3D12).
- Train a model (#train): run `anneal train <model>` to train end to end on your GPU. Choose from:
  - `mlp`: a tiny multilayer perceptron (the fastest first run).
  - `conv`: a small ConvNet.
  - `nanogpt`: a char-level GPT transformer on tinyshakespeare.
  - `llama`: a char-level Llama-style decoder (RMSNorm, grouped-query attention with RoPE, SwiGLU, tied embeddings); the modern-decoder level-up over nanoGPT.
  - `bert`: a char-level BERT encoder trained with masked language modeling; encoder-only and bidirectional, the contrast to the causal nanoGPT/Llama decoders.
  - `vit`: a tiny Vision Transformer on a synthetic image dataset.
  - `resnet9`: ResNet-9 on CIFAR-10.
  - `gpt2`: fine-tune GPT-2-small from HuggingFace weights.
  - `moe`: a small GPT whose per-block feed-forward is a router plus several expert FFNs (soft routing with a load-balance loss); a teaching version of sparse Mixture-of-Experts.
  - `dit`: a class-conditional Diffusion Transformer (adaLN-zero, classifier-free guidance) trained with epsilon-prediction on CIFAR-10.
  - `meanflow`: a MeanFlow one-step generative model on the DiT backbone; its average-velocity objective uses a forward-mode JVP for the total time-derivative (the first use of forward-mode autodiff in the compiler for a real objective).
- Explore the compiler (#tour): inspect the generated WGSL kernels (`anneal kernels <model>`) and the forward/backward UOp graph (`anneal graph <model>`, `anneal viz`).

Section anchors: #start #lessons #prereqs #install #sanity #train #tour #more #further #troubleshooting #qa
Full tutorial text for ingestion: https://georgebuilds.github.io/anneal/llms-full.txt

## Links

- Homepage: https://georgebuilds.github.io/anneal/
- Tutorial: https://georgebuilds.github.io/anneal/tutorial/ : Step-by-step guide to training ML models from scratch in Go on a real GPU. Install anneal, verify your GPU with `anneal doctor`, then train MLP, ConvNet, nanoGPT, a Llama-style decoder, BERT, ViT, ResNet-9, MoE, a Diffusion Transformer (DiT), and MeanFlow, or fine-tune GPT-2-small; covers WebGPU prereqs, the asset download cache, the BPE tokenizer, and a compiler-level walkthrough.
- Visualizer demo: https://georgebuilds.github.io/anneal/visualizer-demo/ : Static UOp graph visualizer for the MLP model. Forward / backward / fused-kernel coloring; a scrubbable timeline through compiler passes; a rules sub-timeline on the gradient and scheduled stages that animates backward nodes appearing in reverse-topological order per BackwardWithTrace's GradTrace, with SVG tooltips naming the gradient rule and firing seq on each backward node.
- Source: https://github.com/georgebuilds/anneal