콘텐츠로 이동

Developing

이 콘텐츠는 아직 번역되지 않았습니다.

Repo layout, the dev loop, and — most importantly — the kernel-authoring hazards that cause silent or catastrophic failure. If you only read one section, read Kernel-authoring hazards.

Short terms used throughout the docs and the codebase:

Term Meaning
DSL Domain-specific language — the Rust-embedded language you write inside a #[kernel] function.
IR Intermediate representation — Iron’s typed kernel graph (Kernel / Op), produced by the proc-macro and consumed by the codegen.
MSL Metal Shading Language — Apple’s C++-based GPU shader language, and the codegen’s final output.
Kernel A GPU compute function; in this repo, a Rust fn annotated with #[kernel].
Threadgroup A block of GPU threads that share threadgroup memory and can synchronise with a barrier.
Simdgroup 32 threads (“lanes”) within a threadgroup that execute in lockstep — Apple’s SIMD width. Cross-lane ops are the simd_* intrinsics.
TPG Threads per threadgroup — the threadgroup dimension of a dispatch. Must be a multiple of 32 for reduction kernels.
PSO Pipeline state object — a compiled Metal compute pipeline. The runtime caches these so a kernel compiles once.
Pass One transformation stage in the codegen pipeline (e.g. ConstFold, Vectorize).

The workspace is seven crates, layered from the shared data model up to the CLI binary:

Crate What it is
wh-iron-core The shared data model — the Kernel / Op IR types, DType, Shape, and error types. Pure data structures with no logic, so every layer above can speak the same vocabulary. Used by every other crate.
wh-iron-macros The compiler front end — the #[kernel] proc-macro and its body parser, which turn a Rust function into Iron IR at compile time. Owns the DSL grammar and its compile-error diagnostics. Used by wh-iron (re-exported as the public #[kernel] attribute).
wh-iron-codegen The optimizing compiler — lowers Iron IR through the 14-pass pipeline and emits MSL. The largest crate; owns every pass and the MSL emitter. Used by wh-iron-runtime, wh-iron-std, and wh-iron-cli.
wh-iron-runtime The GPU execution layer — compiles emitted MSL into Metal PSOs (with a PSO cache) and dispatches kernels through Context. Owns all Metal-framework / objc2 interop. Used by wh-iron and wh-iron-std.
wh-iron The facade — re-exports core, macros, codegen, and runtime behind one prelude so downstream code and external users depend on a single crate. No logic of its own. Used by wh-iron-std and wh-iron-cli.
wh-iron-std The kernel standard library — the actual #[kernel] definitions (mlx/, iron/), their BenchSpecs, the bench harness, and the GPU correctness tests. Where new kernels land. Used by wh-iron-cli.
wh-iron-cli The iron binary — bench / build / inspect / device / snap / diff, the developer-facing entry point. The top of the dependency graph; nothing depends on it.

The compile pipeline: #[kernel] fnwh-iron-macros parses the body into Iron IRwh-iron-codegen runs the optimization passes and emits MSLwh-iron-runtime dispatches it on the GPU.

Terminal window
make build # debug build
make test # workspace tests — codegen, runtime, GPU correctness (GPU on a Mac)
make clippy # lint with -D warnings
make fmt # format
make fmt-check # check formatting without writing
make typos # spell-check
make coverage # HTML coverage report (needs cargo-llvm-cov)
make bench # full benchmark suite vs MLX (macOS + Metal)
make clean # remove target/

Prefer make over raw cargo — it centralises flags and always passes --workspace. See the CLI reference for the iron binary.

CUDA’s default cooperative matmul uses inline mma.sync for eligible F16/BF16 tiles on sm_80+. CudaDevice selects software on older devices; F32 and ineligible tiles also keep their software fallback. Offline CudaGenerator::new() assumes sm_80+, while for_compute_capability(major) selects for a known device. Use IRON_CUDA_MMA=software, =wmma, or =mmasync for explicit comparisons. Explicit tensor-core choices retain their architecture requirements. When constructing a profile directly, pin MmaStrategy::SoftwareLocalC for a software control instead of inheriting TargetProfile::cuda().mma. Always use the selected generator’s shared_bytes with the actual threadgroup size.

