Skip to content

Commit f4b973c

Browse files
committed
feat(vector): fuse the BBQ rerank distance into SimdRuntime's kernel table
Reranking against 1-bit BBQ candidates materialized a `Vec<f32>` per candidate and per query; at an oversample x ef candidate count that dominated the pass and defeated SIMD. The kernel now reads the centred query straight from the prepared payload bytes: zero allocation per candidate and per query, one pass. It is a `SimdRuntime` field (`l2_bbq`), selected in `detect()` beside the f16/bf16 fused kernels — one dispatch system, one place that decides a tier for a host. The per-tier kernels sit with their siblings in `distance/simd/{avx512,avx2,neon,wasm_simd128}.rs`, the scalar reference in `distance/simd/bbq.rs`, and the rerank codec calls `runtime().l2_bbq`. Evidence: parity against an f64 oracle over dims 0..768 (`distance/simd/bbq.rs`), 18 tests pass; the whole detection table is covered by `dispatch_selects_a_tier_that_matches_the_reference`. The 512-bit tier has no hardware in this fleet and runs under Intel SDE; NEON under qemu; the wasm simd128 kernel compiles for `wasm32-wasip1` with `-C target-feature=+simd128`. The AVX10/256 arm the first revision carried is gone: `AVX512VL` implies `AVX512F`, so no host could reach it while AVX10 detection is absent from `std::arch`. Its test went with it.
1 parent 1ff3551 commit f4b973c

11 files changed

Lines changed: 582 additions & 53 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎nodedb-vector/Cargo.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,3 +44,8 @@ rand = { workspace = true }
4444
tempfile = { workspace = true }
4545
libc = { workspace = true }
4646
nodedb-wal = { workspace = true }
47+
fluxbench = { workspace = true }
48+
49+
[[bench]]
50+
name = "bbq_kernel"
51+
harness = false
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
3+
//! BBQ kernel benches: the zero-copy fused path against the
4+
//! reconstruct-and-measure path it replaces.
5+
//!
6+
//! The unfused bench mirrors the removed path shape: decode the prepared
7+
//! payload into a `Vec<f32>` per candidate, then reconstruct each dimension.
8+
//! (Its residual scale is an approximation: the bench measures the allocation
9+
//! and pass shape, not the codec's exact corrective factor.)
10+
//!
11+
//! Run with: cargo bench -p nodedb-vector --bench bbq_kernel
12+
13+
use fluxbench::bench;
14+
use fluxbench::prelude::*;
15+
use std::hint::black_box;
16+
17+
use nodedb_vector::rerank::codec::{PreparedQuery, RerankCodec};
18+
use nodedb_vector::rerank::codecs::BbqRerank;
19+
20+
/// Installs fluxbench's tracking allocator so the harness reports heap bytes
21+
/// and allocation counts per benchmark — the fused path must show zero
22+
/// allocations, the replaced path must show one `Vec<f32>` per candidate.
23+
#[global_allocator]
24+
static ALLOC: fluxbench::TrackingAllocator = fluxbench::TrackingAllocator;
25+
26+
const OVERSAMPLE: u8 = 4;
27+
const CANDIDATES: usize = 256;
28+
29+
fn det_vec(i: usize, dim: usize) -> Vec<f32> {
30+
(0..dim)
31+
.map(|j| (((i * 31 + j) % 100) as f32 / 100.0) - 0.5)
32+
.collect()
33+
}
34+
35+
fn setup(dim: usize) -> (BbqRerank, PreparedQuery, Vec<Vec<u8>>) {
36+
let vecs: Vec<Vec<f32>> = (0..CANDIDATES).map(|i| det_vec(i, dim)).collect();
37+
let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
38+
let mut codec = BbqRerank::new(dim, OVERSAMPLE);
39+
codec.train(&refs).expect("train");
40+
let prepared = codec.prepare_query(&vecs[0]).expect("prepare_query");
41+
let encoded: Vec<Vec<u8>> = vecs
42+
.iter()
43+
.map(|v| codec.encode(v).expect("encode"))
44+
.collect();
45+
(codec, prepared, encoded)
46+
}
47+
48+
/// The replaced path: decode the prepared payload to `Vec<f32>` per candidate,
49+
/// then reconstruct each dimension. `encoded` carries a 32-byte quant header
50+
/// before the sign bits.
51+
fn unfused_l2(payload: &[u8], encoded: &[u8], dim: usize) -> f32 {
52+
let centered: Vec<f32> = payload[4..]
53+
.as_chunks::<4>()
54+
.0
55+
.iter()
56+
.map(|b| f32::from_le_bytes(*b))
57+
.collect();
58+
let scale = 1.0f32 / (dim as f32).sqrt();
59+
let mut acc = 0.0f32;
60+
for i in 0..dim {
61+
let bit = (encoded[32 + i / 8] >> (7 - (i % 8))) & 1;
62+
let recon = if bit != 0 { scale } else { -scale };
63+
let d = centered[i] - recon;
64+
acc += d * d;
65+
}
66+
acc.sqrt()
67+
}
68+
69+
#[bench(id = "bbq_fused_128", group = "bbq_kernel")]
70+
fn bbq_fused_128(b: &mut Bencher) {
71+
let (codec, prepared, encoded) = setup(128);
72+
b.iter(|| {
73+
let mut acc = 0.0f32;
74+
for e in &encoded {
75+
acc += codec.distance_prepared(&prepared, e).expect("distance");
76+
}
77+
black_box(acc)
78+
});
79+
}
80+
81+
#[bench(id = "bbq_unfused_128", group = "bbq_kernel")]
82+
fn bbq_unfused_128(b: &mut Bencher) {
83+
let (_codec, prepared, encoded) = setup(128);
84+
let payload = match &prepared {
85+
PreparedQuery::Bytes(b) => b.as_slice(),
86+
_ => panic!("bbq prepared form is Bytes"),
87+
};
88+
b.iter(|| {
89+
let mut acc = 0.0f32;
90+
for e in &encoded {
91+
acc += unfused_l2(payload, e, 128);
92+
}
93+
black_box(acc)
94+
});
95+
}
96+
97+
#[bench(id = "bbq_fused_768", group = "bbq_kernel")]
98+
fn bbq_fused_768(b: &mut Bencher) {
99+
let (codec, prepared, encoded) = setup(768);
100+
b.iter(|| {
101+
let mut acc = 0.0f32;
102+
for e in &encoded {
103+
acc += codec.distance_prepared(&prepared, e).expect("distance");
104+
}
105+
black_box(acc)
106+
});
107+
}
108+
109+
#[bench(id = "bbq_unfused_768", group = "bbq_kernel")]
110+
fn bbq_unfused_768(b: &mut Bencher) {
111+
let (_codec, prepared, encoded) = setup(768);
112+
let payload = match &prepared {
113+
PreparedQuery::Bytes(b) => b.as_slice(),
114+
_ => panic!("bbq prepared form is Bytes"),
115+
};
116+
b.iter(|| {
117+
let mut acc = 0.0f32;
118+
for e in &encoded {
119+
acc += unfused_l2(payload, e, 768);
120+
}
121+
black_box(acc)
122+
});
123+
}
124+
125+
fn main() {
126+
if let Err(e) = fluxbench::run() {
127+
eprintln!("Error: {e}");
128+
std::process::exit(1);
129+
}
130+
}

