RFC: per-device op tuning cache (einsum first)

[czoli1976]

Problem. Codegen picks one implementation per op with fixed heuristics, and no fixed rule wins everywhere. Measured on an M4 against main: threading a matmul is right for Llama decode GEMVs (+28–36% tg128) and wrong for FastEnhancer's small AMX convs (serial is 35–48% faster multithreaded), and a topology rule that helped some models cost others 6–25% (#2982, closed). The right choice depends on shape, kernel and machine.

Proposal. Keep today's heuristics as the default. Add an opt-in autotuning step that records, per device, which alternative each optimisable op should use, and let codegen consult it.

• Optimisable ops expose a tuning key and their legal alternatives. Einsum first: strategize() already computes every suitable kernel and keeps one; it returns them all instead, heuristic pick first.
• Op identity, not node identity. The key is built from the op and its concrete facts, e.g. einsum_nk,km->mn_F32_F32_F32_b1_m288_k96_n64_aconst_nstrided. Every node with that key shares one decision, across layers and across models, so the file is a per-device cache (FFTW-wisdom style), not a per-model artefact.
• A choice is kernel + packing, and later thread count (the threading decision moves from a runtime gate to codegen and is stored on the op).
• Symbolic dims are tuned per scenario: symbols are made concrete from the scenario's assertions, with a hint where they don't pin a value (an LLM's prefill S), reusing the SymbolValues hints the GPU arena sizing already takes.
• tract … tune builds each optimisable op in isolation per scenario, benches its alternatives, writes or merges the entries, then checks the whole model end-to-end and drops entries that don't win.
• At codegen, a key that hits uses the tuned choice; a miss falls back to the heuristic. No tuning file means no behaviour change.

Design refinements (from review and measurements).

• A choice is kernel × thread count. AMX on M1–M3 and SME on M4 are one unit per cluster, so the best kernel family can flip with the thread count. Tuning runs at the deployment's thread pool and stream concurrency, and both are part of the key, so a single-thread entry never applies to a multi-thread run.
• Symbolic dims: nearest tuned shape. Tune each scenario at a few concrete sizes (an LLM: decode S=1, prefill S ∈ {64, 256, 1024}); a key whose size was not tuned takes the nearest tuned size of the same op signature, else the heuristic.
• Never worse than the heuristic. An entry is accepted only if the whole model gets faster with it, in back-to-back pairs against the current table (3/3 confirmations), and the finished table is dropped unless it beats no table end to end. Isolated op benches are only used to prune candidates, never to accept them.
• Spend the budget on heavy ops. Whole-model acceptance cannot see a win on an op that is a small share of runtime, so tune profiles first and screens keys by their share of the model's time.
• Tuning chooses, it doesn't invent. It only picks among alternatives that exist, so it complements implementation work rather than replacing it. Example: linalg: pack activation panels on the executor #2978 parallelises activation packing — new code that no tuner would have produced, and a fixed threshold that wins on mobilenet_v2 (−18% mt) but loses on inceptionv3 (+6% mt on the M1 Max). With packing ops exposed as optimisable ops (serial vs parallel per key), whole-model tuning would keep it where it wins and drop it where it doesn't, without retuning the threshold. Its profiling step would also surface such bottlenecks (serial packing next to parallel matmuls) as candidates for new implementations.
• Cache location. Reuse cuda's user_cache_dir() (XDG_CACHE_HOME / ~/.cache, Library/Caches on Apple): <cache>/tract/<version>/tuning/<cpu>-t<threads>.tune. The version in the path invalidates entries on upgrade, as for cubins.
• Cross-model reuse is a prior, not a fact. Entries record which models validated them; a model loading a shared cache validates the entries it uses and masks those that don't help it.

Evaluation protocol (acceptance criteria). Corpus: 11 noise-suppression models, PP-OCRv6 (det/rec), mobilenet/inception, LLMs at a few sequence lengths. Devices: M1, M4, x86, Cortex-A53/A55 (bench bot). Conditions: single- and multi-thread, idle and loaded. Tune on half the models, apply the shared cache to the held-out half, and report improved / neutral / regressed against the heuristic with paired statistics. Pass: no held-out regression beyond noise, gains on some models, bounded tuning time.

First measurements (M1 Pro, f32, paired whole-model method, single thread). FastEnhancer-large, GTCRN, DFN3 df_dec and PP-OCRv6 small_rec: no candidate beats the heuristic (0 of 16/17/7/25 keys); the guard keeps today's behaviour. In isolation one PP-OCR shape (960×48×96) runs 1.1–2.8× faster on AMX than the heuristic's NEON pick across thread counts, but in the model it is a 1×1 conv with a different layout taking ~0.5% of runtime, below what whole-model acceptance can see. Large f32 GEMMs: AMX wins at every thread count (no NEON crossover), and AMX+NEON co-execution has a lower ceiling than AMX on all cores. Single short runs swing ±10%, which is why acceptance needs paired confirmation.

Prior art. Op-signature keys and offline tuning with a persistent cache are the established pattern: TensorRT's timing cache (tactics timed at engine build, invalidated by GPU/TensorRT version), TVM's tuning records keyed by op workload (with pre-tuned logs shipped per target), ONNX Runtime's TunableOp (results saved or embedded in the model, guarded by device/version validators), XLA's GPU autotune cache. Serving engines add two ideas adopted above: vLLM ships tuned kernel configs per device name and uses the nearest tuned shape at run time, and TensorRT-LLM's min/opt/max optimization profiles match "scenario + hint". CPU engines (XNNPACK, OpenBLAS, oneDNN) mostly ship per-microarchitecture tables instead of autotuning, which Apple's AMX/SME/NEON mix and P/E cores make harder to get right by hand.

This PR (first slice). Behaviour-neutral plumbing, einsum only:

1. kernel_selection::candidates() returns the legal strategies, the strategize() pick first; strategize() itself is unchanged.
2. kernel_selection::tuning_key() for fully concrete shapes: expression, operand/accumulator types, batch, m, k, n, const-A, strided m/n.
3. tract_core::tuning: a key → (kernel, packing) table, process default or thread-scoped. Codegen's select() takes the table's choice when it names a legal candidate, the heuristic otherwise.
4. CLI: --tuning <file> loads a table; tuning-candidates lists each matmul's key and candidates. Keys are recorded from codegen itself rather than recomputed: by the time codegen runs, earlier rules have often swapped operands (the constant moves to A), so a key derived from the decluttered graph would never match.

On FastEnhancer-large (M1): 50 matmul nodes share 16 keys, each with 23 legal candidates; one table line moves the 5 nodes sharing a key from apple_amx_mmm_f32_32x32 to arm64simd_mmm_f32_8x8_gen.

Next slices: thread count as part of the choice, a tune subcommand that benches candidates and writes the table, scenarios and hints for symbolic dims, file versioning.

Open questions.

• Key format and versioning: what must invalidate an entry (tract version, kernel set, CPU model, pool size)?
• Scope after einsum: conv/im2col strategy, depthwise path, element-wise threading gates; layout/geometry alternatives would be pre-declutter and are left out.
• Isolated benches can mislead (cache state, shared units like AMX); the end-to-end check is the guard, at the cost of tuning time.
• WASM. One binary runs on unknown CPUs and engines (V8, JSC, SpiderMonkey compile simd128 differently), browser timers are coarsened and other tabs/throttling change conditions minute to minute, so measured or on-device tuning goes stale. Proposed instead: robust defaults per build (simd128/relaxed-simd), chosen offline from a cross-engine benchmark matrix (idle and loaded) by smallest worst-case slowdown and shipped as a table; host hints (engine, threads available, page hidden) passed in by the embedding app; conservative threading. The table API takes and returns bytes, not paths, so this stays possible.

Stacked on #3010 (first commit); the RFC is the second commit.

🤖 Generated with Claude Code

🍍

[czoli1976]

@kali editing description as I go with the experiments

[czoli1976]

First results with this slice and a scripted tune (whole-model paired acceptance, final never-worse guard), f32, ms per inference, 5 interleaved rounds. "main" = no table, i.e. identical codegen.

model	machine	threads	main	tuned	Δ
FastEnhancer-large	M1 Pro	all	2.726	1.584	−41.9%
FastEnhancer-large	M4	1	0.751	0.695	−7.5%
FastEnhancer-large	M4	all	1.461	0.764	−47.7%
mobilenet_v2	M1 Pro	1	17.246	14.813	−14.1%
mobilenet_v2	M4	1	10.870	9.426	−13.3%
mobilenet_v2	M4	all	7.567	7.302	−3.5%
DFN3 df_dec	M4	all	2.325	2.009	−13.6%

Every kept table was faster in 5/5 rounds. The other 13 cases (GTCRN and PP-OCRv6 small_rec everywhere, the rest of DFN3/FastEnhancer/mobilenet) found nothing better and keep main's codegen; one PP-OCR table that won only 2/5 rounds was discarded by the guard. So 7/20 faster, 13 unchanged, 0 slower.

The wins come from shapes where the heuristic picks the wrong kernel family: NEON where AMX/SME is faster, and on the M4 (which offers SME, AMX and NEON) sometimes SME where AMX or NEON is faster. Tables were tuned and validated on the same model and device; the held-out evaluation from the protocol is next. Thread count is not yet part of the choice.