Rust Efficiency in Iron

This guide covers efficient Rust and generated GPU work in Iron. The goal is not shorter source or fewer launches by itself. The goal is lower production cost while preserving arithmetic, layout, backend, and maintenance contracts.

Use Rapid Software Testing in Iron to design the evidence for any optimization.

Agent quick contract

1. Locate the measured bottleneck before changing code.
2. Separate compiler-host cost, runtime-dispatch cost, and GPU-kernel cost.
3. Record a correct current baseline at the production dtype and shape.
4. State the mechanism and the maximum end-to-end gain it could provide.
5. Preserve deterministic code generation where output order matters.
6. Keep allocation and repeated string/map work out of hot loops.
7. Treat precision, reduction order, codec behavior, and output rounding as API.
8. Inspect registers, spills, occupancy, traffic, and launch geometry.
9. Re-run a real-device oracle and the consuming model.
10. Keep only a measured production win; record negative results.

Identify the cost domain

Domain	Typical evidence	Typical action
DSL/macro expansion	compile profile, generated token count	simplify expansion, share canonical walkers
IR passes	--time-passes, op-count deltas, allocation profile	pre-size collections, avoid repeated traversal, improve representation
Source emission	deterministic snapshots, string/allocation profile	stream output, cache stable fragments, preserve order
Runtime dispatch	CPU trace, cache hit rate, binding count	cache compiled state, precompute bindings, reuse buffers
GPU kernel	device timing, roofline, registers, occupancy	change tiling, vectorization, fusion, specialization
Model composition	Butter phase trace and token clock	change the operation boundary or decline the local optimization

Do not use a GPU-only metric to claim a compiler win, or a microkernel metric to claim faster decode.

Efficient Rust in the compiler and runtime

Allocate with a reason

• Use Vec::with_capacity when the pass or registry knows the upper bound.
• Use SmallVec only for measured small modal cardinality. Large inline arrays increase stack and move cost on every value.
• Reuse output buffers and descriptor storage across repeated dispatch.
• Keep tiny scalar uploads as typed stack bytes where the backend supports it.
• Avoid format!, String::from, .to_vec(), and cloned maps in per-kernel or per-dispatch loops.

An allocation removed from setup is usually irrelevant. An allocation removed from every decode dispatch may be valuable. Measure the actual call frequency.

Choose map semantics deliberately

• Fast hash maps suit internal numeric IDs and already-computed fingerprints.
• Pre-size maps when the input block length gives a useful bound.
• Keep ordered maps/sets when iteration controls parameter indices, emitted source, snapshots, content hashes, or user-visible output.
• Use stable content hashing for identities that cross processes or runs. A fast in-process map has a different contract.

A blanket map replacement can make codegen nondeterministic even when it makes lookups faster.

Keep one canonical structural walker

If multiple passes independently match every Op or IR variant, they will drift as the enum grows. Share canonical value/effect/block traversal helpers and test new variants once at that boundary.

Every transforming pass must handle the entry body and nested blocks. A pass that removes a body must also repair its enclosing control structure.

Cache at stable boundaries

Useful cache keys include all inputs that change output:

• kernel identity;
• backend and target capability;
• dtype and compile-time values;
• source or IR fingerprint;
• launch/tactic facts baked into the artifact.

Do not cache a context-owned object beyond its context. Do not reuse a graph, pipeline, or descriptor when a captured pointer or branch changed.

Locks and sharing

• Use Rc for values proven thread-local and Arc for cross-thread ownership.
• Keep lock scope below compilation, device synchronization, and disk I/O when possible.
• Use single-flight compilation for identical artifacts.
• Measure concurrent cache misses and hits. An uncontended lock result does not prove the compile storm or serving path.

Comments and errors

Prefer names and types over comments that restate code. Keep comments for:

• codegen ordering constraints;
• dispatch and layout invariants;
• backend API quirks;
• precision or rounding boundaries;
• lifetime and synchronization rules;
• pinned benchmark methodology.

Use expect("specific invariant") at genuinely infallible internal sites. At public or data-dependent boundaries, return a useful error instead of panicking.

Build profiles

Iron keeps downstream-friendly release settings and moves aggressive whole-program optimization to benchmarks:

[profile.release]
codegen-units = 16
panic = "abort"

[profile.bench]
inherits = "release"
lto = "fat"
codegen-units = 1
debug = true

Use the bench profile for performance evidence. Keep debug information there for profiling. Do not force fat LTO on published library consumers without a measured build/runtime tradeoff.

Generated kernel efficiency

Start from the production shape

Decode, prefill, and speculative verification create different launch regimes. Benchmark the exact dtype, dimensions, batch width, and backend used by the consumer. A large-matrix result does not predict a one-row projection, and a short-context attention result does not predict a long-context result.

Preserve the numerical boundary

The following can change generated tokens even when a local tensor looks close:

• reduction partition or order;
• accumulator dtype;
• packed-code special values;
• scale axis and scale dtype;
• output dtype and store rounding point;
• recurrence state read/write order;
• tie handling in reductions or top-k.

