Iron Architecture

How a #[kernel] becomes a compiled GPU shader, and how iron bench / iron test / iron build run and measure it today. Companion docs: TOOLCHAIN_DESIGN.md (the #[kernel] / #[kernel(variants(...))] / #[bench] / #[test_kernel] macro surface), BENCH_METRICS_SPEC.md (metric definitions), KERNEL_CONSOLIDATION_PLAN.md (the kernel restructure roadmap), the backend specs (CUDA / AMD / VULKAN / ANE), cli.md (command flags), developing.md (kernel-authoring hazards).

Two things drive the current shape of the runtime:

1. The runner is a subprocess. iron spawns a generated __iron_runner binary (linked against the project’s wh-iron-std, so its kernel inventory is populated) and streams results back as ProtocolMessage JSON lines; see Subprocess execution.
2. Codegen is multi-backend. IR lowers through a CodegenBackend to Metal (default), CUDA, HIP/ROCm, or SPIR-V/Vulkan. Metal and CUDA execute production workloads; HIP and Vulkan remain narrower feature-gated bring-up paths: see Multi-backend codegen.

Read the page in this order if you are new to the project: first the ownership boundary below, then Crates, From source to shader, Command dispatch, and finally the test and benchmark runners. The registry and subprocess sections explain machinery that matters when a kernel disappears or a runner executes stale code.

Where Iron stops

Iron is the kernel and device layer. It is not the inference engine.

Iron owns	The consuming engine owns
Kernel math and write/read semantics	Model graph and layer schedule
Dtype, packed layout, strides, offsets, and tails	Checkpoint interpretation and tensor ownership
IR passes and backend lowering	Request admission and compatible batching
Module/pipeline creation and launch geometry	KV and recurrent-state lifecycle
Component oracles and kernel measurements	Prefill, decode, speculation, sampling, and serving metrics

The boundary requires evidence from both sides. A component oracle can prove that one kernel implements its declared operation. It cannot prove that the engine chose the right tensor, state version, position, or execution order. Likewise, a coherent model sample does not prove an untested packed layout or tail path.

Crates

flowchart TD
    macros["wh-iron-macros<br/>#[kernel] · variants · #[bench] · #[test_kernel]"]
    core["wh-iron-core<br/>IR + wire protocol<br/>(no GPU)"]
    codegen["wh-iron-codegen<br/>passes + backends<br/>(MSL · CUDA · HIP · SPIR-V)"]
    runtime["wh-iron-runtime<br/>device dispatch<br/>(Metal · CUDA · HIP · Vulkan)"]
    facade["wh-iron (facade)<br/>harness/ + runner/"]
    std["wh-iron-std<br/>kernel stdlib<br/>(mlx · iron · convolution · quant)"]
    cli["wh-iron-cli<br/>tile binary (thin)"]

    macros --> core
    codegen --> core
    runtime --> core
    facade --> core
    facade --> codegen
    facade --> runtime
    std --> facade
    cli --> facade
    cli --> codegen
    cli --> std
    cli --> core

Crate	Responsibility
wh-iron-core	IR (Op, Kernel) + the protocol wire types (ProtocolMessage, ProfileInfo). Pure; no GPU or tooling dependencies.
wh-iron-macros	#[kernel] (lowers a DSL fn to IR) + #[kernel(variants(...))] (compile-time specialisation, stamping one kernel per tuple of int/type/float params); #[bench] / #[test_kernel] (register a setup callback via inventory).
wh-iron-codegen	Optimization passes (const-fold, vectorize, unroll, fusion, DCE, …) + the CodegenBackend seam: msl/ (Metal, default) and the cuda/ / hip/ / spirv/ generators. backend.rs holds the Target enum, TargetProfile, and MmaStrategy.
wh-iron-runtime	Per-backend device, buffers, PSO/module cache, dispatch + timing. device/metal_device.rs is the default; device/{cuda,hip,vulkan}/ are feature-gated (--features cuda|hip|vulkan).
wh-iron (facade)	Re-exports the above; hosts harness/ (the kernel/bench/test registries) and runner/ (the __iron_runner engine: RunnerHarness, GpuRunner, per-backend dispatch, arg parsing, protocol emit, profiling, device specs).
wh-iron-std	The kernel standard library. Modules: mlx/ (kernels with an upstream metal reference), iron/ (model-specific), convolution/ (consolidated 1D/2D/3D/depthwise/winograd + steel_conv/), quant/ (the codec + format + gguf precision layer), plus probes/ and utils. Every #[kernel]/#[bench]/#[test_kernel] lives here.
wh-iron-cli	The iron binary: config, command dispatch, result rendering. It spawns the runner subprocess rather than doing GPU work itself.

From source to shader

