src/stree512.rs
//! S-tree over sorted i64 keys with **B = 16 and AVX-512 node compares** —
//! the "third variant" hypothesis from the 2026-08-10 ragnar brief.
//!
//! Same algorithm as [`crate::stree`] (Ragnar Groot Koerkamp's static search
//! tree, <https://curiouscoding.nl/posts/static-search-tree/>), different
//! geometry: a node holds 16 i64 keys (128 B, two cache lines) and is compared
//! with two `_mm512_cmpgt_epi64_mask` instructions — the same *instruction*
//! count per node as ragnar's original u32/B=16 node, at double the fanout of
//! the AVX2 i64/B=8 tree. log₁₇(8e6) ≈ 5.6 levels vs log₉(8e6) ≈ 7.3.
//!
//! What was traded, stated plainly: a node is now TWO cache lines instead of
//! one. Every level costs two line fills instead of one when cold; the bet is
//! that batching keeps both in flight simultaneously so the *latency* per
//! level is unchanged while the *level count* drops ~25%. The 2026-06-11
//! experiment that rejected i64/B=16 was AVX2 (four compares per node —
//! doubled compute); this variant pays the same two compares as B=8.
//!
//! Machine note (t14s, Strix-class Zen 5 mobile): AVX-512 executes on a
//! 256-bit datapath, so the two 512-bit compares decode to ~four 256-bit µops
//! — compute per node is NOT halved relative to AVX2/B=16 here. Any win below
//! comes from the shorter tree, not from wider ALUs.
//!
//! Internal-node buffer is 64-byte aligned by construction (`NodeBuf16`) —
//! the H4 lesson applied from birth rather than retrofitted.
const B16: usize = 16; // keys per node = branching factor - 1
const MAX64: i64 = i64::MAX;
/// Resolved-once kernel selector — see [`crate::stree_simd`]. This module is
/// the one place where the AVX-512 arm is not a free instruction swap: a B=16
/// i64 node is 128 B, so it needs TWO zmm compares (or four ymm ones), and the
/// geometry itself is what the 2026-06-11 experiment refuted.
#[cfg(target_arch = "x86_64")]
use crate::stree_simd::SimdPath;
// ── shape helpers (same recurrences as stree.rs, B = 16) ──────────────────────
fn blocks16(n: usize) -> usize {
n.div_ceil(B16)
}
fn prev_keys16(n: usize) -> usize {
blocks16(n).div_ceil(B16 + 1) * B16
}
fn height16(n: usize) -> usize {
if n <= B16 { 1 } else { height16(prev_keys16(n)) + 1 }
}
fn layer_size16(mut n: usize, h: usize, height: usize) -> usize {
for _ in h..height - 1 {
n = prev_keys16(n);
}
n
}
fn layer_sizes16(n: usize, height: usize) -> Vec<usize> {
(0..height)
.map(|h| layer_size16(n, h, height).div_ceil(B16))
.collect()
}
#[inline]
fn id_at(data: &[u8], idx: usize, stride: usize) -> i64 {
let off = idx * stride;
i64::from_le_bytes(data[off..off + 8].try_into().unwrap())
}
#[inline]
fn prefetch_ptr<T>(ptr: *const T) {
#[cfg(target_arch = "x86_64")]
unsafe {
std::arch::x86_64::_mm_prefetch(ptr as *const i8, std::arch::x86_64::_MM_HINT_T0);
}
#[cfg(not(target_arch = "x86_64"))]
{
let _ = ptr;
}
}
/// Best-effort `madvise(MADV_HUGEPAGE)` (private copy of stree.rs's helper —
/// same rationale: the tree is the random-access hot structure).
#[inline]
fn advise_hugepage<T>(slice: &[T]) {
#[cfg(target_os = "linux")]
unsafe {
let _ = libc::madvise(
slice.as_ptr() as *mut libc::c_void,
std::mem::size_of_val(slice),
libc::MADV_HUGEPAGE,
);
}
#[cfg(not(target_os = "linux"))]
{
let _ = slice;
}
}
// ── 64-byte-aligned node storage ──────────────────────────────────────────────
//
// `Vec<[i64; 16]>` is align-8 by type; a multi-MB allocation lands at
// glibc's mmap base + 0x10, so every 128-B node would straddle THREE lines
// instead of two. Over-allocate a plain Vec<i64> and start at the first
// 64-byte boundary — safe code, no custom allocator.
pub struct NodeBuf16 {
buf: Vec<i64>,
off: usize, // in i64 units
len: usize, // in nodes
}
impl NodeBuf16 {
fn alloc(n_nodes: usize) -> Self {
// Advise-before-first-touch, same as stree::NodeBuf — a hint after
// faulting waits on khugepaged and delivers nothing (measured).
let len = n_nodes * B16 + 8;
let mut buf: Vec<i64> = Vec::with_capacity(len);
crate::stree::advise_hugepage_raw(buf.as_ptr() as *const u8, len * 8);
buf.resize(len, 0);
let addr = buf.as_ptr() as usize;
let off = ((64 - (addr & 63)) & 63) / 8;
buf[off..off + n_nodes * B16].fill(MAX64);
NodeBuf16 { buf, off, len: n_nodes }
}
}
impl std::ops::Deref for NodeBuf16 {
type Target = [[i64; B16]];
#[inline(always)]
fn deref(&self) -> &[[i64; B16]] {
unsafe {
std::slice::from_raw_parts(self.buf.as_ptr().add(self.off) as *const [i64; B16], self.len)
}
}
}
impl std::ops::DerefMut for NodeBuf16 {
#[inline(always)]
fn deref_mut(&mut self) -> &mut [[i64; B16]] {
unsafe {
std::slice::from_raw_parts_mut(
self.buf.as_mut_ptr().add(self.off) as *mut [i64; B16],
self.len,
)
}
}
}
// ── node compare ──────────────────────────────────────────────────────────────
#[inline(always)]
fn count_lt_scalar16(node: &[i64; B16], q: i64) -> usize {
node.iter().filter(|&&x| x < q).count()
}
/// Two 512-bit compares cover the whole 16-key node; the two 8-bit masks
/// popcount to the branch index. `cmpgt(q, x)` == `x < q`.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
#[inline]
unsafe fn count_lt_avx512(node: &[i64; B16], q: i64) -> usize {
use std::arch::x86_64::*;
unsafe {
let qv = _mm512_set1_epi64(q);
let a = _mm512_loadu_epi64(node.as_ptr());
let b = _mm512_loadu_epi64(node.as_ptr().add(8));
let m1: __mmask8 = _mm512_cmpgt_epi64_mask(qv, a);
let m2: __mmask8 = _mm512_cmpgt_epi64_mask(qv, b);
(m1.count_ones() + m2.count_ones()) as usize
}
}
/// AVX2 fallback for a machine that has no AVX-512 — **oden (Zen 2) is exactly
/// that machine**, and it is the reason this function exists.
///
/// Four 256-bit `cmpgt` per node instead of two 512-bit ones. This is precisely
/// the configuration the 2026-06-11 experiment refuted (two cache lines AND
/// doubled compute), so nothing here is recommended; the point is that one
/// binary must stay *correct and sane* everywhere, and a 16-key scalar
/// `filter().count()` per level is neither. Keeping this arm honest also keeps
/// the A/B meaningful: `ZNIPPY_STREE_SIMD=avx2` now measures what oden would
/// actually run, on t14s.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn count_lt_avx2_16(node: &[i64; B16], q: i64) -> usize {
use std::arch::x86_64::*;
unsafe {
let q_v = _mm256_set1_epi64x(q);
let mut bits = 0u32;
let mut i = 0;
while i < B16 {
// Pairwise reduction — see `crate::stree::count_lt_avx2` for why the
// two compares are packed into one 8-bit mask instead of two
// `vpmovmskb` + two `popcnt` + a `/8` (2026-08-11, measured on oden).
let lo = _mm256_cmpgt_epi64(
q_v,
_mm256_loadu_si256(node.as_ptr().add(i) as *const __m256i),
);
let hi = _mm256_cmpgt_epi64(
q_v,
_mm256_loadu_si256(node.as_ptr().add(i + 4) as *const __m256i),
);
let packed = _mm256_packs_epi32(lo, hi);
bits += (_mm256_movemask_ps(_mm256_castsi256_ps(packed)) as u32).count_ones();
i += 8;
}
bits as usize
}
}
// ── the tree ──────────────────────────────────────────────────────────────────
/// Internal-nodes-only S-tree with B=16/AVX-512 geometry; the sorted key
/// column in `data` (one i64 every `stride` bytes) is the leaf layer, exactly
/// like [`crate::stree::STree64Mmap`].
pub struct STree512Mmap {
tree: NodeBuf16,
offsets: Vec<usize>,
pub count: usize,
pub stride: usize,
}
impl STree512Mmap {
pub fn new_with_stride(data: &[u8], count: usize, stride: usize) -> Self {
assert!(count > 0);
assert!(stride >= 8);
let h = height16(count);
let lsizes = layer_sizes16(count, h);
let ni = h - 1;
let n_blocks: usize = lsizes[..ni].iter().sum();
let mut offsets = Vec::with_capacity(ni);
let mut acc = 0usize;
for &ls in &lsizes[..ni] {
offsets.push(acc);
acc += ls;
}
// Serial fill: the internal layer set is ~1/250th of the keyspace
// (≈4 MB for 8e6 keys) — milliseconds, not worth a fan-out.
let mut tree = NodeBuf16::alloc(n_blocks.max(1));
for lvl in (0..ni).rev() {
let oh = offsets[lvl];
for blk in 0..lsizes[lvl] {
let node = &mut tree[oh + blk];
for j in 0..B16 {
let mut k = blk * (B16 + 1) + j + 1;
for _ in lvl..h - 2 {
k *= B16 + 1;
}
node[j] = if k * B16 < count { id_at(data, k * B16, stride) } else { MAX64 };
}
}
}
advise_hugepage(&tree[..]);
Self { tree, offsets, count, stride }
}
#[inline]
pub fn find_exact(&self, q: i64, data: &[u8]) -> Option<usize> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.find_exact_avx512(q, data) },
SimdPath::Avx2 => return unsafe { self.find_exact_avx2(q, data) },
SimdPath::Scalar => {}
}
self.find_exact_impl(q, data, count_lt_scalar16)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn find_exact_avx512(&self, q: i64, data: &[u8]) -> Option<usize> {
self.find_exact_impl(q, data, |n, qq| unsafe { count_lt_avx512(n, qq) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn find_exact_avx2(&self, q: i64, data: &[u8]) -> Option<usize> {
self.find_exact_impl(q, data, |n, qq| unsafe { count_lt_avx2_16(n, qq) })
}
#[inline(always)]
fn find_exact_impl<F: Fn(&[i64; B16], i64) -> usize>(
&self,
q: i64,
data: &[u8],
cnt: F,
) -> Option<usize> {
let mut k = 0usize;
for &o in &self.offsets {
let jump = cnt(&self.tree[o + k], q);
k = k * (B16 + 1) + jump;
}
let stride = self.stride;
for t in 0..=B16 {
let pos = k * B16 + t;
if pos >= self.count {
break;
}
match id_at(data, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => return Some(pos),
std::cmp::Ordering::Greater => break,
std::cmp::Ordering::Less => {}
}
}
None
}
/// Single-phase pipelined batch: route P queries level-synchronously with
/// next-level prefetch, prefetch the leaf lines at the last internal
/// level, then scan leaves immediately — no routed Vec, no global sort,
/// results written positionally. The in-RAM-section shape.
pub fn lookup_batch_fused<const P: usize>(&self, ids: &[i64], data: &[u8]) -> Vec<Option<usize>> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => {
return unsafe { self.lookup_batch_fused_avx512::<P>(ids, data) };
}
SimdPath::Avx2 => return unsafe { self.lookup_batch_fused_avx2::<P>(ids, data) },
SimdPath::Scalar => {}
}
self.lookup_batch_fused_impl::<P, _>(ids, data, count_lt_scalar16)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn lookup_batch_fused_avx512<const P: usize>(
&self,
ids: &[i64],
data: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused_impl::<P, _>(ids, data, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn lookup_batch_fused_avx2<const P: usize>(
&self,
ids: &[i64],
data: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused_impl::<P, _>(ids, data, |n, q| unsafe { count_lt_avx2_16(n, q) })
}
#[inline(always)]
fn lookup_batch_fused_impl<const P: usize, F: Fn(&[i64; B16], i64) -> usize>(
&self,
ids: &[i64],
data: &[u8],
cnt: F,
) -> Vec<Option<usize>> {
let n = ids.len();
let mut out = vec![None; n];
let n_levels = self.offsets.len();
let stride = self.stride;
let mut i = 0usize;
while i + P <= n {
let mut kk = [0usize; P];
for h in 0..n_levels.saturating_sub(1) {
let o = self.offsets[h];
let o2 = self.offsets[h + 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], ids[i + j]);
kk[j] = kk[j] * (B16 + 1) + jump;
let p = &self.tree[o2 + kk[j]] as *const [i64; B16];
prefetch_ptr(p);
prefetch_ptr((p as *const u8).wrapping_add(64));
}
}
if n_levels > 0 {
let o = self.offsets[n_levels - 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], ids[i + j]);
kk[j] = kk[j] * (B16 + 1) + jump;
// Candidate range is B16+1 records from kk[j]*B16; prefetch
// the three lines that can cover it at stride 8.
let byte = (kk[j] * B16).min(self.count - 1) * stride;
let base = unsafe { data.as_ptr().add(byte.min(data.len() - 1)) };
prefetch_ptr(base);
prefetch_ptr(base.wrapping_add(64));
prefetch_ptr(base.wrapping_add(128));
}
}
for j in 0..P {
let q = ids[i + j];
let k = kk[j];
for t in 0..=B16 {
let pos = k * B16 + t;
if pos >= self.count {
break;
}
match id_at(data, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => {
out[i + j] = Some(pos);
break;
}
std::cmp::Ordering::Greater => break,
std::cmp::Ordering::Less => {}
}
}
}
i += P;
}
while i < n {
out[i] = self.find_exact(ids[i], data);
i += 1;
}
out
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Every kernel this CPU can run must agree. `stree512` is the one variant
/// whose AVX2 arm is brand new (a Zen-2 box previously fell all the way to
/// scalar here), so this is its first correctness evidence.
#[test]
fn every_simd_path_agrees_across_kernels() {
crate::stree_simd::with_each_path(|p| {
let mut state = 0x9e37_79b9_7f4a_7c15u64;
let mut next = move || {
state ^= state >> 12;
state ^= state << 25;
state ^= state >> 27;
state.wrapping_mul(0x2545_F491_4F6C_DD1D)
};
for &n in &[1usize, 16, 17, 289, 30_000] {
let mut keys: Vec<i64> = (0..n).map(|_| next() as i64).collect();
keys.extend_from_slice(&[i64::MIN, i64::MAX, -1, 0]);
keys.sort_unstable();
keys.dedup();
let data = keybuf(&keys);
let t = STree512Mmap::new_with_stride(&data, keys.len(), 8);
for (i, &k) in keys.iter().enumerate() {
assert_eq!(t.find_exact(k, &data), Some(i), "{p:?} n={n} k={k}");
}
let b = t.lookup_batch_fused::<16>(&keys, &data);
for i in 0..keys.len() {
assert_eq!(b[i], Some(i), "{p:?} fused n={n} i={i}");
}
}
});
}
fn keybuf(keys: &[i64]) -> Vec<u8> {
let mut v = Vec::with_capacity(keys.len() * 8);
for k in keys {
v.extend_from_slice(&k.to_le_bytes());
}
v
}
#[test]
fn matches_binary_search_across_sizes() {
let mut state = 0x1234_5678_9abc_def0u64;
let mut next = move || {
state ^= state >> 12;
state ^= state << 25;
state ^= state >> 27;
state.wrapping_mul(0x2545_F491_4F6C_DD1D)
};
for &n in &[1usize, 5, 16, 17, 100, 4096, 100_000] {
let mut keys: Vec<i64> = (0..n).map(|_| (next() >> 1) as i64).collect();
keys.sort_unstable();
keys.dedup();
let data = keybuf(&keys);
let t = STree512Mmap::new_with_stride(&data, keys.len(), 8);
// present keys
let qs: Vec<i64> = (0..200.min(keys.len())).map(|_| keys[(next() as usize) % keys.len()]).collect();
for &q in &qs {
let want = keys.binary_search(&q).ok();
assert_eq!(t.find_exact(q, &data), want, "serial n={n} q={q}");
}
let batch = t.lookup_batch_fused::<8>(&qs, &data);
for (i, &q) in qs.iter().enumerate() {
assert_eq!(batch[i], keys.binary_search(&q).ok(), "fused n={n} i={i}");
}
// absent keys
let absent: Vec<i64> = (0..100).map(|_| (next() >> 1) as i64 | 1).filter(|k| keys.binary_search(k).is_err()).collect();
let miss = t.lookup_batch_fused::<8>(&absent, &data);
assert!(miss.iter().all(Option::is_none), "n={n} absent hit");
}
}
}