src/stree_simd.rs
//! **One-time SIMD capability resolution for the whole stree family.**
//!
//! One compiled artifact has to run optimally on both machines in the
//! constellation: t14s (Zen 5 mobile, `avx512f/bw/vl/dq`) and oden (Zen 2,
//! **no AVX-512 at all**). That means no 512-bit instruction may be reachable
//! unless CPUID said so at run time — and it means the check must not sit in a
//! lookup loop. `is_x86_feature_detected!` is cheap (a cached `AtomicU32` in
//! std) but it is not free, and `stree::batch_stream` used to pay it once per
//! P-query chunk.
//!
//! So the resolution happens **once**, on the first query, and is then a plain
//! relaxed byte load:
//!
//! ```text
//! detected() OnceLock<SimdPath> pure CPUID; the hardware truth, immutable
//! auto_path() policy what `auto` means; <= detected(), by choice
//! path() AtomicU8 the ACTIVE selection, always <= detected()
//! force(p) clamps to detected() the A/B knob; can never enable a missing ISA
//! ```
//!
//! `force` clamping to `detected()` is the safety property that matters: a
//! benchmark (or a stray env var) asking for AVX-512 on oden gets AVX2 back and
//! runs correctly, because [`SimdPath`] is ordered `Scalar < Avx2 < Avx512` and
//! the selection is a `min`. No caller can talk the dispatcher into executing an
//! instruction the CPU does not have.
//!
//! ## What this does NOT dispatch: geometry
//!
//! Instruction choice is a property of the CPU **reading** the tree, so it is
//! runtime-dispatchable — that is this module's whole job. Node *geometry*
//! (B=8 i64 vs B=16 i64 vs hi32's 16×u32) is a property of the **tree**, and a
//! reader cannot dispatch over it.
//!
//! For the `*Mmap` trees here that is a non-issue, and it is worth stating why
//! rather than leaving it to be rediscovered: **the internal nodes are never
//! serialized.** znippy's `__gunnar_oid__` section stores only the sorted i64
//! key column (`oid_index.rs`: "the internal nodes are not in this section"),
//! and `new_with_stride` rebuilds the tree from that column at open time. So
//! the geometry is chosen by the *reader*, per open, on the reader's own CPU —
//! an archive written on t14s carries no AVX-512-shaped decision into oden.
//!
//! If a future format ever does persist nodes, the geometry must be recorded in
//! the header and honoured by every reader; it must never be inferred from the
//! writer's CPUID. The 2026-06-11 experiment already refuted the configuration
//! that mistake would produce (i64 B=16 read with AVX2: two cache lines per
//! node, four 256-bit compares) — see [`crate::stree512`]'s header.
//!
//! ## Env override
//!
//! `ZNIPPY_STREE_SIMD=scalar|avx2|avx512|auto` (default `auto`). Read exactly
//! once, during resolution. Unparseable values fall back to `auto` — this is a
//! benchmarking knob, not a configuration surface worth failing a lookup over.
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::OnceLock;
/// Widest node-compare kernel a caller may use. **Ordered**: `Scalar < Avx2 <
/// Avx512`, which is what makes [`force`]'s clamp a `min`.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
#[repr(u8)]
pub enum SimdPath {
/// Portable fallback. Also what a non-x86_64 build always gets.
Scalar = 0,
/// 256-bit `_mm256_cmpgt_*`. Two compares for a 64-byte node.
Avx2 = 1,
/// 512-bit `_mm512_cmpgt_*_mask`. **One** compare for a 64-byte node.
Avx512 = 2,
}
impl SimdPath {
#[inline(always)]
fn from_u8(v: u8) -> Self {
match v {
2 => SimdPath::Avx512,
1 => SimdPath::Avx2,
_ => SimdPath::Scalar,
}
}
/// `"scalar"` / `"avx2"` / `"avx512"` — the spelling [`force_from_str`]
/// accepts, so a report can round-trip what it measured.
pub fn as_str(self) -> &'static str {
match self {
SimdPath::Scalar => "scalar",
SimdPath::Avx2 => "avx2",
SimdPath::Avx512 => "avx512",
}
}
}
/// Sentinel for "the atomic has not been initialised from `detected()` yet".
const UNRESOLVED: u8 = u8::MAX;
static DETECTED: OnceLock<SimdPath> = OnceLock::new();
static ACTIVE: AtomicU8 = AtomicU8::new(UNRESOLVED);
/// Pure CPUID, cached forever. The ceiling on everything else here.
///
/// Requires the full `f/bw/vl/dq` set for [`SimdPath::Avx512`], not just
/// `avx512f`: the stree kernels only need `f`, but reporting a coarser
/// capability than the constellation's other AVX-512 users (`vann`) check for
/// would make the two disagree about what "this machine has AVX-512" means.
pub fn detected() -> SimdPath {
*DETECTED.get_or_init(|| {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx512f")
&& is_x86_feature_detected!("avx512bw")
&& is_x86_feature_detected!("avx512vl")
&& is_x86_feature_detected!("avx512dq")
{
return SimdPath::Avx512;
}
if is_x86_feature_detected!("avx2") {
return SimdPath::Avx2;
}
}
SimdPath::Scalar
})
}
/// The **active** path — what every stree entry point branches on.
///
/// Hot-path cost after the first call: one relaxed byte load and a compare.
/// The initialising branch is `#[cold]`, so the resolved path is the straight
/// line through this function.
#[inline(always)]
pub fn path() -> SimdPath {
match ACTIVE.load(Ordering::Relaxed) {
UNRESOLVED => resolve_cold(),
v => SimdPath::from_u8(v),
}
}
/// **The `auto` policy** — what an unpinned process runs, and the one place a
/// per-target narrowing belongs.
///
/// `auto` has always meant *"the widest kernel this CPU can execute"*, which is
/// [`detected`]. That is only the right answer while every kernel is actually
/// faster than the one below it, and on 2026-08-11 that turned out to be false
/// on oden: the hand-written AVX2 node compare lost to the "scalar" fallback on
/// all six stree bench arms — up to 1.58x on `stree_1gb`'s serial path — and
/// `auto` picked it anyway. **The dispatch was not the defect** (the mechanism
/// was a fat reduction tail in `count_lt_avx2`, since fixed — see
/// [`crate::stree::count_lt_avx2`]); but nothing in this module *said* what
/// `auto` promises, so nothing could notice it had stopped being true.
///
/// It says it here now, as a property with a name:
///
/// > **`auto` must never select a kernel that is slower than a narrower one on
/// > the CPU it is running on.**
///
/// Measured verdict as of 2026-08-11, best-of-N interleaved, six stree arms:
///
/// | machine | ceiling | scalar | avx2 | avx512 | `auto` picks |
/// |---|---|---|---|---|---|
/// | oden (Zen 2, znver2 baseline) | avx2 | fast — LLVM auto-vectorises the fallback well | ties/wins by 0–8 % after the fix | n/a | avx2 ✔ |
/// | t14s (Zen 5, znver5 baseline) | avx512 | **3x slower** — LLVM extracts the `k`-mask bit by bit | 1.3–1.5x slower than avx512 | best | avx512 ✔ |
///
/// So the policy is still "the ceiling", and that is now a *measured* claim
/// rather than an assumption. A narrowing (a `matches!` on a microarchitecture,
/// a size threshold, a `cfg!(target_feature)` test) goes **here**, and nowhere
/// else: this returns a value that [`force`] then clamps to [`detected`] again,
/// so it can only ever narrow an already-safe choice — it can never widen one.
///
/// Note the trap this function exists to document: **"scalar" is not scalar.**
/// Under `-C target-cpu=native` the portable `count_lt_scalar` auto-vectorises
/// to whatever the build's baseline allows, and the quality of what LLVM emits
/// varies enormously between targets (excellent on znver2, pathological on
/// znver5). A policy keyed on "is the fallback good here?" is therefore not
/// knowable at compile time, which is why the fix for the oden regression was
/// to make the kernel better rather than to teach `auto` to avoid it.
pub fn auto_path() -> SimdPath {
detected()
}
#[cold]
#[inline(never)]
fn resolve_cold() -> SimdPath {
let ceiling = detected();
let want = match std::env::var("ZNIPPY_STREE_SIMD").ok().as_deref() {
Some("scalar") => SimdPath::Scalar,
Some("avx2") => SimdPath::Avx2,
Some("avx512") => SimdPath::Avx512,
// "auto", unset, or anything unrecognised.
_ => auto_path(),
};
let eff = want.min(ceiling);
// Racing callers all compute the same value from the same inputs, so a
// plain store is enough — there is no "first writer wins" to protect.
ACTIVE.store(eff as u8, Ordering::Relaxed);
eff
}
/// Pin the active path for an A/B measurement; returns what actually took
/// effect after clamping to [`detected`].
///
/// Intended for benchmarks and for a startup-time policy decision — **not**
/// for flipping mid-run. It is memory-safe at any time (the clamp guarantees
/// the ISA exists, and each kernel is behind its own `#[target_feature]`
/// wrapper), but a concurrent reader may observe either value, so a mid-flight
/// batch could be timed across two kernels.
pub fn force(req: SimdPath) -> SimdPath {
let eff = req.min(detected());
ACTIVE.store(eff as u8, Ordering::Relaxed);
eff
}
/// [`force`] from the same spellings `ZNIPPY_STREE_SIMD` accepts. `None` on an
/// unrecognised name, so a bench harness can tell a typo from a clamp.
///
/// `"auto"` is **not** a synonym for `"avx512"` here, even though the two
/// resolve to the same thing on a machine whose ceiling is AVX-512: it routes
/// through [`auto_path`], so a bench that asks for `auto` measures the policy
/// an unpinned process would actually run, not a pin that happens to coincide
/// with it today.
pub fn force_from_str(s: &str) -> Option<SimdPath> {
let p = match s {
"scalar" => SimdPath::Scalar,
"avx2" => SimdPath::Avx2,
"avx512" => SimdPath::Avx512,
"auto" => auto_path(),
_ => return None,
};
Some(force(p))
}
/// One-line capability banner for bench output / `# machine` comments.
pub fn describe() -> String {
format!(
"stree_simd: active={} detected={}",
path().as_str(),
detected().as_str()
)
}
/// Serialises the tests that pin a path — [`force`] is process-global, and
/// `cargo test` runs test fns concurrently. Held for the whole sweep, not per
/// path, so no other forcing test can interleave a different kernel underneath.
///
/// Non-forcing tests are deliberately *not* required to take it: they must be
/// correct under whichever path happens to be pinned, which is exactly the
/// property being asserted.
#[cfg(test)]
pub(crate) static TEST_PATH_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
/// Run `f` once per kernel this CPU can actually execute, then restore `auto`.
///
/// The gate is `p <= detected()`, so on oden this runs Scalar+Avx2 and simply
/// never reaches the AVX-512 arm — the same clamp the production dispatch uses,
/// exercised by the same call.
#[cfg(test)]
pub(crate) fn with_each_path(f: impl Fn(SimdPath)) {
let _g = TEST_PATH_LOCK.lock().unwrap_or_else(|e| e.into_inner());
for p in [SimdPath::Scalar, SimdPath::Avx2, SimdPath::Avx512] {
if p > detected() {
continue;
}
assert_eq!(force(p), p);
f(p);
}
force(auto_path());
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn active_never_exceeds_detected() {
let _g = TEST_PATH_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let d = detected();
assert!(path() <= d, "active {:?} exceeds detected {:?}", path(), d);
// Asking for more than the hardware has must clamp, not enable.
assert_eq!(force(SimdPath::Avx512), d.min(SimdPath::Avx512));
assert_eq!(force(SimdPath::Scalar), SimdPath::Scalar);
assert_eq!(force(SimdPath::Avx2), d.min(SimdPath::Avx2));
// Restore for any other test in this binary.
force(auto_path());
}
#[test]
fn str_round_trip() {
let _g = TEST_PATH_LOCK.lock().unwrap_or_else(|e| e.into_inner());
for p in [SimdPath::Scalar, SimdPath::Avx2, SimdPath::Avx512] {
assert_eq!(force_from_str(p.as_str()), Some(p.min(detected())));
}
assert_eq!(force_from_str("sse2"), None);
// "auto" is its own request, routed through the policy — not a third
// spelling of "avx512".
assert_eq!(force_from_str("auto"), Some(auto_path()));
force(auto_path());
}
#[test]
fn ordering_is_capability_ordering() {
assert!(SimdPath::Scalar < SimdPath::Avx2);
assert!(SimdPath::Avx2 < SimdPath::Avx512);
}
/// The `auto` policy is a *narrowing* of [`detected`], never a widening —
/// whatever [`auto_path`] grows into later. Cheap, hermetic, and the thing
/// that keeps the safety property true if someone edits the policy.
#[test]
fn auto_is_a_narrowing_of_detected() {
let _g = TEST_PATH_LOCK.lock().unwrap_or_else(|e| e.into_inner());
assert!(
auto_path() <= detected(),
"auto_path() {:?} exceeds detected() {:?} — the policy may only narrow",
auto_path(),
detected()
);
// An unpinned process must land exactly on the policy.
force(auto_path());
assert_eq!(path(), auto_path());
}
/// **The regression that shipped, as an assertion.** `auto` may not be
/// slower than a narrower kernel on the CPU running the test.
///
/// `#[ignore]` on purpose: this is a *timing* assertion, and a hermetic
/// suite that goes red because another agent started a build is worse than
/// no assertion at all. It is meant to be run deliberately, on a quiet box,
/// as part of a bench session:
///
/// ```text
/// cargo test --release -- --ignored auto_is_not_slower_than_a_narrower_kernel --nocapture
/// ```
///
/// The workload is a 128 MB `STree64` walked by the un-pipelined
/// `find_exact` — deliberately the shape that has *no* software prefetch, so
/// a fat node-compare shows up undiluted, and deliberately large enough to
/// leave cache. Size matters to the bar: the 2026-08-11 oden regression
/// measured 1.06x at 8 MB, **1.33x at 128 MB** and 1.58x at 1 GB, because
/// the cost is an out-of-order-window effect that only bites once the tree
/// is RAM-resident. 128 MB is the smallest arm that clears the 1.25x bar
/// while still building in ~0.2 s. Kernels are **interleaved** and reported
/// best-of-N, the same method the box demands.
#[test]
#[ignore = "timing assertion — run deliberately on a quiet box, see the doc comment"]
fn auto_is_not_slower_than_a_narrower_kernel() {
use std::hint::black_box;
use std::time::Instant;
let _g = TEST_PATH_LOCK.lock().unwrap_or_else(|e| e.into_inner());
// 16 M i64 keys = a 128 MB tree: past LLC on both constellation
// machines, which is where the regression this guards against shows up.
let n = 1usize << 24;
let keys: Vec<i64> = {
let mut seed = 0xD15EA5E_u64;
let mut cur = 1i64;
(0..n)
.map(|_| {
seed ^= seed >> 12;
seed ^= seed << 25;
seed ^= seed >> 27;
cur += 1 + (seed.wrapping_mul(0x2545F4914F6CDD1D) % 15) as i64;
cur
})
.collect()
};
let queries: Vec<i64> = {
let mut seed = 0xC0FFEE_u64;
(0..1usize << 18)
.map(|_| {
seed ^= seed >> 12;
seed ^= seed << 25;
seed ^= seed >> 27;
keys[(seed.wrapping_mul(0x2545F4914F6CDD1D) as usize) % n]
})
.collect()
};
let tree = crate::stree::STree64::new(&keys);
let paths: Vec<SimdPath> = [SimdPath::Scalar, SimdPath::Avx2, SimdPath::Avx512]
.into_iter()
.filter(|&p| p <= detected())
.collect();
let mut best = vec![f64::MAX; paths.len()];
// Interleaved: the kernel is the INNER loop, so drift hits all of them
// equally. rep 0 is a warmup and is discarded.
for rep in 0..5 {
for (i, &p) in paths.iter().enumerate() {
assert_eq!(force(p), p);
let t = Instant::now();
let mut acc = 0usize;
for &q in &queries {
acc ^= black_box(tree.find_exact(q)).unwrap_or(0);
}
black_box(acc);
let ns = t.elapsed().as_nanos() as f64 / queries.len() as f64;
if rep > 0 && ns < best[i] {
best[i] = ns;
}
}
}
force(auto_path());
for (i, &p) in paths.iter().enumerate() {
eprintln!(" {:>6}: {:6.2} ns/query", p.as_str(), best[i]);
}
let auto_i = paths.iter().position(|&p| p == auto_path()).unwrap();
for (i, &p) in paths.iter().enumerate() {
if p >= auto_path() {
continue;
}
assert!(
best[auto_i] <= best[i] * 1.25,
"auto ({}) is {:.2} ns/query but the narrower {} is {:.2} — \
auto must not select a kernel that loses on this CPU",
auto_path().as_str(),
best[auto_i],
p.as_str(),
best[i]
);
}
}
}