The iron wrappers (make emit-all, make bench-vv, make inspect-*, make time-passes, make iron …) also rebuild the __iron_runner subprocess binary first via make runner. If you invoke cargo run -p wh-iron-cli by hand instead, you skip that step and can emit stale codegen — see “A stale __iron_runner silently emits old codegen” under Kernel-authoring hazards.

New kernels can declare their correctness test and benchmark inline with the #[test_kernel] / #[bench] attributes (run via iron test / iron bench) instead of, or alongside, the hand-written tests/*_gpu_correctness.rs files. Both styles coexist during the migration — see testing.md and crates/wh-iron-std/src/mlx/arange.rs for the template.

make hooks installs three git hooks (under scripts/hooks/, opt-in per clone via git config core.hooksPath) that mirror the CI checks locally so the slow round-trip to GitHub doesn’t surface trivially-catchable failures:

Hook When Checks Time
pre-commit every git commit make fmt-check + make typos ~2-5s
commit-msg every git commit AI-attribution scan — Co-Authored-By: & other trailer keys, 🤖 Generated with … footers — + Git-trailer-shape scan (Word:\s lookalikes). Naming files like CLAUDE.md / .cursor/, or disclosing AI use in prose, is fine. ~50ms
pre-push every git push make clippy + cargo test … --test kernel_registry_consistency ~30-90s

Run make hooks once after cloning. Bypass an individual hook with --no-verify (e.g. git commit --no-verify) when you know what you’re doing. Uninstall with make hooks-uninstall.

What’s caught locally: the cheap wins — formatting, common typos, AI-attribution trailers / footers in commit messages, trailer-shape lookalikes, clippy warnings, and the #[kernel] pub fninventory::submit! consistency check.

What still only runs in CI: the GPU correctness suite (needs Metal), cargo test --workspace (full test matrix), iron build --emit all (multi-minute), and the full bench harness. Pre-push handles the part of CI that fails most often without GPU.

Branch Purpose
main Stable releases only. Commits here are tagged (v0.1.0, …).
dev Integration branch for the next release. Feature PRs merge here.
feat/* fix/* perf/* docs/* Short-lived topic branches cut from dev.

Cut a topic branch from dev, PR back into dev, squash- or rebase-merge after review + green CI. See Publishing for the devmain release flow.

PR titles follow Conventional Commits so .github/workflows/auto-label.yml can categorise them for release notes and .github/workflows/pr.yml can validate the format:

feat: add softmax vector path for small N
fix(codegen): correct version gate for half2 stores
perf(runtime): cache PSO lookups by function signature
docs: update CLI install instructions
test(core): add scan correctness test
chore: bump nightly toolchain

Add ! for breaking changes (feat!: …) and describe them in the PR body.

Want Command
IR before any passes iron inspect <kernel> --ir
Final MSL iron inspect <kernel>
IR after one pass iron inspect <kernel> --pass <name> (or --pass all)
Per-pass op-count deltas iron inspect <kernel> --stats
Which pass is slow iron build --time-passes --filter <kernel>
Emit every kernel’s MSL iron build --emit all -o <dir>

When a kernel regresses, --stats before/after the change shows which pass changed the op count; --pass all dumps the IR at every stage.

These are not style preferences. Each one has bitten this project; each one fails silently (wrong output, no error) or catastrophically (frozen machine). Read all of them before writing a kernel.

⚠️ A wrong dispatch can freeze the machine

Section titled “⚠️ A wrong dispatch can freeze the machine”

Metal compute dispatches are non-preemptive — once a threadgroup starts, the GPU runs it to completion. An infinite loop inside a kernel never yields: the WindowServer compositor starves of GPU time, the screen locks at the last frame, and a hard power-cycle is the only recovery.

The concrete trap: reduction-mode kernels compute the simdgroup count as n_simd = lsize / 32 (integer division). A loop strided by n_simdfor _t in range(sg, n_kv, n_simd) — becomes an infinite GPU loop when n_simd == 0, i.e. when the kernel is dispatched with fewer than 32 threads per threadgroup. A 4-thread dispatch of a 1024-thread kernel once froze a dev machine for a full day. The kernel was correct; the dispatch geometry was not.

Rules for any kernel that uses simd_* / threadgroup_*:

  • Threads-per-threadgroup must be a multiple of 32 and ≥ 32 (one full simdgroup).
  • The dispatch geometry is part of the kernel’s contract — derive it from the kernel’s invariants, never from an unrelated “number of elements” count.
  • GPU correctness tests and BenchSpecs set the threadgroup size from the kernel side, so they are safe. The danger is any consumer that turns a caller-supplied dimension into a dispatch shape — guard those.

Context::dispatch_with_grid(kernel, buffers, constexprs, grid_xyz, tg_xyz) calls dispatchThreadgroups. grid_xyz is counted in threadgroups, not threads — total threads = grid.{x·y·z} · tg.{x·y·z}.

  • Grid3D — one thread per output element, no cross-thread cooperation. program_id::<i>() lowers to the thread index.
    #[kernel] fn mul<T>(a: Tensor<T>, b: Tensor<T>, out: Tensor<T>) {
    let i = program_id::<0>();
    store(out[i], load(a[i]) * load(b[i]));
    }
    // dispatch: grid=[1,1,1] tg=[N,1,1] (or grid=[ceil(N/TPG),1,1] tg=[TPG,1,1])
  • Reduction — uses simd_* / threadgroup_*. program_id::<i>() lowers to the threadgroup index; threads within a group cooperate.
    #[kernel] fn rms_norm<T>(x: Tensor<T>, /* … */) {
    let row = program_id::<0>(); // = threadgroup index, one TG per row
    // … reduce_sum across the threadgroup …
    }
    // dispatch: grid=[rows,1,1] tg=[TPG,1,1]

