SIMD programming

Deep dive

How it works

In: arrays of f32, u8 or similar, processed in chunks of 4, 8 or 16 lanes · Out: lane-wise results (or one reduced value) · State: none between calls, apart from a cached CPU-feature check for dispatch · Cost: AVX2 peak 32 f32 flops per cycle per core with 2 FMA ports.

1. Lanes, not loops: load a vector, compute lane-wise, store. Cross-lane work (horizontal sums, shuffles, compaction) is the expensive part. Keep it out of the inner loop.
2. Dependency chains: an FMA has 4-cycle latency on current x86 cores. One accumulator gives one FMA per 4 cycles. Use k independent accumulators and add them at the end.
3. Dispatch: compile the kernel with #[target_feature(enable = "...")] so LLVM may emit those instructions inside it. Check with is_x86_feature_detected! before calling it. Calling it on a CPU without the feature is undefined behavior, hence unsafe.
4. Tails: finish the last n % width elements with scalar code, a masked load (AVX-512 _mm512_maskz_loadu_ps, SVE predicates), or a padded buffer.
5. Compile-time alternative: -C target-cpu=x86-64-v3 (AVX2, FMA, BMI2) or x86-64-v4 (AVX-512 F/BW/CD/DQ/VL) makes the whole binary use them, and cfg!(target_feature = "avx2") becomes true. The binary then faults on older CPUs.

Decision table

Code

Dot product with AVX-512, AVX2+FMA, NEON and scalar kernels, selected once at runtime. Checked against an f64 reference for lengths 0-1M. Builds for aarch64-apple-darwin and x86_64-linux.

#[cfg(target_arch = "aarch64")]
use std::arch::aarch64::*;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

/// Portable fallback. 8 independent sums let LLVM vectorize (it may not reorder
/// a single float sum) and keep several FP adds in flight.
pub fn dot_scalar(a: &[f32], b: &[f32]) -> f32 {
    let mut acc = [0.0f32; 8];
    let (ca, cb) = (a.chunks_exact(8), b.chunks_exact(8));
    let tail: f32 = ca.remainder().iter().zip(cb.remainder()).map(|(x, y)| x * y).sum();
    for (x, y) in ca.zip(cb) {
        for k in 0..8 { acc[k] += x[k] * y[k]; }
    }
    acc.iter().sum::<f32>() + tail
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
unsafe fn dot_avx2(a: &[f32], b: &[f32]) -> f32 {
    unsafe {
        let (n, pa, pb) = (a.len(), a.as_ptr(), b.as_ptr());
        let mut acc = [_mm256_setzero_ps(); 4]; // 4 chains hide the 4-cycle FMA latency
        let mut i = 0;
        while i + 32 <= n {
            for k in 0..4 {
                let x = _mm256_loadu_ps(pa.add(i + 8 * k)); // unaligned load: same speed if no line split
                let y = _mm256_loadu_ps(pb.add(i + 8 * k));
                acc[k] = _mm256_fmadd_ps(x, y, acc[k]);
            }
            i += 32;
        }
        let v = _mm256_add_ps(_mm256_add_ps(acc[0], acc[1]), _mm256_add_ps(acc[2], acc[3]));
        // Horizontal sum 8 -> 4 -> 2 -> 1 lanes.
        let s = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps::<1>(v));
        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
        let s = _mm_add_ss(s, _mm_shuffle_ps::<0b01>(s, s));
        _mm_cvtss_f32(s) + dot_scalar(&a[i..], &b[i..])
    }
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn dot_avx512(a: &[f32], b: &[f32]) -> f32 {
    unsafe {
        let (n, pa, pb) = (a.len(), a.as_ptr(), b.as_ptr());
        let (mut acc0, mut acc1) = (_mm512_setzero_ps(), _mm512_setzero_ps());
        let mut i = 0;
        while i + 32 <= n {
            acc0 = _mm512_fmadd_ps(_mm512_loadu_ps(pa.add(i)), _mm512_loadu_ps(pb.add(i)), acc0);
            acc1 = _mm512_fmadd_ps(_mm512_loadu_ps(pa.add(i + 16)), _mm512_loadu_ps(pb.add(i + 16)), acc1);
            i += 32;
        }
        while i < n {
            // Masked tail: lanes past the end are neither loaded nor faulted.
            let left = (n - i).min(16);
            let m: __mmask16 = if left == 16 { 0xFFFF } else { (1u16 << left) - 1 };
            let x = _mm512_maskz_loadu_ps(m, pa.add(i));
            let y = _mm512_maskz_loadu_ps(m, pb.add(i));
            acc0 = _mm512_fmadd_ps(x, y, acc0);
            i += 16;
        }
        _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1))
    }
}

#[cfg(target_arch = "aarch64")]
fn dot_neon(a: &[f32], b: &[f32]) -> f32 {
    // NEON is part of the AArch64 baseline: no runtime detection needed.
    unsafe {
        let (n, pa, pb) = (a.len(), a.as_ptr(), b.as_ptr());
        let mut acc = [vdupq_n_f32(0.0); 4];
        let mut i = 0;
        while i + 16 <= n {
            for k in 0..4 {
                // vfmaq_f32(acc, x, y) = acc + x * y, fused
                acc[k] = vfmaq_f32(acc[k], vld1q_f32(pa.add(i + 4 * k)), vld1q_f32(pb.add(i + 4 * k)));
            }
            i += 16;
        }
        let v = vaddq_f32(vaddq_f32(acc[0], acc[1]), vaddq_f32(acc[2], acc[3]));
        vaddvq_f32(v) + dot_scalar(&a[i..], &b[i..]) // vaddvq: horizontal add of 4 lanes
    }
}