Apply the component oracle first, then the consuming model’s deterministic trajectory and state oracle. A faster non-exact mode must be explicit and separately gated.

Vectorize transfers without widening the math change

Wider aligned loads can improve transaction utilization while retaining the same scalar computation and accumulation order. Prove:

• alignment and tail behavior for every production shape;
• tensor view offsets are honored;
• packed element boundaries are not crossed incorrectly;
• special-code decode matches the canonical codec;
• output bits or documented tolerance remain valid.

Do not combine transfer vectorization, reduction restructuring, and dtype changes in one experiment. That destroys causal and numerical diagnosis.

Specialize stable degenerate shapes

Compile-time or dispatch-time specialization is useful when an invariant is common and removes real control or data work, such as one-token decode, fixed head dimensions, or a fixed speculative verify width.

Keep the general path. Gate the specialized path by exact shape/capability and test the boundary immediately outside the specialization.

Fusion rules

Fusion is promising when a producer’s intermediate would otherwise be written to memory and immediately read by one consumer. This is vertical fusion.

Horizontal fusion of independent operations usually saves only launch setup while increasing bindings, live values, and register pressure.

Before writing a fused kernel, answer:

1. Which intermediate write and read disappear?
2. How many consumers use the value?
3. Are threadgroup and reduction shapes compatible?
4. Does fusion preserve the reduction tree and rounding point?
5. How many extra pointers and live values are introduced?
6. What happens to registers, spills, shared memory, and resident blocks?
7. Is the baseline already in one encoder, stream, or captured graph?
8. What production phase share bounds the end-to-end gain?

Prefer a producer-side epilogue when it computes an expensive transform once at the store boundary. Avoid moving that transform into a fan-out consumer that recomputes it for every output element.

Keep reductions separate when their cooperative shapes differ. A two-pass pipeline can be more efficient and clearer than forced fusion.

Launch and occupancy

Dispatch geometry is part of the kernel ABI.

• Distinguish threads from threadgroups in every backend.
• Document grid and threadgroup invariants beside the kernel.
• Check shared-memory size, alignment, and tail guards with checked arithmetic.
• Inspect compiler register count and local/spill bytes.
• Sweep meaningful occupancy discontinuities rather than assuming more threads or fewer registers is always better.
• Keep enough independent work in flight. Reuse that doubles accumulators may lose more occupancy than it saves in loads.

Shared-memory staging is not automatically faster. It can add barriers and reduce occupancy when the cache already serves the reuse. Measure it against the best unfused/resident baseline.

Removing a barrier is useful only when the memory-order proof is sound and the production clock moves. A provably redundant barrier may still be hidden below noise.

Benchmark discipline

Use resident inputs for values constant across iterations and warm the device before measurement. Record:

• exact commit, backend, device, OS/driver/compiler;
• dtype, shape, launch geometry, and bytes/FLOPs;
• build profile;
• correctness oracle and known-bad mutation;
• warmup, iterations, trials, median, and dispersion;
• timer boundary;
• registers, spills, shared memory, occupancy, and bottleneck where relevant;
• consuming Butter result.

The default command surface is:

# Correctness first
make test
make iron test

# Kernel metrics
make bench-v
make bench-vv

# Compiler and generated-source inspection
make time-passes
make inspect-stats KERNEL=<name>
make inspect-ir KERNEL=<name>
iron build --emit all -o /tmp/iron-smoke

For one kernel, use a filter through the CLI or the documented Cargo correctness/performance companion in Testing.

Treat these as harness warnings:

• latency does not scale when bytes/work scale materially;
• a result changes with test order;
• the first cell is consistently slower;
• an unsupported path is reported as a pass;
• output is empty but scored equivalent;
• the candidate has no trace/path marker;
• microbench improves while the consumer is flat or slower.

Optimization loop

1. Baseline: current shippable path, correct and reproducible.
2. Profile: assign cost to compiler, dispatch, transfer, kernel, or model.
3. Bound: multiply candidate phase share by plausible local improvement.
4. Hypothesis: one mechanism, one expected observation.
5. Microgate: exact production shape and complete operation.
6. Correctness: CPU/codec/layout oracle plus known-bad mutation.
7. Resources: registers, spills, occupancy, traffic, and launch proof.
8. Consumer: Butter component and multi-token model gates.
9. Performance: same-binary A/B where possible, multiple warm medians.
10. Decision: keep, revise, or record as a negative result.

Review checklist

• The measured bottleneck and production caller are named.
• No new repeated allocation, formatting, hashing, or cloning exists on the hot path.
• Ordered output remains deterministic where required.
• Cache keys include every output-affecting input and respect context lifetime.
• Dispatch, layout, alignment, and tail contracts are documented and tested.
• Numerical boundaries and special-code behavior are preserved.
• The checker was demonstrated against a known bad case.
• Target hardware executed the path.
• Registers, spills, occupancy, and traffic were inspected when applicable.
• The Butter production clock improves outside variance.
• Negative findings and unsupported backends are recorded.

Related guides

• Rapid Software Testing in Iron.
• Testing.
• Kernel style guide.
• Architecture.
• Bench metrics.
• Developing.