Wrong: grid=[N,1,1] tg=[N,1,1] for a Grid3D kernel — that is threads in flight, most with garbage indices. The product grid · tg must equal exactly the thread count the kernel expects.

⚠️ Inner macro_rules! silently empties the kernel body

Section titled “⚠️ Inner macro_rules! silently empties the kernel body”

To share a body across non-generic variants (bit widths, fixed group sizes), wrap the entire #[kernel] fn declaration in an outer macro_rules!never put a macro_rules! call inside a #[kernel] body. The proc-macro does not expand inner declarative macros: it sees the call as opaque tokens, drops it, and emits a kernel with no body. xcrun metal compiles it fine and it ships all-zeros output. The proc-macro now rejects the inner-body shape with a compile error — heed it rather than working around it.

Right — wrap the whole declaration; the compiler expands the outer macro before the #[kernel] proc-macro runs, so the body parser sees concrete tokens with $bits already substituted:

macro_rules! dequant_gather_kernel {
($name:ident, $bits:literal, $subop:literal) => {
#[kernel]
pub fn $name<T>(/* params */) {
let bit_off = d * $bits; // $bits already substituted
// …
}
inventory::submit! { /* BenchSpec */ };
};
}
dequant_gather_kernel!(dequant_gather_int4, 4u32, "int4");

Wrong — inner macro inside the body → empty MSL (now a compile error):

macro_rules! body { ($bits:literal) => { /* … */ }; }
#[kernel] pub fn dequant_gather_int4<T>(/* … */) { body!(4); }