/// Pick the best kernel once; call through the pointer in hot loops.
pub fn select_dot() -> fn(&[f32], &[f32]) -> f32 {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx512f") {
            return |a, b| unsafe { dot_avx512(a, b) }; // safe: feature checked above
        }
        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
            return |a, b| unsafe { dot_avx2(a, b) };
        }
    }
    #[cfg(target_arch = "aarch64")]
    return dot_neon;
    #[allow(unreachable_code)]
    dot_scalar
}

pub fn dot(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len()); // the kernels read b[..a.len()] unchecked
    select_dot()(a, b)
}

On an M4 Pro (data in L1), the NEON kernel does about 28 GFLOP/s and the auto-vectorized scalar fallback about 18.

Complexity and constants

• Peak per core per cycle (f32): AVX2 with 2 FMA ports = 2 × 8 × 2 = 32 flops. AVX-512 with 2 FMA ports = 64.
• Loads: 2-3 per cycle on current x86. A dot product needs 2 loads per FMA, so it is load-bound at L1 before it is FMA-bound.
• Lane-crossing shuffles (vpermps, vpermd) cost order of 3 cycles latency. In-lane shuffles cost 1. Gathers are order of one element per cycle: not a free random load.
• AVX-512 frequency: Skylake-SP lowers its clock under heavy 512-bit code. Ice Lake and later server parts much less. Zen 4 runs 512-bit ops on 256-bit units without a clock drop.

Efficiency

• Cost: peak per core and cycle is 32 f32 flops with AVX2 and 2 FMA ports, 64 with AVX-512. A dot product needs 2 loads per FMA, so it is load-bound at L1 (2-3 loads per cycle) before FMA-bound. Bound: FMA ports or L1 loads in cache, memory bandwidth beyond it (roofline).
• Levers:
  • k independent accumulators — one FMA per 4 cycles → near 2 per cycle with about 8 accumulators — float reductions
  • -C target-cpu=x86-64-v3 or native — whole-program AVX2/FMA, no dispatch cost — servers you control
  • runtime dispatch with is_x86_feature_detected! — one binary reaches AVX2 or AVX-512 — unknown x86 fleets; many Intel client CPUs lack AVX-512
  • compare plus movemask on bytes — gigabytes per second for JSON parsing (Langdale and Lemire) — parsing, search
• Naive → standard → state of the art:
• naive: scalar loop, one accumulator · 0.5 flop per cycle (2 flops per 4-cycle FMA)
• standard: AVX2 with 8 accumulators · up to 32 flops per cycle per core
• state of the art: AVX-512 with mask tails and runtime dispatch · up to 64 flops per cycle, capped by L1 loads or memory bandwidth

Implementation notes

• Rust: x86 SSE/AVX/AVX2/FMA intrinsics are stable, AArch64 NEON since 1.59, AVX-512 since 1.89. std::simd is still nightly-only.
• Inside a #[target_feature] function, recent Rust treats intrinsics without pointer arguments as safe. Loads and stores through raw pointers stay unsafe.
• is_x86_feature_detected! caches its result, but resolve the kernel once per call site anyway (fn pointer or OnceLock). Branching per element prevents inlining.
• Unaligned loads (loadu, vld1q) are as fast as aligned ones unless they cross a cache line. Align hot buffers to 64 B (128 B on Apple) to avoid splits.
• Mixing legacy SSE and AVX encodings costs a state transition. Compilers insert vzeroupper. Hand-written asm must do it too.
• AArch64 NEON idioms: vaddvq_f32 (horizontal add), vqtbl1q_u8 (16-byte table lookup), vshrn_n_u16 + vget_lane_u64 as a movemask substitute.
• Integer dot products: AVX-512 VNNI vpdpbusd and NEON sdot/udot (vdotq_s32, needs the dotprod feature) give int8 → int32 in one instruction. They are the basis of quantized inference.

Pitfalls

• Different summation order gives different float results per ISA. Tests need a tolerance, and outputs must not depend on bit-exact sums (lockstep, hashes).
• Reading past the slice end in the tail "because it's in the same page": UB in Rust, and a real fault at a page boundary.
• Forgetting fma when enabling avx2: they are separate CPUID bits.
• Measuring with data that fits in L1, then deploying on DRAM-sized arrays.
• Denormals: some x86 operations are order of 100× slower on subnormal inputs. Flush them to zero (MXCSR FTZ/DAZ) only if the semantics allow it.

Variants and alternatives

In production

simdjson, Hyperscan, Rust's memchr crate (SSE2/AVX2/NEON), Stockfish NNUE (AVX2/AVX-512/NEON paths), ClickHouse, and every BLAS (OpenBLAS, MKL, Apple Accelerate).

Implementations

Sources

• Intel, Intel Intrinsics Guide.
• Arm, Arm C Language Extensions (ACLE) and Neon Intrinsics Reference. ACLE
• Agner Fog, Instruction tables and Optimizing software in C++. tables · optimizing
• Geoff Langdale, Daniel Lemire, Parsing Gigabytes of JSON per Second, VLDB Journal, 2019. paper · code
• Wojciech Muła, Nathan Kurz, Daniel Lemire, Faster Population Counts Using AVX2 Instructions, The Computer Journal, 2018. paper