‎nodedb-vector/src/distance/simd/avx2.rs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,3 +114,62 @@ unsafe fn hsum256(v: std::arch::x86_64::__m256) -> f32 {
114114
let sums2 = _mm_add_ss(sums, shuf2);
115115
_mm_cvtss_f32(sums2)
116116
}
117+
118+
use super::bbq::{l2_scalar_from_bytes, recon_scale};
119+
/// Safe entry for `SimdRuntime`; the feature guard lives in `SimdRuntime::detect`.
120+
pub fn l2_bbq(centered: &[u8], packed: &[u8], residual_norm: f32, dim: usize) -> f32 {
121+
// SAFETY: selected only when `detect()` observed this tier's features.
122+
unsafe { l2_bbq_impl(centered, packed, residual_norm, dim) }
123+
}
124+
125+
#[target_feature(enable = "avx2,fma")]
126+
unsafe fn l2_bbq_impl(centered: &[u8], packed: &[u8], residual_norm: f32, dim: usize) -> f32 {
127+
use std::arch::x86_64::*;
128+
129+
let scale = recon_scale(residual_norm, dim);
130+
let scale_v = _mm256_set1_ps(scale);
131+
let mut acc = _mm256_setzero_ps();
132+
133+
let mut i = 0;
134+
while i + 8 <= dim {
135+
// SAFETY: `i + 8 <= dim` and the caller guarantees
136+
// `centered.len() >= dim * 4`, so the 32-byte unaligned load stays in
137+
// bounds; `i / 8` is in range because eight dims consume one byte.
138+
let q = unsafe { _mm256_loadu_ps(centered.as_ptr().add(i * 4).cast::<f32>()) };
139+
let byte = unsafe { *packed.get_unchecked(i / 8) } as usize;
140+
let signs = unsafe { _mm256_load_ps(SIGN_LANES[byte].0.as_ptr()) };
141+
let recon = _mm256_mul_ps(signs, scale_v);
142+
let d = _mm256_sub_ps(q, recon);
143+
acc = _mm256_fmadd_ps(d, d, acc);
144+
i += 8;
145+
}
146+
147+
let sum = unsafe { hsum256(acc) } + l2_scalar_from_bytes(centered, packed, scale, i, dim);
148+
sum.sqrt()
149+
}
150+
151+
/// `±1.0` lane patterns for every packed byte, MSB-first, 32-byte aligned for
152+
/// an aligned load. Precomputed `reverse_bits` mapping (dim `k` → lane `k`).
153+
#[cfg(all(target_arch = "x86_64", target_endian = "little"))]
154+
#[derive(Clone, Copy)]
155+
#[repr(align(32))]
156+
struct Aligned8([f32; 8]);
157+
158+
#[cfg(all(target_arch = "x86_64", target_endian = "little"))]
159+
static SIGN_LANES: [Aligned8; 256] = build_sign_lanes();
160+
161+
#[cfg(all(target_arch = "x86_64", target_endian = "little"))]
162+
const fn build_sign_lanes() -> [Aligned8; 256] {
163+
let mut table = [Aligned8([0.0; 8]); 256];
164+
let mut byte = 0usize;
165+
while byte < 256 {
166+
let mut lane = 0usize;
167+
while lane < 8 {
168+
let bit = (byte >> (7 - lane)) & 1;
169+
table[byte].0[lane] = if bit == 1 { 1.0 } else { -1.0 };
170+
lane += 1;
171+
}
172+
byte += 1;
173+
}
174+
table
175+
}

