Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
351 changes: 334 additions & 17 deletions Cargo.lock

Large diffs are not rendered by default.

4 changes: 2 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,8 @@ ordered-float = "5"

[dev-dependencies]
# Test/bench corpora are seeded with these; not needed by the library.
rand = "0.8"
rand_chacha = "0.3"
rand = "0.10"
rand_chacha = "0.10"

[features]
# SIMD is always compiled, never feature-gated: x86_64 dispatches AVX-512/AVX2 at
Expand Down
10 changes: 5 additions & 5 deletions examples/bench_rank.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@
//! downstream tooling.

use ordvec::search_asymmetric_byte_lut;
use rand::{Rng, SeedableRng};
use rand::{RngExt, SeedableRng};
Comment thread
Fieldnote-Echo marked this conversation as resolved.
Comment thread
Fieldnote-Echo marked this conversation as resolved.
use rand_chacha::ChaCha8Rng;
use std::time::Instant;
// `RankQuantFastscan` is `#[doc(hidden)]` (optional b=2 scan path);
Expand Down Expand Up @@ -232,8 +232,8 @@ fn load_npy_f32(path: &str) -> (Vec<f32>, usize, usize) {

/// Sample a single standard-normal value.
fn gauss(rng: &mut ChaCha8Rng) -> f32 {
let u1: f32 = rng.gen_range(1e-9..1.0);
let u2: f32 = rng.gen_range(0.0..1.0);
let u1: f32 = rng.random_range(1e-9..1.0);
let u2: f32 = rng.random_range(0.0..1.0);
(-2.0 * u1.ln()).sqrt() * (std::f32::consts::TAU * u2).cos()
}

Expand Down Expand Up @@ -293,14 +293,14 @@ fn make_clustered_corpus(cfg: &Config, seed: u64) -> (Vec<f32>, Vec<f32>, Vec<us

let mut corpus = Vec::with_capacity(cfg.n * d);
for _ in 0..cfg.n {
let c = rng.gen_range(0..cfg.n_clusters);
let c = rng.random_range(0..cfg.n_clusters);
let proto = &protos[c * l..(c + 1) * l];
corpus.extend_from_slice(&make_embedding(proto, noise_doc, &mut rng));
}
let mut queries = Vec::with_capacity(cfg.n_queries * d);
let mut q_clusters = Vec::with_capacity(cfg.n_queries);
for _ in 0..cfg.n_queries {
let c = rng.gen_range(0..cfg.n_clusters);
let c = rng.random_range(0..cfg.n_clusters);
q_clusters.push(c);
let proto = &protos[c * l..(c + 1) * l];
queries.extend_from_slice(&make_embedding(proto, noise_q, &mut rng));
Expand Down
24 changes: 16 additions & 8 deletions src/sign_bitmap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -444,14 +444,14 @@ unsafe fn sign_scan_collect_batched_avx512vpop(
#[cfg(test)]
mod tests {
use super::*;
use rand::{Rng, SeedableRng};
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;

const D: usize = 256;

fn make_corpus(seed: u64, n: usize) -> Vec<f32> {
let mut rng = ChaCha8Rng::seed_from_u64(seed);
(0..n * D).map(|_| rng.gen_range(-1.0..1.0)).collect()
(0..n * D).map(|_| rng.random_range(-1.0..1.0)).collect()
}

fn scalar_hamming(q: &[u64], d: &[u64]) -> u32 {
Expand Down Expand Up @@ -498,7 +498,7 @@ mod tests {
let mut idx = SignBitmap::new(D);
idx.add(&corpus);
let mut rng = ChaCha8Rng::seed_from_u64(11);
let query: Vec<f32> = (0..D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let query: Vec<f32> = (0..D).map(|_| rng.random_range(-1.0..1.0)).collect();
let candidates = idx.top_m_candidates(&query, 10);
assert_eq!(candidates.len(), 10);
// Recompute Hamming distance for each returned candidate and
Expand All @@ -525,7 +525,9 @@ mod tests {
idx.add(&corpus);
let mut rng = ChaCha8Rng::seed_from_u64(99);
let batch: usize = 5;
let queries: Vec<f32> = (0..batch * D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..batch * D)
.map(|_| rng.random_range(-1.0..1.0))
.collect();
for m in [10usize, 30, 100] {
let single: Vec<Vec<u32>> = (0..batch)
.map(|bi| idx.top_m_candidates(&queries[bi * D..(bi + 1) * D], m))
Expand All @@ -552,7 +554,9 @@ mod tests {
const BIG_D: usize = 65_536; // u16::MAX + 1 — the smallest dim above the old cap
let n = 4;
let mut rng = ChaCha8Rng::seed_from_u64(41);
let corpus: Vec<f32> = (0..n * BIG_D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let corpus: Vec<f32> = (0..n * BIG_D)
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let mut original = SignBitmap::new(BIG_D);
original.add(&corpus);

Expand Down Expand Up @@ -586,7 +590,7 @@ mod tests {

// Sanity: same query produces same top-M.
let mut rng = ChaCha8Rng::seed_from_u64(23);
let query: Vec<f32> = (0..D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let query: Vec<f32> = (0..D).map(|_| rng.random_range(-1.0..1.0)).collect();
let orig_top = original.top_m_candidates(&query, 10);
let loaded_top = loaded.top_m_candidates(&query, 10);
assert_eq!(orig_top, loaded_top);
Expand Down Expand Up @@ -616,10 +620,14 @@ mod tests {
const PROD_D: usize = 1024;
let n = 256;
let mut rng = ChaCha8Rng::seed_from_u64(31);
let corpus: Vec<f32> = (0..n * PROD_D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let corpus: Vec<f32> = (0..n * PROD_D)
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let mut idx = SignBitmap::new(PROD_D);
idx.add(&corpus);
let queries: Vec<f32> = (0..3 * PROD_D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..3 * PROD_D)
.map(|_| rng.random_range(-1.0..1.0))
.collect();
// Batched (AVX-512 dispatched at qpv=16) must agree with scalar
// reference computed via simple Hamming.
let batched = idx.top_m_candidates_batched(&queries, 32);
Expand Down
6 changes: 3 additions & 3 deletions src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -454,7 +454,7 @@ impl TopK {
#[cfg(test)]
mod tests {
use super::{and_popcount, checked_new_len, xor_popcount, TopK};
use rand::{Rng, SeedableRng};
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;

fn naive_and(d: &[u64], q: &[u64]) -> u32 {
Expand All @@ -476,8 +476,8 @@ mod tests {
let mut rng = ChaCha8Rng::seed_from_u64(0xC0FFEE);
for qpv in [1usize, 2, 3, 4, 7, 8, 15, 16, 17, 31] {
for _ in 0..64 {
let d: Vec<u64> = (0..qpv).map(|_| rng.gen()).collect();
let q: Vec<u64> = (0..qpv).map(|_| rng.gen()).collect();
let d: Vec<u64> = (0..qpv).map(|_| rng.random()).collect();
let q: Vec<u64> = (0..qpv).map(|_| rng.random()).collect();
assert_eq!(and_popcount(&d, &q), naive_and(&d, &q), "AND qpv={qpv}");
assert_eq!(xor_popcount(&d, &q), naive_xor(&d, &q), "XOR qpv={qpv}");
}
Expand Down
36 changes: 20 additions & 16 deletions tests/index/bitmap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

use ordvec::rank::rank_transform;
use ordvec::{Bitmap, RankQuant};
use rand::{Rng, SeedableRng};
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::{make_corpus, D, N};
Expand All @@ -24,7 +24,7 @@ fn rank_io_round_trip_bitmap_index() {
assert_eq!(loaded.n_top(), idx.n_top());

let mut rng = ChaCha8Rng::seed_from_u64(142);
let q: Vec<f32> = (0..D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let q: Vec<f32> = (0..D).map(|_| rng.random_range(-1.0..1.0)).collect();
let r1 = idx.search(&q, 10);
let r2 = loaded.search(&q, 10);
assert_eq!(r1.indices_for_query(0), r2.indices_for_query(0));
Expand Down Expand Up @@ -61,7 +61,7 @@ fn bitmap_then_subset_recovers_exact_when_m_eq_n() {
rq.add(&corpus);

let mut rng = ChaCha8Rng::seed_from_u64(99_999);
let query: Vec<f32> = (0..D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let query: Vec<f32> = (0..D).map(|_| rng.random_range(-1.0..1.0)).collect();

// Stage 1 with M = N: candidate set is every doc.
let cands = bitmap.top_m_candidates(&query, N);
Expand Down Expand Up @@ -130,7 +130,7 @@ fn bitmap_top_m_candidates_uses_no_ground_truth() {
let mut bitmap = Bitmap::new(D, n_top);
bitmap.add(&corpus);
let mut rng = ChaCha8Rng::seed_from_u64(50);
let query: Vec<f32> = (0..D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let query: Vec<f32> = (0..D).map(|_| rng.random_range(-1.0..1.0)).collect();
let a = bitmap.top_m_candidates(&query, 50);
let b = bitmap.top_m_candidates(&query, 50);
assert_eq!(a, b, "candidate selection must be deterministic");
Expand All @@ -150,21 +150,21 @@ fn bitmap_top_m_candidates_deterministic_at_ties() {
// First 150 docs are exact duplicates → all score identically
// against any query. Remaining 50 are random.
let mut rng = ChaCha8Rng::seed_from_u64(404);
let duplicate_vec: Vec<f32> = (0..TIE_D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let duplicate_vec: Vec<f32> = (0..TIE_D).map(|_| rng.random_range(-1.0..1.0)).collect();
let mut corpus: Vec<f32> = Vec::with_capacity(TIE_N * TIE_D);
for _ in 0..150 {
corpus.extend_from_slice(&duplicate_vec);
}
for _ in 0..50 {
for _ in 0..TIE_D {
corpus.push(rng.gen_range(-1.0..1.0));
corpus.push(rng.random_range(-1.0..1.0));
}
}
let _ = rank_transform(&duplicate_vec); // assert symbol is in scope
let n_top = TIE_D / 4;
let mut bitmap = Bitmap::new(TIE_D, n_top);
bitmap.add(&corpus);
let query: Vec<f32> = (0..TIE_D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let query: Vec<f32> = (0..TIE_D).map(|_| rng.random_range(-1.0..1.0)).collect();

// Repeated calls must produce identical candidate sets — the
// composite key forces a unique partition even when 150 docs
Expand All @@ -180,7 +180,9 @@ fn bitmap_top_m_candidates_deterministic_at_ties() {
// Batched path agrees with single-query (the batched-equivalence
// guarantee from `bitmap_batched_matches_single_query` extended
// to the high-tie regime).
let queries: Vec<f32> = (0..3 * TIE_D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..3 * TIE_D)
.map(|_| rng.random_range(-1.0..1.0))
.collect();
for q in [
&queries[..TIE_D],
&queries[TIE_D..2 * TIE_D],
Expand Down Expand Up @@ -211,10 +213,10 @@ fn bitmap_batched_avx512_production_dim() {
const BATCH: usize = 5;
let mut rng = ChaCha8Rng::seed_from_u64(7);
let corpus: Vec<f32> = (0..N_DOCS * PROD_D)
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let queries: Vec<f32> = (0..BATCH * PROD_D)
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let n_top = PROD_D / 4;
let mut bitmap = Bitmap::new(PROD_D, n_top);
Expand Down Expand Up @@ -247,10 +249,10 @@ fn bitmap_batched_hot_plus_tail_split() {
const BATCH: usize = 11;
let mut rng = ChaCha8Rng::seed_from_u64(101);
let corpus: Vec<f32> = (0..N_DOCS * PROD_D)
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let queries: Vec<f32> = (0..BATCH * PROD_D)
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let n_top = PROD_D / 4;
let mut bitmap = Bitmap::new(PROD_D, n_top);
Expand Down Expand Up @@ -286,7 +288,7 @@ fn bitmap_batched_edge_cases() {

// m == 0: each per-query slot is an empty Vec.
let mut rng = ChaCha8Rng::seed_from_u64(202);
let queries: Vec<f32> = (0..3 * D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..3 * D).map(|_| rng.random_range(-1.0..1.0)).collect();
let res = bitmap.top_m_candidates_batched(&queries, 0);
assert_eq!(res.len(), 3);
for c in &res {
Expand Down Expand Up @@ -333,10 +335,10 @@ fn bitmap_batched_avx512_high_qpv_no_panic() {
const BATCH: usize = 3;
let mut rng = ChaCha8Rng::seed_from_u64(123);
let corpus: Vec<f32> = (0..N_DOCS * HIGH_D)
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let queries: Vec<f32> = (0..BATCH * HIGH_D)
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let n_top = HIGH_D / 4;
let mut bitmap = Bitmap::new(HIGH_D, n_top);
Expand Down Expand Up @@ -367,7 +369,9 @@ fn bitmap_batched_matches_single_query() {
bitmap.add(&corpus);
let mut rng = ChaCha8Rng::seed_from_u64(99);
let batch: usize = 7; // intentionally non-power-of-2
let queries: Vec<f32> = (0..batch * D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..batch * D)
.map(|_| rng.random_range(-1.0..1.0))
.collect();
for m in [10usize, 50, 100] {
let single: Vec<Vec<u32>> = (0..batch)
.map(|bi| bitmap.top_m_candidates(&queries[bi * D..(bi + 1) * D], m))
Expand Down
14 changes: 7 additions & 7 deletions tests/index/fastscan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use std::sync::Arc;
use std::thread;

use ordvec::{RankQuant, RankQuantFastscan};
use rand::{Rng, SeedableRng};
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::{make_corpus, D, N};
Expand All @@ -30,8 +30,8 @@ fn fastscan_b2_top10_matches_avx512_kernel() {
const FN: usize = 100;
let mut rng = ChaCha8Rng::seed_from_u64(31337);

let docs: Vec<f32> = (0..FN * FD).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..3 * FD).map(|_| rng.gen_range(-1.0..1.0)).collect();
let docs: Vec<f32> = (0..FN * FD).map(|_| rng.random_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..3 * FD).map(|_| rng.random_range(-1.0..1.0)).collect();

// Reference: the production RankQuant asym kernel.
let mut idx = RankQuant::new(FD, 2);
Expand Down Expand Up @@ -74,7 +74,7 @@ fn fastscan_handles_k_zero() {
// parallel scan (Codex stop-hook regression, source c4fd4d6).
let corpus = make_corpus(250);
let mut rng = ChaCha8Rng::seed_from_u64(251);
let queries: Vec<f32> = (0..(2 * D)).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..(2 * D)).map(|_| rng.random_range(-1.0..1.0)).collect();

let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
Expand Down Expand Up @@ -116,7 +116,7 @@ fn fastscan_handles_k_greater_than_n_vectors() {
const N_SMALL: usize = 5;
let mut rng = ChaCha8Rng::seed_from_u64(261);
let corpus: Vec<f32> = (0..(N_SMALL * D))
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
Expand All @@ -140,7 +140,7 @@ fn fastscan_handles_k_greater_than_n_vectors() {
fn fastscan_search_is_thread_safe() {
let corpus = make_corpus(262);
let mut rng = ChaCha8Rng::seed_from_u64(263);
let queries: Vec<f32> = (0..(4 * D)).map(|_| rng.gen_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..(4 * D)).map(|_| rng.random_range(-1.0..1.0)).collect();

let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
Expand Down Expand Up @@ -174,7 +174,7 @@ fn fastscan_dim_boundary_matrix() {
const N_SMALL: usize = 16;
let mut rng = ChaCha8Rng::seed_from_u64(270 + dim as u64);
let corpus: Vec<f32> = (0..(N_SMALL * dim))
.map(|_| rng.gen_range(-1.0..1.0))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let mut fs = RankQuantFastscan::new(dim);
fs.add(&corpus);
Expand Down
4 changes: 2 additions & 2 deletions tests/index/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use std::io::Write;

use ordvec::rank::{bucket_centre, bucket_ranks, rank_norm, rank_transform, rankquant_norm};
use ordvec::{Bitmap, Rank, RankQuant, SignBitmap};
use rand::{Rng, SeedableRng};
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;

mod bitmap;
Expand All @@ -41,7 +41,7 @@ pub fn make_corpus(seed: u64) -> Vec<f32> {
let mut rng = ChaCha8Rng::seed_from_u64(seed);
let mut v = vec![0.0f32; N * D];
for x in v.iter_mut() {
*x = rng.gen_range(-1.0..1.0);
*x = rng.random_range(-1.0..1.0);
}
v
}
Expand Down
4 changes: 2 additions & 2 deletions tests/index/multi_bucket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

use ordvec::rank::{bucket_centre, bucket_ranks, rank_transform};
use ordvec::MultiBucketBitmap;
use rand::{Rng, SeedableRng};
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::{make_corpus, D, N};
Expand All @@ -21,7 +21,7 @@ fn multi_bucket_bilinear_equals_symmetric_rankquant_inner(bits: u8) {
let w = mb.outer_product_weights();

let mut rng = ChaCha8Rng::seed_from_u64(700 + bits as u64);
let query: Vec<f32> = (0..D).map(|_| rng.gen_range(-1.0..1.0)).collect();
let query: Vec<f32> = (0..D).map(|_| rng.random_range(-1.0..1.0)).collect();
let q_bitmaps = mb.query_bitmaps_from_ranks(&query);

let q_ranks = rank_transform(&query);
Expand Down
Loading
Loading