flowchart LR
    K["#[kernel]<br/>fn iron_exp&lt;T&gt;(..)"] --> IR["IR<br/>(Op variants)"]
    V["#[kernel(variants(...))]"] -. "stamps N kernels" .-> IR
    IR --> Passes["codegen passes<br/>const-fold · vectorize · unroll · FMA · DCE · …"]
    Passes --> Backend{CodegenBackend<br/>Target}
    Backend -->|Metal default| MSL[".metal → kernels.metallib"]
    Backend -->|Cuda| CU[".cu (NVRTC/PTX)"]
    Backend -->|Hip| HIP[".hip (hipRTC)"]
    Backend -->|Spirv| SPV[".comp → SPIR-V"]
    B["#[bench] / #[test_kernel]"] -. "inventory::submit!" .-> Reg["harness registries<br/>all_kernels / all_benches / all_tests"]

#[kernel] lowers the DSL function to IR; #[kernel(variants(...))] stamps out one specialised kernel per parameter tuple before lowering. The codegen passes optimise the IR (backend-independent), then a CodegenBackend emits the target source. Metal is the default backend, and MSL emission produces a .metal source that xcrun metal compiles to a metallib. The feature-gated CUDA runtime is also a fully executing production path: it compiles emitted CUDA through NVRTC, loads PTX or CUBIN through the driver API, and supports graph, matrix-library, and AOT helper paths used by Butter. #[bench] and #[test_kernel] are optional annotations on the same function that register a setup callback (BenchSetup / TestSetup) into an inventory registry the runner iterates.

One operation from definition to production

The full lifecycle is longer than “write a kernel and benchmark it”:

1. Define the semantic contract. Name inputs, outputs, dtype, packed layout, shape domain, offsets, tails, reduction order, and non-finite policy.
2. Lower and inspect. Run pure IR/codegen tests and inspect the generated target source for the intended variant.
3. Prove the component. Execute the registered test or backend corpus against an independent CPU/component oracle, including boundary shapes and a known-bad mutation.
4. Integrate through a semantic operation. The consuming engine selects the route only for the exact qualified backend, dtype, shape, and graph mode.
5. Prove composition. Run checkpoint, model-state, eager/captured, and multi-token tests in the consuming engine.
6. Measure both clocks. Record isolated kernel time and the affected end-to-end prefill, decode, or serving metric with a rollback arm.
7. Promote or reject. Keep the selector and receipt for a winner. Preserve the failed gate and reopen condition for a rejection.

This sequence prevents a fast microbenchmark from becoming a model route without layout, state, and user-visible evidence.

Command dispatch

Every subcommand is a struct implementing the IronCommand trait (cmd/mod.rs); main.rs parses args, builds a Harness (the loaded IronConfig), and dispatches:

flowchart TD
    main["main.rs<br/>parse args"] --> cfg["ConfigLoader<br/>defaults → iron.toml → IRON_* env → CLI args"]
    cfg --> harness["Harness (owns IronConfig)"]
    harness --> cmd{IronCommand}
    cmd --> bench["bench"]
    cmd --> test["test"]
    cmd --> build["build (+ --emit)"]
    cmd --> inspect["inspect"]
    cmd --> other["device · snap · diff · clean · config · update · init<br/>(pure CPU / IO; stay in-process)"]
    bench --> pr["ProjectRunner<br/>spawn __iron_runner"]
    test --> pr
    build --> pr
    inspect --> pr

bench / test / build / inspect route through ProjectRunner, which spawns the runner subprocess. device (GPU query), snap (save baseline), diff (compare baselines), clean (remove build artifacts), config (show resolved config), update (self-update), and init (scaffold a new project) are pure CPU / IO and run directly.

emit is a build flag, not a command. iron build --emit msl|metallib|swift|ir|all --out <dir> writes per-kernel .metal, the compiled kernels.metallib, the IronKernels.swift bindings, and/or the manifest.json IR descriptor. build / inspect always emit MSL (the --backend flag below applies only to bench / test).

Kernel registry

Registration deliberately spans three crates, and the split isn’t obvious from the call sites:

• wh-iron-core re-exports the inventory crate, so the macro-expanded inventory::submit! calls have a single canonical path to submit to.
• wh-iron-codegen owns KernelEntry + all_kernels() (src/kernel_registry.rs), placed next to KernelInlinePass, its only consumer (the inliner needs the full kernel set to resolve cross-kernel primitive calls).
• wh-iron (facade) holds the bench/test registries in harness/registry.rs (all_benches / all_tests), consumed by the runner.

The load-bearing detail: inventory statics live in a linker section and are garbage-collected if nothing references the library. The __iron_runner bin in wh-iron-std (bin/runner.rs) exists largely to do extern crate wh_iron_std;; that one line forces the linker to keep every submit! static, so the registries are non-empty inside the child process. iron init scaffolds a per-project copy of this bin for downstream projects. Deleting the extern crate line as “dead code” silently empties the registries because it is intentional, not cruft.

Subprocess execution

The iron CLI does no GPU work itself. ProjectRunner spawns __iron_runner, a binary linked against the project’s wh-iron-std, so the #[kernel]/#[bench]/#[test_kernel] inventory is populated inside the child and streams its stdout, parsing each line as a ProtocolMessage.

flowchart LR
    cli["iron CLI<br/>(thin protocol parser)"] -- "spawn + args" --> proc["__iron_runner<br/>RunnerHarness + GpuRunner + inventory"]
    proc -- "ProtocolMessage JSON lines (stdout)" --> cli

