notes
about things notes.zzstoatzz.io
notes
Star
Fork
Watch
notes languages ziglang language cpu-kernels.md
4.7 kB · 112 lines
Markdown
at main

cpu-specific kernels #

a SIMD or matrix kernel beyond the build's baseline needs three things: the assembler to accept it, a runtime check that the CPU has it, and a check that it computes what the portable path computes. zig handles each differently per architecture and per backend. learned on 0.16.

the assemblers disagree about features #

LLVM's aarch64 assembler refuses instructions the target lacks:

error: <inline asm>:1:2: instruction requires: sme
        rdsvl x25, #1

so an SME kernel needs the feature enabled for the whole build. zig's native detection does not help here: on Apple M4/M5 it reports apple_m1 with SME off, so the build needs -Dcpu=native+sme+sme2 and the code needs a comptime gate (builtin.cpu.has(.aarch64, .sme)) to keep a portable path for everyone else.

LLVM's x86 assembler accepts AVX-512 BF16 (vdpbf16ps, vcvtne2ps2bf16) in an x86_64_v3 build. one binary can carry the kernel and pick it at runtime, which is what an x86 fleet with mixed CPU generations needs. the runtime check is CPUID feature bits plus XCR0 (the OS must enable opmask and ZMM state: xgetbv & 0xE6 == 0xE6), not CPUID alone.

the self-hosted x86 backend cannot assemble them #

in 0.16, Debug builds for x86_64 Linux use zig's self-hosted backend. its inline-asm parser rejects AT&T memory operands that LLVM accepts:

error: invalid memory operand: '(%[bias])'
error: unknown size: '(%%rsi)'

ReleaseSafe/ReleaseFast use LLVM, and so does Debug on macOS, so a Mac test run and a release build both pass while Linux CI's Debug test step fails. gate such kernels on the backend and let the portable path run elsewhere:

pub const compiled = builtin.cpu.arch == .x86_64 and
    builtin.zig_backend == .stage2_llvm;

then run the tests in ReleaseSafe as well as Debug, or the kernel is never compiled by CI at all.

streaming mode zeroes the vector registers #

smstart/smstop (entering and leaving SME streaming mode) reset the SVE and NEON registers. keep the whole streaming section in one asm block and clobber v0–v31; v8–v15 are callee-saved, and without the clobbers the compiler assumes their values survive the block.

size tiles for the register file you ship to #

a tile that fits aarch64's 32 vector registers spills on AVX2's 16. a four-query attention tile holding four 32-lane query vectors and four 32-lane accumulators needs 32 ymm registers; on an x86_64_v3 build it produced 359 stack spills in the kernel and no speedup. broadcasting the scalar operands from memory (vbroadcastss, no register held) and accumulating in 16-lane slices brought it to 40 with the same FMA count.

counting spills in the disassembly is a seconds-long proxy for x86 performance that works when the development machine is ARM:

zig build -Doptimize=ReleaseFast -Dtarget=x86_64-linux-musl -Dcpu=x86_64_v3
objdump -d --no-show-raw-insn zig-out/bin/x \
  | awk '/<ops.attention>:/{f=1;next} f&&/^$/{exit} f' \
  | rg -c 'ymm[0-9]+.*\(%rsp\)|\(%rsp\).*ymm'

vector builtins that are not vectorized #

@exp on a @Vector lowers to one scalar libcall per lane on these targets; a Cephes-style range reduction and polynomial is several times faster and within a few ulp. when a stage is unexpectedly slow, look for a libcall symbol (ldexp, expf) in the profile before tuning the loop around it.

prove the kernel against an emulation #

keep a plain-zig function that computes the kernel's exact numerics (same rounding, same packing, same accumulation type) and use it three ways:

  • tests on any machine, for the packing and the numerics
  • a load-time self test on the machine that will run the kernel: run a small block through both, and only enable the kernel if they agree (bit for bit for integer kernels, within tolerance for float)
  • accuracy measurement of a reduced-precision path on hardware that lacks it: the emulation is slow but gives the same answers

this is what makes runtime dispatch safe: a CPU where the asm misbehaves falls back to the portable path instead of producing plausible wrong output.

  • build/cross-compilation — -Dtarget without -Dcpu drops to generic x86_64
  • fibers — another place where clobber lists decide correctness
  • instance-sizing — fly hosts differ in CPU generation, which is why the dispatch is at runtime

sources #

  • embedz — src/sme.zig, src/bf16.zig, src/int8.zig, src/ops.zig; SME fmopa, AVX-512 BF16 and AVX2 vpmaddwd kernels with self tests, 2026-09-30