Canonical reference: crates/wh-iron-std/src/kernels/sdpa/sdpa_prefill_sink.rs (an outer shim that only injects a per-family buffer identifier around a #[kernel(variants(...))] width axis — prefer variants(...) alone whenever the variants are purely numeric). For hand-unrolled tree reductions, replace *_step! macros with a DSL for loop over the halving strides — identical MSL, survives the proc-macro.

⚠️ threadgroup_alloc is hoisted to function scope — names must be globally unique

Section titled “⚠️ threadgroup_alloc is hoisted to function scope — names must be globally unique”

Every threadgroup_alloc("name", …) call is lifted to a single function-scope declaration, regardless of which branch it sits in. Two branches that each call threadgroup_alloc("tg_max", …) therefore collide into a duplicate threadgroup float tg_max[…] — a Metal compile error. This one at least fails loudly, but the fix is non-obvious: give every threadgroup buffer a name unique across the whole kernel, not just within its branch (e.g. a greedy path’s scratch is tg_gmax, the sampling path’s is tg_max).

⚠️ Empty-body MSL also slips through pass ordering

Section titled “⚠️ Empty-body MSL also slips through pass ordering”

The macro trap above is one way the codegen emits a kernel with a valid function/loop header but no body. The other is pass ordering: a pass eliminates a loop body but leaves the loop header, or a Const a later pass needs is still rolled inside a BinOp so the trip count is invisible. The result is for (…) { } — an empty loop — and again the kernel ships all-zeros output that xcrun metal accepts.

Invariants for any codegen pass you write or touch:

  • A pass that rewrites blocks must walk both kernel.body and every entry in kernel.blockskernel.body is the entry block, not part of the map.
  • A pass that removes a loop body must also remove the loop header.
  • A pass that consumes a Const must run after the pass that produces it.

Detection — emit every kernel (via make, so the runner is rebuilt first), then scan for empty bodies:

Terminal window
make emit-all OUT=/tmp/iron-smoke
awk '
/for \(.*\) \{$/ { f=1; fn=FILENAME; l=FNR; next }
f && /^[[:space:]]*\}$/ { print fn":"l": empty for-loop body"; f=0; next }
f { f=0 }
/^kernel void [A-Za-z_0-9]+\(/ { k=1; fn=FILENAME; l=FNR; next }
k && /^\{$/ { next }
k && /^\}$/ { print fn":"l": empty kernel body"; k=0; next }
k { k=0 }
' /tmp/iron-smoke/Resources/kernels/*.metal

Empty output = clean. Any hit = ship-stopper. Neither xcrun metal nor MSL snapshots catch an empty body — only a GPU correctness test does (see Testing).

⚠️ A stale __iron_runner silently emits old codegen

Section titled “⚠️ A stale __iron_runner silently emits old codegen”

iron is a thin protocol parser. The codegen — MSL emission, render_swift_wrappers, the metallib build — plus every bench and test dispatch runs inside __iron_runner, a separate [[bin]] of wh-iron-std (crates/wh-iron-std/bin/runner.rs) that iron spawns as a sibling of its own executable (target/release/__iron_runner, see runner_binary() in crates/wh-iron-cli/src/project_runner.rs).

The trap: cargo run -p wh-iron-cli rebuilds only iron. It never rebuilds __iron_runner, and iron does not check that the runner it finds is fresh. Edit crates/wh-iron-codegen/src/emit.rs (or any kernel .rs), run the CLI by hand, and the emitted .metal / metallib / IronKernels.swift come from the previous runner build — no error, no warning. Downstream in Butter this surfaces as a runtime fatalError or all-zeros kernel output, never as a build failure. It has also produced stale test verdicts: a correct kernel fix reported as still failing because the runner under test predated it.

Rules:

  • Go through make. Every CLI wrapper in the Makefile depends on make runner, which runs cargo build --release -p wh-iron-std --bin __iron_runner before spawning iron. Cargo’s fingerprint check makes it free when nothing changed.
  • If you must call cargo run -p wh-iron-cli … or target/*/iron directly, rebuild the runner yourself first, with the same profile (--release vs debug) — the sibling lookup only finds a runner in the same target/<profile>/ directory as the iron you are running.
  • After any codegen change, DO run the empty-body scan above on freshly emitted output — it is the cheapest way to confirm the artifacts you are looking at came from the code you just changed.

Document the dispatch contract: ## DISPATCH INVARIANTS

Section titled “Document the dispatch contract: ## DISPATCH INVARIANTS”

A reduction kernel’s threadgroup geometry is part of its API, but the kernel cannot enforce it at runtime. Make the contract explicit in the kernel’s .rs doc comment so anyone dispatching it has something to verify against:

//! ## DISPATCH INVARIANTS
//!
//! - **TPG: 1024 threads** (32 simdgroups × 32 lanes).
//! - **Grid: 1 threadgroup per row** (1D grid, program_id<0> = row).
//! - **head_dim == 128.** Each lane owns 4 consecutive elements; loads are
//! unconditional — other head dims read out of bounds / pin the GPU.

This is a young convention — most kernels don’t carry a block yet. Add one whenever you write or touch a reduction kernel; four lines, and it is the only place the geometry contract is written down.

⚠️ select does not guard a load — an optional buffer needs an if

Section titled “⚠️ select does not guard a load — an optional buffer needs an if”

A kernel that takes an optional buffer behind a constexpr flag and documents “with the flag off the buffer is never read (any 1-element buffer satisfies the binding)” must keep the load inside runtime control flow. A select is a function call: both operands are evaluated, so select(has_sink == 1u32, load(sink[q_head]), neg_infinity()) still reads sink[q_head] for every head, and with the documented one-element placeholder every head past the first reads out of bounds. Nothing reports it in a normal run — the discarded value is garbage, or (under shader validation) zero.

Right — the load is issued only when the feature is on:

let mut sink_score = neg_infinity();
if has_sink == 1u32 {
sink_score = load(sink[q_head]);
}

The if-conversion pass leaves this shape alone (an assignment in the arm is not predicated), so the emitted MSL is a real branch. Pin it: an in-source #[cfg(test)] that runs the kernel through MslGenerator and calls crate::utils::msl_check::assert_indexed_loads_guarded_by(&msl, "sink", "has_sink"), plus a device case in crates/wh-iron-std/tests/optional_buffer_placeholder_gpu.rs that binds the one-element placeholder with the feature off under Metal shader validation (see Testing).

Repeated reductions must finish reading shared partials

Section titled “Repeated reductions must finish reading shared partials”

A barrier before a shared-memory read does not protect that read from writes in the next loop iteration. In a two-level reduction, one SIMD group can finish its broadcast read and overwrite the same scratch while another group still needs the previous result. Ordinary test inputs and scheduling can hide this race.

The MSL emitter lets every SIMD group reduce the shared partials into a register before the second threadgroup barrier. The next iteration can then reuse the scratch safely. Custom reductions that broadcast through a shared slot instead need an additional barrier after every group has consumed it. SIMD reduction helpers must also return the complete result in every lane; shuffle-down Product needs a lane-zero broadcast.

crates/wh-iron-std/tests/msl_reduction_loop_gpu.rs checks every lane against an independent CPU fold, including bounded scheduling skew and output guards.

⚠️ Perf numbers that look “too flat to be physical”

Section titled “⚠️ Perf numbers that look “too flat to be physical””

A latency that does not scale with input size is almost always the harness contaminating the measurement, not an exceptional kernel. Before publishing a bench number, confirm the curve has the shape physics predicts. Two harness fixes belong in every perf measurement:

  • Use resident buffers for inputs constant across iterations — otherwise the bench re-uploads them every iteration and you measure host→GPU traffic, not the kernel.
  • Dummy-dispatch once to warm the GPU clock — cold DVFS gives the first measured shape a ~2× bandwidth deficit.

A “flat ~215 µs regardless of context length” sliding-window-attention claim turned out to be upload overhead drowning the kernel; switching to resident buffers dropped the floor to ~98 µs and revealed the real curve.

The full how-to lives in the Kernel Style Guide — file shape, naming, the #[kernel(variants(...))] axis, shared primitives, the CPU oracle, and the bench. Read it before adding a kernel; the principles below are the short version. (Where a kernel file belongs is the consolidation plan.)

  • Improve the compiler, don’t hand-write MSL. If the DSL can’t express a pattern, extend the codegen (body parser → IR → MSL emit). Don’t bypass it.
  • One generic <T> kernel beats five precision-specific copies — f32 / f16 / bf16 all flow through the same #[kernel] fn.
  • Every non-trivial kernel ships a GPU correctness test in the same commit. See Testing — it is the only layer that catches the empty-body and numeric-correctness bugs above.
  • Dispatch safety is the caller’s job — for now. Nothing in the kernel itself stops a degenerate dispatch from hanging the GPU (see the freeze hazard). A future direction is a runtime-checked dispatch path — kernel-level escape hatches that trap a runaway threadgroup before it pins the device — but until that lands, the consumer constructing the dispatch geometry is the only line of defense.