From 8091e2627a979bd102e93fe8ec544ae2c0110afc Mon Sep 17 00:00:00 2001 From: Nelson Spence Date: Mon, 15 Jun 2026 23:08:31 -0500 Subject: [PATCH 1/5] feat(beir-bench): within-query-threaded sign baseline + DistL2 HNSW Add `sign-rq2-threaded`: a within-query-threaded SignBitmap -> RankQuant b=2 method. Doc-major sign codes are scanned in parallel over doc-stripes with hardware VPOPCNTDQ, then an exact-m candidate selection (count desc, id asc tie-break) reproduces the serial `sign-rq2` set byte-for-byte (top-100 overlap 1.0000) -- isolating the within-query threading speedup from any quality change. Harden HNSW to DistL2 (was DistDot). Embeddings are unit-normalized, so min-L2 == max-dot == max-cosine (identical neighbors), but DistL2 avoids anndists' DistDot `1-dot` distance assertion that panics on near-duplicate pairs whose float dot rounds just past 1.0 -- rare at 171K docs, frequent near 1M. Score becomes `-distance` (nearer = higher), preserving nearest-first order. beir_eval: add the family_meta entry for ordvec-sign-rq2-threaded. Signed-off-by: Nelson Spence --- benchmarks/beir-bench/src/main.rs | 309 +++++++++++++++++++++++++++++- benchmarks/beir/beir_eval.py | 6 + 2 files changed, 309 insertions(+), 6 deletions(-) diff --git a/benchmarks/beir-bench/src/main.rs b/benchmarks/beir-bench/src/main.rs index a6d63e45..13eb7520 100644 --- a/benchmarks/beir-bench/src/main.rs +++ b/benchmarks/beir-bench/src/main.rs @@ -863,8 +863,13 @@ fn main() { write_topk, &mut timing_writer, ), + "sign-rq2-threaded" => run_sign_threaded( + cfg.candidates, corpus, &queries, dim, n_docs, n_queries, cfg.top_k, cfg.batch, + threads_resolved, &query_pool, &cfg, corpus_ids, &query_ids, &simd, &encoder_sha, + write_topk, &mut timing_writer, + ), other => panic!( - "unknown method '{other}'. Supported: flat, hnsw, rq2, rq4, bitmap-rq2, sign-rq2" + "unknown method '{other}'. Supported: flat, hnsw, rq2, rq4, bitmap-rq2, sign-rq2, sign-rq2-threaded" ), } } @@ -935,7 +940,9 @@ fn run_flat( } // --------------------------------------------------------------------------- -// Method: hnsw (pure-Rust HNSW, hnsw_rs; DistDot on unit-norm vectors) +// Method: hnsw (pure-Rust HNSW, hnsw_rs; DistL2 ≡ max-dot on unit-norm vectors). +// Score is `-distance` (nearer = smaller L2 = higher score), so the eval ranks +// nearest-first; for unit vectors this is the identical ordering DistDot gives. // --------------------------------------------------------------------------- #[allow(clippy::too_many_arguments)] @@ -959,12 +966,16 @@ fn run_hnsw( ) { let slug = "hnsw"; eprintln!(" building HNSW M={HNSW_M} ef_c={HNSW_EF_CONSTRUCTION} ({n_docs} docs) ..."); - let hnsw: Hnsw = Hnsw::new( + // DistL2 (not DistDot): embeddings are unit-normalized, so min-L2 ≡ max-dot ≡ + // max-cosine — identical neighbors — but DistL2 avoids anndists' DistDot + // `1-dot` distance assert, which panics on near-duplicate pairs whose float + // dot rounds just past 1.0 (rare at 171K, frequent at ~1M). + let hnsw: Hnsw = Hnsw::new( HNSW_M, n_docs, HNSW_MAX_LAYER, HNSW_EF_CONSTRUCTION, - DistDot {}, + DistL2 {}, ); // Insert (build uses all cores via the global pool). let doc_refs: Vec<(&[f32], usize)> = (0..n_docs) @@ -994,7 +1005,7 @@ fn run_hnsw( .map(|qi| { hnsw.search(query_rows[qi], top_k, HNSW_EF_SEARCH) .into_iter() - .map(|nb| (nb.d_id as i64, 1.0 - nb.distance)) + .map(|nb| (nb.d_id as i64, -nb.distance)) .collect() }) .collect() @@ -1006,7 +1017,7 @@ fn run_hnsw( .into_iter() .map(|nbs| { nbs.into_iter() - .map(|nb| (nb.d_id as i64, 1.0 - nb.distance)) + .map(|nb| (nb.d_id as i64, -nb.distance)) .collect() }) .collect() @@ -1257,3 +1268,289 @@ fn run_two_stage( timing_writer, ); } + +/// Deterministic EXACTLY-`m` selection over a `(count, id)` candidate pool, by +/// `(count desc, id asc)` -- mirrors `SignBitmap`'s `select_nth_unstable_by` +/// exact-`m_eff` tie-break. The `>= tau` threshold set is `>= m` (boundary ties +/// overshoot); this trims it to exactly `m` by keeping the highest-agreement +/// docs, tie-broken on smaller id. Output is sorted ascending by id so the serial +/// and threaded paths return byte-identical candidate sets. +fn select_exact_m(pool: &mut [(u32, u32)], m: usize, out: &mut Vec) { + out.clear(); + let m_eff = m.min(pool.len()); + if m_eff == 0 { + return; + } + let cmp = |a: &(u32, u32), b: &(u32, u32)| b.0.cmp(&a.0).then_with(|| a.1.cmp(&b.1)); + pool.select_nth_unstable_by(m_eff - 1, cmp); + out.extend(pool[..m_eff].iter().map(|&(_, id)| id)); + out.sort_unstable(); +} + +/// Within-query-threaded `sign->rq2` baseline: doc-major sign codes (`wpd` +/// u64/doc), per-doc agreement via hardware VPOPCNTDQ over all `dim` bits, +/// parallelized over doc-stripes (parallel agreement scan -> histogram -> global +/// top-M tau -> EXACTLY-m trim via [`select_exact_m`]) with the SAME fixed-`m` +/// budget + (count desc, id asc) tie-break as the serial `sign-rq2` baseline. It +/// is the SAME SignBitmap candidate set as `sign-rq2`, just computed with +/// within-query threads (the serial baseline scans single-threaded per query). +/// `threads=1` reproduces the serial sign scan. +#[allow(clippy::too_many_arguments)] +fn run_sign_threaded( + m: usize, + corpus: &[f32], + queries: &[f32], + dim: usize, + n_docs: usize, + n_queries: usize, + top_k: usize, + batch: usize, + threads: usize, + pool: &rayon::ThreadPool, + cfg: &Config, + corpus_ids: &[String], + query_ids: &[String], + simd: &[String], + encoder_sha: &str, + write_topk: bool, + timing_writer: &mut dyn Write, +) { + let slug = "ordvec-sign-rq2-threaded"; + let wpd = dim.div_ceil(64); + eprintln!( + " building doc-major SIGN codes + RankQuant b=2 (threaded, m={m}, {n_docs} docs) ..." + ); + let mut rq = RankQuant::new(dim, 2); + let t0 = Instant::now(); + let mut codes = vec![0u64; n_docs * wpd]; + for d in 0..n_docs { + let row = &corpus[d * dim..(d + 1) * dim]; + let base = d * wpd; + for (j, &v) in row.iter().enumerate() { + // `> 0.0` -- same threshold as core SignBitmap (zero/NaN with negatives). + if v > 0.0 { + codes[base + (j >> 6)] |= 1u64 << (j & 63); + } + } + } + rq.add(corpus); + let build_seconds = t0.elapsed().as_secs_f64(); + // doc-major sign code is dim bits/doc (= dim/8 bytes) -- identical substrate + // size to the serial SignBitmap baseline. + let bytes_per_vector = (wpd * 8) + rq.bytes_per_vec(); + let index_total_mib = ((codes.len() * 8) + rq.byte_size()) as f64 / 1024.0 / 1024.0; + + let out_k = top_k.min(m).min(n_docs); + let warmup = 5.min(n_queries); + let mut scratch = SubsetScratch::new(); + let mut out_scores = vec![f32::NEG_INFINITY; batch * out_k]; + let mut out_indices = vec![-1i64; batch * out_k]; + let mut cand: Vec = Vec::new(); + let mut agree = vec![0u16; n_docs]; + + let (samples, preds) = pool.install(|| { + time_and_collect(n_queries, batch, warmup, write_topk, |bs, be| { + let nq_batch = be - bs; + let needed = nq_batch * out_k; + if out_scores.len() != needed { + out_scores.resize(needed, f32::NEG_INFINITY); + out_indices.resize(needed, -1); + } + // Stage 1: per-query within-query-threaded sign scan -> exact-m CSR. + let mut offsets = Vec::with_capacity(nq_batch + 1); + let mut cand_flat: Vec = Vec::new(); + offsets.push(0usize); + for qi in bs..be { + let q = &queries[qi * dim..(qi + 1) * dim]; + let qcode = build_query_sign(q, wpd); + sign_scan_topm_par( + &codes, wpd, n_docs, &qcode, m, threads, &mut agree, &mut cand, + ); + cand_flat.extend_from_slice(&cand); + offsets.push(cand_flat.len()); + } + // Stage 2: pooled subset rerank. + let batch_q = &queries[bs * dim..be * dim]; + rq.search_asymmetric_subset_batched_serial_into( + batch_q, + &offsets, + &cand_flat, + top_k, + &mut scratch, + &mut out_scores, + &mut out_indices, + ); + let mut idx = vec![-1i64; nq_batch * top_k]; + let mut sc = vec![0.0f32; nq_batch * top_k]; + for qi in 0..nq_batch { + let si = &out_indices[qi * out_k..(qi + 1) * out_k]; + let ss = &out_scores[qi * out_k..(qi + 1) * out_k]; + let copy = si.len().min(top_k); + idx[qi * top_k..qi * top_k + copy].copy_from_slice(&si[..copy]); + sc[qi * top_k..qi * top_k + copy].copy_from_slice(&ss[..copy]); + } + (idx, sc) + }) + }); + + finalize( + slug, + &samples, + preds, + dim, + n_docs, + n_queries, + top_k, + threads, + batch, + m, + bytes_per_vector, + index_total_mib, + build_seconds, + &cfg.dataset, + &cfg.split, + query_ids, + corpus_ids, + &cfg.out_dir, + simd, + encoder_sha, + timing_writer, + ); +} + +/// Query sign bits, doc-major layout (bit `j` set iff `q[j] > 0.0` -- the SAME +/// threshold as core SignBitmap's `build_query_bitmap`, so zero/NaN group with +/// the negatives and this threaded baseline is candidate-faithful to `sign-rq2`). +fn build_query_sign(q: &[f32], wpd: usize) -> Vec { + let mut c = vec![0u64; wpd]; + for (j, &v) in q.iter().enumerate() { + if v > 0.0 { + c[j >> 6] |= 1u64 << (j & 63); + } + } + c +} + +/// Single-pass parallel sign-agreement scan + EXACTLY-m top selection. Scans the +/// doc-major sign codes ONCE (hardware VPOPCNTDQ, bandwidth-bound -- the same +/// vectorized popcount the optimized scan uses, so the baseline is not unfairly +/// slow) into a per-doc agreement buffer, histograms to a global top-M `tau`, +/// then trims the `>= tau` set to exactly `m` via [`select_exact_m`]. +/// `agree` is a caller-owned reusable scratch so the per-query alloc stays out of +/// the timing. +#[allow(clippy::too_many_arguments)] +fn sign_scan_topm_par( + codes: &[u64], + wpd: usize, + n: usize, + qcode: &[u64], + m: usize, + threads: usize, + agree: &mut [u16], + out: &mut Vec, +) { + use rayon::prelude::*; + let dim = wpd * 64; + let t = threads.max(1).min(n.max(1)); + let chunk = n.div_ceil(t).max(1); + // Phase A: ONE parallel pass over the codes -> per-doc agreement. + agree[..n] + .par_chunks_mut(chunk) + .enumerate() + .for_each(|(ci, slot)| { + let d0 = ci * chunk; + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("avx512vpopcntdq") { + unsafe { scan_agree_avx512(codes, wpd, d0, qcode, slot) }; + return; + } + } + for (li, a) in slot.iter_mut().enumerate() { + let base = (d0 + li) * wpd; + let mut ham = 0u32; + for w in 0..wpd { + ham += (codes[base + w] ^ qcode[w]).count_ones(); + } + *a = (dim as u32 - ham) as u16; + } + }); + // Parallel per-stripe histogram + merge -> global top-M threshold tau. + let hists: Vec> = agree[..n] + .par_chunks(chunk) + .map(|slot| { + let mut h = vec![0u32; dim + 1]; + for &a in slot { + h[a as usize] += 1; + } + h + }) + .collect(); + let mut hist = vec![0u64; dim + 1]; + for h in &hists { + for (c, &v) in h.iter().enumerate() { + hist[c] += v as u64; + } + } + let mut cum = 0u64; + let mut tau = 0u32; + for c in (0..=dim).rev() { + cum += hist[c]; + if cum >= m as u64 { + tau = c as u32; + break; + } + } + // Phase B: parallel extract (agreement, id) for agreement >= tau, then trim to + // EXACTLY m -- same fixed budget + tie-break as the serial SignBitmap baseline + // as the serial sign baseline, so both rerank identical-size candidate sets. + let parts: Vec> = agree[..n] + .par_chunks(chunk) + .enumerate() + .map(|(ci, slot)| { + let d0 = ci * chunk; + let mut local = Vec::new(); + for (li, &a) in slot.iter().enumerate() { + if a as u32 >= tau { + local.push((a as u32, (d0 + li) as u32)); + } + } + local + }) + .collect(); + let mut poolv: Vec<(u32, u32)> = Vec::new(); + for p in parts { + poolv.extend_from_slice(&p); + } + select_exact_m(&mut poolv, m, out); +} + +/// Hardware VPOPCNTDQ sign-agreement scan for docs `[d0, d0+slot.len())`: the same +/// vectorized popcount the optimized scan uses, so the baseline is not unfairly +/// slow. Fills `slot` with `agreement = dim - hamming`. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512vpopcntdq")] +unsafe fn scan_agree_avx512(codes: &[u64], wpd: usize, d0: usize, qcode: &[u64], slot: &mut [u16]) { + use std::arch::x86_64::*; + let dim = (wpd * 64) as u32; + let cp = codes.as_ptr(); + let qp = qcode.as_ptr(); + for (li, a) in slot.iter_mut().enumerate() { + let base = (d0 + li) * wpd; + let mut acc = _mm512_setzero_si512(); + let mut w = 0usize; + while w + 8 <= wpd { + let c = _mm512_loadu_si512(cp.add(base + w) as *const __m512i); + let q = _mm512_loadu_si512(qp.add(w) as *const __m512i); + let pc = _mm512_popcnt_epi64(_mm512_xor_si512(c, q)); + acc = _mm512_add_epi64(acc, pc); + w += 8; + } + let mut ham = _mm512_reduce_add_epi64(acc) as u32; + while w < wpd { + ham += (*cp.add(base + w) ^ *qp.add(w)).count_ones(); + w += 1; + } + *a = (dim - ham) as u16; + } +} diff --git a/benchmarks/beir/beir_eval.py b/benchmarks/beir/beir_eval.py index 13d0c429..f8ef8b5f 100644 --- a/benchmarks/beir/beir_eval.py +++ b/benchmarks/beir/beir_eval.py @@ -566,6 +566,12 @@ def write_csv( "search_type": "candidate-gen + rerank", "headline_role": "candidate", }, + "ordvec-sign-rq2-threaded": { + "family": "ordvec two-stage", + "implementation": "SignBitmap (within-query threaded, VPOPCNTDQ) → RankQuant b=2", + "search_type": "candidate-gen + rerank", + "headline_role": "candidate (within-query-threaded sign scan)", + }, } From 42ae55de55b6578789d45a02e4ab428a5508205a Mon Sep 17 00:00:00 2001 From: Nelson Spence Date: Tue, 16 Jun 2026 10:47:54 -0500 Subject: [PATCH 2/5] fix(beir-bench): u16 dim guard + parallel sign-code build Address PR #244 bot review (gemini/qodo): - qodo: sign_scan_topm_par stored agreement (= dim - hamming) as u16 and cast silently, which would wrap for dim > 65535 into a wrong tau / candidate set. Add a fail-loud assert on the precondition (BEIR / embedding dims are <= ~4096, far below this) covering both the scalar path and the AVX-512 kernel reached only from this fn. - gemini: parallelize the previously sequential doc-major sign-code build loop with rayon (one wpd-word stripe per doc), matching the adjacent rq.add which is already parallel. Benchmark harness only; core ordvec crate untouched. Signed-off-by: Nelson Spence --- benchmarks/beir-bench/src/main.rs | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/benchmarks/beir-bench/src/main.rs b/benchmarks/beir-bench/src/main.rs index 13eb7520..2ba756f7 100644 --- a/benchmarks/beir-bench/src/main.rs +++ b/benchmarks/beir-bench/src/main.rs @@ -1323,15 +1323,19 @@ fn run_sign_threaded( let mut rq = RankQuant::new(dim, 2); let t0 = Instant::now(); let mut codes = vec![0u64; n_docs * wpd]; - for d in 0..n_docs { - let row = &corpus[d * dim..(d + 1) * dim]; - let base = d * wpd; - for (j, &v) in row.iter().enumerate() { - // `> 0.0` -- same threshold as core SignBitmap (zero/NaN with negatives). - if v > 0.0 { - codes[base + (j >> 6)] |= 1u64 << (j & 63); + // One mutable stripe per doc (`wpd` words) -> parallel sign-code build, matching + // the adjacent `rq.add` which is already parallel over the corpus. + { + use rayon::prelude::*; + codes.par_chunks_mut(wpd).enumerate().for_each(|(d, code)| { + let row = &corpus[d * dim..(d + 1) * dim]; + for (j, &v) in row.iter().enumerate() { + // `> 0.0` -- same threshold as core SignBitmap (zero/NaN with negatives). + if v > 0.0 { + code[j >> 6] |= 1u64 << (j & 63); + } } - } + }); } rq.add(corpus); let build_seconds = t0.elapsed().as_secs_f64(); @@ -1451,6 +1455,15 @@ fn sign_scan_topm_par( ) { use rayon::prelude::*; let dim = wpd * 64; + // Agreement (= dim - hamming) is stored as u16, so dim must fit in u16. BEIR / + // embedding dims are far below this (<= ~4096), but assert the precondition so a + // pathological dim fails loud here instead of silently wrapping the cast into a + // wrong `tau` and wrong candidate set (covers both the scalar path below and the + // AVX-512 kernel, which is only reached from this function). + assert!( + dim <= u16::MAX as usize, + "sign_scan_topm_par: dim {dim} exceeds the u16 agreement range (65535)" + ); let t = threads.max(1).min(n.max(1)); let chunk = n.div_ceil(t).max(1); // Phase A: ONE parallel pass over the codes -> per-doc agreement. From 750e97c20ff77cd25a0dd14b6a4a4ce613709a3f Mon Sep 17 00:00:00 2001 From: Nelson Spence Date: Tue, 16 Jun 2026 12:47:26 -0500 Subject: [PATCH 3/5] fix(beir-bench): u32 agreement, alloc-free timed path, rerank out_k Address PR #244 re-review (gemini HIGH + qodo) on 42ae55d: - qodo (u16 truncation): widen the agreement buffer + AVX-512 kernel from u16 to u32, eliminating the (dim - hamming) cast entirely (the prior assert didn't satisfy the static check). u32 is ~free here -- the agreement buffer is a few % of the codes traffic the scan streams. - gemini (per-query allocs in the timed hot path): make sign_scan_topm_par allocation-free by threading caller-owned reusable scratch -- a per-stripe histogram buffer (hists_buf) and the candidate pool (poolv). Histogram tau now sums per-stripe columns directly (no merge buffer); Phase B stays PARALLEL via par_extend into the reused poolv (rejected gemini's sequential-Phase-B suggestion: it would serialize an O(n) pass and hurt the threaded baseline's latency at 1M+ docs). - qodo (rerank buffer-size mismatch): the rerank was called with top_k but buffers are sized out_k = min(top_k, m/candidates, n_docs); pass out_k instead in BOTH the threaded and the serial sign-rq2 baselines, so the length assert can't panic when the candidate budget is below top_k. Candidate-equivalence preserved (select_exact_m imposes a strict total order, so the parallel extract order is irrelevant). Verified on scifact full corpus: sign-rq2 vs sign-rq2-threaded top-100 overlap = 1.0000 (min 1.0000), 300/300 exact-identical ranked lists. Benchmark harness only; core ordvec crate untouched. Signed-off-by: Nelson Spence --- benchmarks/beir-bench/src/main.rs | 135 +++++++++++++++++------------- 1 file changed, 75 insertions(+), 60 deletions(-) diff --git a/benchmarks/beir-bench/src/main.rs b/benchmarks/beir-bench/src/main.rs index 2ba756f7..b51816aa 100644 --- a/benchmarks/beir-bench/src/main.rs +++ b/benchmarks/beir-bench/src/main.rs @@ -1220,11 +1220,14 @@ fn run_two_stage( }; // Stage 2: pooled subset rerank (allocation-free). + // Rerank for `out_k` (= top_k capped by the candidate budget + corpus), + // matching the `batch * out_k` buffers; passing `top_k` would mis-size the + // buffers and panic the length assert when the budget is below `top_k`. rq.search_asymmetric_subset_batched_serial_into( batch_q, &offsets, &cand_flat, - top_k, + out_k, &mut scratch, &mut out_scores_buf, &mut out_indices_buf, @@ -1350,7 +1353,12 @@ fn run_sign_threaded( let mut out_scores = vec![f32::NEG_INFINITY; batch * out_k]; let mut out_indices = vec![-1i64; batch * out_k]; let mut cand: Vec = Vec::new(); - let mut agree = vec![0u16; n_docs]; + // Per-query scratch, allocated once so the timed path is allocation-free: + // `agree` is the u32 per-doc agreement buffer, `hists_buf` holds one (dim+1)-bin + // histogram per thread-stripe, `poolv` is the reused >= tau candidate pool. + let mut agree = vec![0u32; n_docs]; + let mut hists_buf = vec![0u32; threads.max(1) * (wpd * 64 + 1)]; + let mut poolv: Vec<(u32, u32)> = Vec::with_capacity(m * 2); let (samples, preds) = pool.install(|| { time_and_collect(n_queries, batch, warmup, write_topk, |bs, be| { @@ -1368,18 +1376,29 @@ fn run_sign_threaded( let q = &queries[qi * dim..(qi + 1) * dim]; let qcode = build_query_sign(q, wpd); sign_scan_topm_par( - &codes, wpd, n_docs, &qcode, m, threads, &mut agree, &mut cand, + &codes, + wpd, + n_docs, + &qcode, + m, + threads, + &mut agree, + &mut hists_buf, + &mut poolv, + &mut cand, ); cand_flat.extend_from_slice(&cand); offsets.push(cand_flat.len()); } - // Stage 2: pooled subset rerank. + // Stage 2: pooled subset rerank. Rerank for `out_k` (= top_k capped by `m` + // + corpus) to match the `batch * out_k` buffers; passing `top_k` would + // mis-size them and panic the length assert when `m < top_k`. let batch_q = &queries[bs * dim..be * dim]; rq.search_asymmetric_subset_batched_serial_into( batch_q, &offsets, &cand_flat, - top_k, + out_k, &mut scratch, &mut out_scores, &mut out_indices, @@ -1440,8 +1459,9 @@ fn build_query_sign(q: &[f32], wpd: usize) -> Vec { /// vectorized popcount the optimized scan uses, so the baseline is not unfairly /// slow) into a per-doc agreement buffer, histograms to a global top-M `tau`, /// then trims the `>= tau` set to exactly `m` via [`select_exact_m`]. -/// `agree` is a caller-owned reusable scratch so the per-query alloc stays out of -/// the timing. +/// `agree` (u32 per-doc agreement), `hists_buf` (one `dim+1`-bin histogram per +/// thread-stripe) and `poolv` (the `>= tau` candidate pool) are all caller-owned +/// reusable scratch, so the timed path performs no per-query allocation. #[allow(clippy::too_many_arguments)] fn sign_scan_topm_par( codes: &[u64], @@ -1450,23 +1470,19 @@ fn sign_scan_topm_par( qcode: &[u64], m: usize, threads: usize, - agree: &mut [u16], + agree: &mut [u32], + hists_buf: &mut [u32], + poolv: &mut Vec<(u32, u32)>, out: &mut Vec, ) { use rayon::prelude::*; let dim = wpd * 64; - // Agreement (= dim - hamming) is stored as u16, so dim must fit in u16. BEIR / - // embedding dims are far below this (<= ~4096), but assert the precondition so a - // pathological dim fails loud here instead of silently wrapping the cast into a - // wrong `tau` and wrong candidate set (covers both the scalar path below and the - // AVX-512 kernel, which is only reached from this function). - assert!( - dim <= u16::MAX as usize, - "sign_scan_topm_par: dim {dim} exceeds the u16 agreement range (65535)" - ); + let hlen = dim + 1; let t = threads.max(1).min(n.max(1)); let chunk = n.div_ceil(t).max(1); - // Phase A: ONE parallel pass over the codes -> per-doc agreement. + let stripes = n.div_ceil(chunk); + // Phase A: ONE parallel pass over the codes -> per-doc agreement. Stored as u32 + // so the `dim - hamming` count never truncates regardless of dim. agree[..n] .par_chunks_mut(chunk) .enumerate() @@ -1485,57 +1501,56 @@ fn sign_scan_topm_par( for w in 0..wpd { ham += (codes[base + w] ^ qcode[w]).count_ones(); } - *a = (dim as u32 - ham) as u16; + *a = dim as u32 - ham; } }); - // Parallel per-stripe histogram + merge -> global top-M threshold tau. - let hists: Vec> = agree[..n] - .par_chunks(chunk) - .map(|slot| { - let mut h = vec![0u32; dim + 1]; + // Parallel per-stripe histogram into the reused `hists_buf` (stripe `ci` owns + // `hists_buf[ci*hlen .. (ci+1)*hlen]`); zeroed per query but never reallocated. + let used = stripes * hlen; + hists_buf[..used].fill(0); + hists_buf[..used] + .par_chunks_mut(hlen) + .zip(agree[..n].par_chunks(chunk)) + .for_each(|(h, slot)| { for &a in slot { h[a as usize] += 1; } - h - }) - .collect(); - let mut hist = vec![0u64; dim + 1]; - for h in &hists { - for (c, &v) in h.iter().enumerate() { - hist[c] += v as u64; - } - } + }); + // Global top-M threshold tau: walk agreement high->low, summing the per-stripe + // histogram columns until the cumulative count reaches m (no merge buffer). let mut cum = 0u64; let mut tau = 0u32; - for c in (0..=dim).rev() { - cum += hist[c]; + 'tau: for c in (0..=dim).rev() { + for s in 0..stripes { + cum += hists_buf[s * hlen + c] as u64; + } if cum >= m as u64 { tau = c as u32; - break; + break 'tau; } } - // Phase B: parallel extract (agreement, id) for agreement >= tau, then trim to - // EXACTLY m -- same fixed budget + tie-break as the serial SignBitmap baseline - // as the serial sign baseline, so both rerank identical-size candidate sets. - let parts: Vec> = agree[..n] - .par_chunks(chunk) - .enumerate() - .map(|(ci, slot)| { - let d0 = ci * chunk; - let mut local = Vec::new(); - for (li, &a) in slot.iter().enumerate() { - if a as u32 >= tau { - local.push((a as u32, (d0 + li) as u32)); - } - } - local - }) - .collect(); - let mut poolv: Vec<(u32, u32)> = Vec::new(); - for p in parts { - poolv.extend_from_slice(&p); - } - select_exact_m(&mut poolv, m, out); + // Phase B: parallel extract (agreement, id) for agreement >= tau into the reused + // `poolv` (clear() keeps the capacity; `par_extend` stays parallel), then trim to + // EXACTLY m via `select_exact_m` -- same fixed budget + (count desc, id asc) + // tie-break as the serial sign baseline, so both rerank identical candidate sets. + // Extract order is irrelevant: select_exact_m imposes a strict total order. + poolv.clear(); + poolv.par_extend( + agree[..n] + .par_chunks(chunk) + .enumerate() + .flat_map_iter(|(ci, slot)| { + let d0 = ci * chunk; + slot.iter().enumerate().filter_map(move |(li, &a)| { + if a >= tau { + Some((a, (d0 + li) as u32)) + } else { + None + } + }) + }), + ); + select_exact_m(poolv, m, out); } /// Hardware VPOPCNTDQ sign-agreement scan for docs `[d0, d0+slot.len())`: the same @@ -1543,7 +1558,7 @@ fn sign_scan_topm_par( /// slow. Fills `slot` with `agreement = dim - hamming`. #[cfg(target_arch = "x86_64")] #[target_feature(enable = "avx512f,avx512vpopcntdq")] -unsafe fn scan_agree_avx512(codes: &[u64], wpd: usize, d0: usize, qcode: &[u64], slot: &mut [u16]) { +unsafe fn scan_agree_avx512(codes: &[u64], wpd: usize, d0: usize, qcode: &[u64], slot: &mut [u32]) { use std::arch::x86_64::*; let dim = (wpd * 64) as u32; let cp = codes.as_ptr(); @@ -1564,6 +1579,6 @@ unsafe fn scan_agree_avx512(codes: &[u64], wpd: usize, d0: usize, qcode: &[u64], ham += (*cp.add(base + w) ^ *qp.add(w)).count_ones(); w += 1; } - *a = (dim - ham) as u16; + *a = dim - ham; } } From f95f2944b8e08308f24d3644e35a9757bc851419 Mon Sep 17 00:00:00 2001 From: Nelson Spence Date: Tue, 16 Jun 2026 13:03:17 -0500 Subject: [PATCH 4/5] fix(beir-bench): checked_mul for hists_buf length qodo re-review (750e97c): guard the hists_buf allocation length with checked_mul so a pathological threads/dim can't wrap usize into a too-small buffer in release (unreachable on 64-bit native, but matches the core crate's util::checked_* convention and fails loud). Benchmark harness only; core ordvec crate untouched. Signed-off-by: Nelson Spence --- benchmarks/beir-bench/src/main.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/benchmarks/beir-bench/src/main.rs b/benchmarks/beir-bench/src/main.rs index b51816aa..19bba7c6 100644 --- a/benchmarks/beir-bench/src/main.rs +++ b/benchmarks/beir-bench/src/main.rs @@ -1357,7 +1357,14 @@ fn run_sign_threaded( // `agree` is the u32 per-doc agreement buffer, `hists_buf` holds one (dim+1)-bin // histogram per thread-stripe, `poolv` is the reused >= tau candidate pool. let mut agree = vec![0u32; n_docs]; - let mut hists_buf = vec![0u32; threads.max(1) * (wpd * 64 + 1)]; + // checked_mul so a pathological threads/dim can't wrap usize into a too-small + // buffer in release (matches the core crate's `util::checked_*` convention); + // `sign_scan_topm_par` slices `hists_buf[..stripes * (dim + 1)]`. + let hists_buf_len = threads + .max(1) + .checked_mul(wpd * 64 + 1) + .expect("hists_buf length overflow"); + let mut hists_buf = vec![0u32; hists_buf_len]; let mut poolv: Vec<(u32, u32)> = Vec::with_capacity(m * 2); let (samples, preds) = pool.install(|| { From 4aa646f0ade2dc1fe83aa3f80d3f47d2620abf15 Mon Sep 17 00:00:00 2001 From: Nelson Spence Date: Tue, 16 Jun 2026 13:11:23 -0500 Subject: [PATCH 5/5] fix(beir-bench): guard avx512f in scan_agree_avx512 dispatch Codex stop-gate: the dispatch only checked is_x86_feature_detected!( "avx512vpopcntdq") but scan_agree_avx512 is #[target_feature(enable = "avx512f,avx512vpopcntdq")]. Calling a target_feature fn is only sound when the caller verifies EVERY enabled feature at runtime, so the avx512f check was missing. Guard both, matching the core crate's dispatch convention (lib.rs / multi_bucket.rs). Re-verified candidate-equivalence on scifact (top-100 overlap 1.0000, 300/300 exact). Benchmark harness only; core crate untouched. Signed-off-by: Nelson Spence --- benchmarks/beir-bench/src/main.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/benchmarks/beir-bench/src/main.rs b/benchmarks/beir-bench/src/main.rs index 19bba7c6..2d52b5ea 100644 --- a/benchmarks/beir-bench/src/main.rs +++ b/benchmarks/beir-bench/src/main.rs @@ -1497,7 +1497,13 @@ fn sign_scan_topm_par( let d0 = ci * chunk; #[cfg(target_arch = "x86_64")] { - if std::is_x86_feature_detected!("avx512vpopcntdq") { + // Guard EVERY feature the kernel enables via `#[target_feature]` + // (`avx512f` + `avx512vpopcntdq`) -- detecting only vpopcntdq would + // call into an under-verified target. Mirrors the core crate's + // dispatch (e.g. `lib.rs` / `multi_bucket.rs`). + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512vpopcntdq") + { unsafe { scan_agree_avx512(codes, wpd, d0, qcode, slot) }; return; }