Piece	Where	Purpose
ProtocolMessage (+ runner_version)	wh-iron-core::protocol	Versioned JSON-line wire format (CLI ↔ runner).
RunnerArgs	wh-iron::runner::args	Subprocess CLI arg parsing (from_env_args), incl. --backend.
RunnerHarness	wh-iron::runner::harness	Orchestrates bench / test / build / inspect, emitting protocol messages.
runner::backend	wh-iron::runner::backend	Routes --backend cuda|hip|vulkan through the matching feature-gated device.
ProjectRunner	wh-iron-cli::project_runner	Spawns __iron_runner, streams + parses its stdout.

Multi-backend codegen

backend.rs defines Target { Metal, Cuda, Hip, Spirv } and the CodegenBackend trait; MslGenerator is the Metal impl, with cuda::CudaGenerator / hip / spirv alongside. A TargetProfile carries the per-backend knobs the IR-to-source lowering needs: SIMD/warp lane width (32 on Metal simdgroup / CUDA warp / RDNA wave32, 64 on CDNA wave64, variable on Vulkan subgroup) and the MmaStrategy (Metal simdgroup_matrix 8×8, CUDA/CDNA/RDNA tensor-core paths, Vulkan VK_KHR_cooperative_matrix).

• Codegen (emit) supports all four targets; build/inspect emit MSL.
• Execution is Metal by default. iron bench|test --backend cuda|hip|vulkan routes through the feature-gated CudaDevice / HipDevice / VulkanDevice in wh-iron-runtime. CUDA is production-used and must be validated against independent CPU/component oracles plus its registered kernel corpus; it is not merely a GPU-to-Metal comparison path. HIP and Vulkan use the same inventory mechanism but currently cover a narrower subset. Relevant corpora live under tests/{cuda,hip}_kernel_corpus.rs and tests/vulkan_sdpa_multi.rs.

See the backend specs (CUDA / AMD / VULKAN / ANE) for each target’s design and hazards.

Bench runner

flowchart TD
    RH["RunnerHarness (in __iron_runner)"] --> loop["for each #[bench] × dtype (sequential)"]
    loop --> emit["backend emit + compile (PSO/module cache)"]
    emit --> timed["⏱ GpuRunner.bench(warmup, iters)<br/>→ BenchStats (min/mean µs)"]
    timed --> dm["metrics from stats: to_gflops · estimate_profile ·<br/>classify_bottleneck · device_specs::lookup (CPU-only, AFTER timing)"]
    dm --> res["ProfileInfo { GFLOP/s · %-peak · bottleneck }"]
    res --> proto["ProtocolMessage → CLI SuitePrinter<br/>(GB/s · GFLOP/s · %-peak · bottleneck)"]

The bench run loop is sequential: GPU dispatch + timing is serialized on the device, so running benches concurrently would corrupt timings. (The CPU-only work that can parallelize, such as iron build MSL emit and iron test oracles, uses rayon; the bench run does not.)

Timing isolation. GpuRunner.bench(…) is the only timed region; it returns the finalized BenchStats. Metric derivation, including GFLOP/s (gpu::to_gflops), the roofline / bottleneck verdict (profile::estimate_profile

• profile::classify_bottleneck), and %-of-peak (device_specs::lookup), runs strictly afterward and consumes those stats; it never dispatches to the GPU. So metric computation cannot skew the measured kernel performance.

Test runner

iron test iterates the #[test_kernel] registry and dispatches each setup, comparing GPU output against the test’s CPU oracle within its tolerance. The CPU oracle pass is rayon-parallel (order-preserving via collect); GPU dispatch of the survivors is sequential. Under --backend, the same inventory runs on the selected device and is compared against the declared independent oracle. A Metal side-by-side is a separate optional comparison, not the correctness authority for CUDA, HIP, or Vulkan.

Kernel profiling

The roofline / occupancy metrics shown under iron bench -v / -vv:

flowchart LR
    stats["BenchStats<br/>(measured µs)"] --> g["gpu::to_gflops"]
    ir["kernel IR"] --> prof["profile::estimate_profile<br/>+ classify_bottleneck<br/>(occupancy · registers · bottleneck)"]
    specs["device_specs::lookup<br/>(peak BW + FP32/FP16 TFLOPS)"] --> PI
    g --> PI
    prof --> PI
    PI["protocol::ProfileInfo<br/>latency µs · GFLOP/s · %-peak BW/compute · arith intensity · bottleneck"]

• profile::estimate_profile runs the pass pipeline + register/occupancy analysis on the IR (pure CPU).
• device_specs::lookup returns per-chip peak ceilings (bandwidth + FP32 TFLOPS, with FP16 = 2× FP32 on the SIMD pipe, and the M5 Neural-Accelerator FP16 ceiling where applicable). Unknown devices return None, which leaves roofline columns blank rather than raising an error.
• protocol::ProfileInfo is the serializable wire form (GFLOP/s, %-peak BW / compute, arithmetic intensity, bottleneck) the runner streams back to the CLI for rendering.

See BENCH_METRICS_SPEC.md for the metric formulas and device-spec sourcing.