Notes about implementing and working with the kernels

Kernels in tract-linalg are built from templated assembly, called through extern "C".

The templates live in linalg/$arch/$family — linalg/arm64/arm64simd, linalg/x86_64/avx512, and so on. A .S.j2 file is a Jinja2 template producing one kernel's assembly; a bare .j2 is a macro or include partial, and a .S.raw is raw assembly included as-is (brace-escaped). In general the file and the main entrypoint share a name stem. However, the proc name has a suffix based on the package version. In order to skip maintaining this the extern_kernel! macro — generated by build.rs into OUT_DIR — declares the matching function and re-exports it sans suffix.

Kernels work like a VM. When dispatching a kernel there's a list of instructions from FusedKerSpec that's dispatched in a jump table. For example; as of writing a MatMatMul is roughly encoded as [Clear, AddMatMul, Store, Done]. The assembly label that instructions return to is .non_linear_loop.

When iterating on assembly; building the code and looking at the generated assembly under target/$profile/build/tract-linalg-*/out/*.S can be much easier than tracking the flow through each macro.

If one needs to debug a kernel a useful workflow is to simply insert a mov rNN, [0] at the appropriate point, and configure GDB with handle SIGSEGV stop nopass. This'll pause in GDB but not send the signal to the program.

Benchmarking a kernel

tract hwbench <M,K,N[,dt]>... times every matmul micro-kernel this machine can run at each shape and prints their flop/s sorted fastest-first, marking the one the dispatcher currently picks with <--. This is the tool for both checking a dispatch decision and calibrating one. Add --json to parse the results, or --assert (with --tolerance <pct>, default 5) to fail when a pick lags the fastest kernel — the basis for a CI that guards kernel selection. The portable Rust kernels are excluded by default and come back with --include-generic, which is worth doing on wasm, where what the compiler autovectorises rivals the hand-written ones.

Both arm and x86 pick through an analytic LinearCostModel fitted from on-device timings — per Cortex model and Apple generation on arm (cortex_a53_linear, apple_m4_linear, …), per vendor and instruction set on x86 (intel_fma_linear, amd_avx512_linear, …), with separate fits for the matvec role where n == 1. Regenerating one after a kernel change is scripted via tract cost-model gather|fit|regen — see doc/cost-model.md.

Where no model applies — the avx-only tier, and x86 parts that are neither Intel nor AMD — the fallback is pick_mmm in linalg/src/x86_64/mmm.rs, reached after a hand-written table for the small n that some tile shape fits exactly. It maximises scale * m_util * n_util over a KernelChoice table (AVX_CHOICES, FMA_CHOICES, X86_F32_CHOICES), where scale is a kernel's relative throughput once tile-fill is equal. Left at 1.0, sub-1% tile-padding differences alone decide between otherwise-equal kernels, which can route large-N / small-K GEMMs onto a narrower, slower tile. Populate it from an hwbench run at a padding-neutral N (one divisible by every nr, e.g. 120 — the default 512 unfairly favours power-of-two nr): tract hwbench 512,512,120, then normalise the column to the fastest kernel.

Inspecting dispatch

tract selection dumps what mmm selection answers for every machine, not just this one: the declared kernels, each architecture's runnable set, the resolved tier ladder, and per accumulator and shape the tier's answer, the answer selection honours, the pick, and the set that survives retain_best. The shape sweep is hard-coded so that two dumps always compare, which is the point — dump before and after a dispatch change on the same host and diff it, rather than arguing about what moved.

It is not a golden file. A few tiers read the machine rather than the instruction set — the Apple chip generation, the Cortex model behind the a53/a55 boosts, the AMX virtualisation heuristic, TRACT_AMX_BF16 — so one revision dumps differently on two hosts. Diff two builds on one host, and run it on the board when the board is the question.

The sweep is synthetic and each row is filtered by its own instruction set, so it covers every cohort with no knob set — including features this host lacks, which is the point of auditing from a dev machine. TRACT_CPU_ISA steers what a run or an hwbench uses, not this dump.

Nothing conjures a kernel this toolchain never assembled: SVE and SME sit behind build.rs assembler probes, and the DECL rows say built=false for a tree that was only declared — which is why the aarch64 rows above sve2 offer nothing new on a host whose assembler refused it.

tract routines does the same for the single-winner kernels — the element-wise functions, reductions and binaries, which have no shape to weigh and so pick one implementation per machine. It draws the function × machine grid as one coloured cell each: green where the machine runs a kernel written for its own rung of the ladder, white where it runs correct code written for something else — portable Rust, or an architecture kernel from a rung below — red where a portable f16 kernel converts to f32 and back around every operation, blue where a settlement declares the cell closed on purpose, and a dot where nothing implements the pair. Columns are one per rung of every architecture's ladder, named by the rung's nickname under its architecture — fma, avx512, fp16, sme, relaxed — plus the host. Which kernel each cell means is the list under the grid, naming every kernel the host runs and what closed the cells that are closed; --isa takes any machine named above the grid and lists that one instead.

Tuning knobs

A handful of TRACT_* env vars steer kernel selection and CPU detection without recompiling — most usefully TRACT_CPU_ISA for editing the probed instruction set, TRACT_LAZY_IM2COL_MIN_KERNEL / TRACT_LAZY_IM2COL_MAX_EAGER_BYTES for the Conv codegen crossover, and TRACT_CPU_AARCH64_KIND / TRACT_CPU_ARM32_NEON for forcing detection on emulated or misreporting targets. Run tract list-knobs for the full, always-current list, or see cli-recipe.md § Configuration knobs for the annotated highlights.