‎nodedb-vector/src/distance/simd/avx512.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,3 +99,39 @@ unsafe fn ip_impl(a: &[f32], b: &[f32]) -> f32 {
9999
-dot
100100
}
101101
}
102+
103+
use super::bbq::{l2_scalar_from_bytes, recon_scale};
104+
/// Safe entry for `SimdRuntime`; the feature guard lives in `SimdRuntime::detect`.
105+
pub fn l2_bbq(centered: &[u8], packed: &[u8], residual_norm: f32, dim: usize) -> f32 {
106+
// SAFETY: selected only when `detect()` observed this tier's features.
107+
unsafe { l2_bbq_impl(centered, packed, residual_norm, dim) }
108+
}
109+
110+
#[target_feature(enable = "avx512f")]
111+
unsafe fn l2_bbq_impl(centered: &[u8], packed: &[u8], residual_norm: f32, dim: usize) -> f32 {
112+
use std::arch::x86_64::*;
113+
114+
let scale = recon_scale(residual_norm, dim);
115+
let pos = _mm512_set1_ps(scale);
116+
let neg = _mm512_set1_ps(-scale);
117+
let mut acc = _mm512_setzero_ps();
118+
119+
let mut i = 0;
120+
while i + 16 <= dim {
121+
// SAFETY: `i + 16 <= dim` and the caller guarantees
122+
// `centered.len() >= dim * 4`; two packed bytes are in range because
123+
// 16 dims consume exactly two bytes.
124+
// SAFETY: in-bounds per the comment above.
125+
let q = unsafe { _mm512_loadu_ps(centered.as_ptr().add(i * 4).cast::<f32>()) };
126+
let b0 = unsafe { *packed.get_unchecked(i / 8) };
127+
let b1 = unsafe { *packed.get_unchecked(i / 8 + 1) };
128+
let mask: __mmask16 = (b0.reverse_bits() as u16) | ((b1.reverse_bits() as u16) << 8);
129+
let recon = _mm512_mask_blend_ps(mask, neg, pos);
130+
let d = _mm512_sub_ps(q, recon);
131+
acc = _mm512_fmadd_ps(d, d, acc);
132+
i += 16;
133+
}
134+
135+
let sum = _mm512_reduce_add_ps(acc) + l2_scalar_from_bytes(centered, packed, scale, i, dim);
136+
sum.sqrt()
137+
}

0 commit comments

Comments
 (0)