Axle v0.14.1

SIMD CPU dispatch

Same binary on a 2010 Westmere CPU and a 2023 Zen 4 — Axle programs that use SIMD pick the best implementation of each kernel at process startup, without any per-call overhead in user code.

This page explains how. The user-facing API is documented in SIMD and auto-vectorisation; here we look at the dispatch model under it.

Newer CPUs add wider and smarter SIMD instructions over time (SSE → AVX2 → AVX-512 on x86-64; NEON on ARM). A binary that uses an AVX2 instruction runs faster on a chip that has it — but crashes on one that doesn’t, with an illegal-instruction fault. Runtime feature detection solves the dilemma : the program asks the CPU what it supports (on x86-64 this is the cpuid instruction, a single query that reports the available feature bits), then, at startup, wires each SIMD operation to the best implementation the host can actually run. One binary, best kernel everywhere.

The problem

Modern x86-64 CPUs span four microarchitecture levels (Red Hat / glibc convention):

LevelFeatures added
Baseline (v1)SSE2 (required by x86-64 itself)
v2SSE3, SSSE3, SSE4.1, SSE4.2, POPCNT
v3AVX, AVX2, FMA, BMI1, BMI2
v4AVX-512 F + BW + CD + DQ + VL

A binary compiled with -march=x86-64-v3 (assuming AVX2 + FMA) runs ~2× faster than -march=x86-64 for SIMD-heavy code — but crashes with SIGILL on an older CPU. The usual answer is to ship per-level binaries (Steam Runtime, RHEL multi-arch packages), but that doubles or triples deployment artefacts.

Axle wants one binary that works on every x86-64 host AND uses the best SIMD on each one.

The model

For each SIMD kernel that actually benefits from CPU-specific instructions (FMA and horizontal reductions), the runtime provides:

  • N variants with different #[target_feature] annotations (baseline scalar loop, AVX2+FMA inline intrinsics, future AVX-512).
  • One trampoline with a stable __axle_simd_<op>_<shape> symbol exposed to generated code.
  • One static AtomicPtr<()> slot per (op, shape) that holds the selected variant’s function pointer.

At process startup, a CPU probe via cpuid selects the highest level the host supports, and install() stores the matching variant into every slot with Ordering::Release. From then on, every SIMD call from generated code is one Ordering::Relaxed load + one indirect call — predicted perfectly by the branch-target predictor after the first call.

The pattern is the same one used by glibc, openssl, ffmpeg, Rust’s multiversion crate, and every modern game engine. On x86-64 it costs 1-2 cycles per call ; usually hidden by instruction-level parallelism behind the surrounding code.

ABI

Kernel trampolines take pointer arguments rather than vector-typed register arguments :

__axle_simd_fma_f64x4(a*, b*, c*, out*)        — 4 pointers
__axle_simd_reduce_fadd_f64x4(in*, out*)       — 2 pointers

This is intentional. x86-64 SysV passes __m256 in YMM registers ; Windows x64 passes it as a hidden pointer. Two different ABIs to maintain per kernel. The pointer-based ABI is identical on every target — and the caller already allocates each operand in an alloca (the user’s let v : f64x4 = …; lowers to one), so the spill is paid for by the source-level binding, not by the call boundary.

When the dispatch runs

__axle_simd_init is called from the generated main as one of the first prologue statements — but only when the program actually uses SIMD. The compiler tracks whether any vector builtin appears in the program and emits the init call exactly when it does. Binaries that never touch SIMD never reference __axle_simd_init in the produced IR — they don’t pay the ~40 cycles of cpuid probing or the 16 atomic stores at startup.

What an FMA call looks like in the emitted IR

When the compiler lowers a.fma(b, c) (an f32x8), the produced LLVM IR is:

%a_ptr   = alloca <8 x float>
%b_ptr   = alloca <8 x float>
%c_ptr   = alloca <8 x float>
%out_ptr = alloca <8 x float>
store <8 x float> %a, ptr %a_ptr
store <8 x float> %b, ptr %b_ptr
store <8 x float> %c, ptr %c_ptr
call void @__axle_simd_fma_f32x8(ptr %a_ptr, ptr %b_ptr,
                                 ptr %c_ptr, ptr %out_ptr)
%result = load <8 x float>, ptr %out_ptr

The four allocas land in the entry block ; LLVM’s mem2reg and SROA passes typically erase three of them (the inputs) when the operands originate from SSA values that don’t need a memory roundtrip — the trampoline ABI demands pointer args but the compile-time cost on the caller side is minimal.

Reductions follow the same shape with one input slot plus a scalar-sized output slot.

What’s not multi-versioned

splat, the literal-list constructor, load, store, select, shuffle, eq, and movemask are pure LLVM vector IR ops — they don’t have CPU-specific intrinsics that would benefit from runtime dispatch. They emit inline, and LLVM’s target-feature legaliser picks the right instruction sequence at compile time.

The same path applies to FMA and reductions for shapes that aren’t in the dispatch table (f32x16, f64x8, integer reductions) — they get the inline LLVM vector intrinsic and rely on the compile-time target-features baseline.

Why pointers instead of register-passed vectors

x86-64 SysV passes __m256 in YMM registers; Windows x64 passes it as a hidden pointer. Two different ABIs to maintain per kernel. The pointer-based ABI is identical on every target — and each operand is already in a stack slot (the user’s let v : f64x4 = …; lowers to one), so the spill is paid for by the source-level binding, not by the call boundary.

Adding a new dispatched shape

Adding a new lane shape to the runtime dispatch is intentionally mechanical: declare the kernel variants (baseline + the higher tiers you target), wire them into the install table, and tell the compiler that this shape now has a runtime kernel. Every remaining piece — symbol naming, init gating, call-site lowering — is shared with the existing shapes.


Limitations

LimitWhy it holds
Two kernels per shape, not one per levelthe table holds a portable scalar variant and an AVX2+FMA one; a v4 host runs the AVX2 kernel
Only FMA and the floating reductions are dispatchedsplat, load, store, select, shuffle, eq and movemask are plain LLVM vector IR, legalised at compile time
max / min reductions have no AVX2 variantthe AVX2 kernels exist for the sums; the maximum and minimum run the portable loop on every host
Outside x86-64 the probe answers Baselineno per-architecture level is wired for NEON, so the portable kernel is the installed one
A program that only annotates loops does not start the dispatchthe init call is emitted when a vector builtin appears in the program; a @vectorize loop is compiled for the build’s target-features
The dispatched shapes are f32x4, f32x8, f64x2, f64x4a wider or narrower shape takes the inline LLVM intrinsic on the compile-time baseline

See also

simdcpu-dispatchperformance