diff --git a/Cargo.lock b/Cargo.lock index 03d9322bc8..982128478d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9406,6 +9406,7 @@ version = "2.3.2" dependencies = [ "bincode 2.0.1", "bytemuck", + "criterion 0.5.1", "getrandom 0.2.17", "memmap2", "parking_lot 0.12.5", diff --git a/crates/ruvector-diskann/Cargo.toml b/crates/ruvector-diskann/Cargo.toml index bf4adcb7f8..d7ebb8794a 100644 --- a/crates/ruvector-diskann/Cargo.toml +++ b/crates/ruvector-diskann/Cargo.toml @@ -36,6 +36,11 @@ simsimd = { workspace = true, optional = true } [dev-dependencies] tempfile = "3.9" +criterion = { workspace = true } + +[[bench]] +name = "distance" +harness = false [target.'cfg(target_arch = "wasm32")'.dev-dependencies] # wasm32 SIMD128 correctness + A/B timing tests (`wasm-pack test --node`). diff --git a/crates/ruvector-diskann/benches/distance.rs b/crates/ruvector-diskann/benches/distance.rs new file mode 100644 index 0000000000..915790c275 --- /dev/null +++ b/crates/ruvector-diskann/benches/distance.rs @@ -0,0 +1,101 @@ +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion}; +use rand::{rngs::StdRng, Rng, SeedableRng}; +use ruvector_diskann::distance::{ + inner_product, l2_squared, pq_asymmetric_distance, scalar_l2_squared, +}; + +const DIMS: [usize; 4] = [128, 384, 768, 1536]; +const SEED: u64 = 0x5EED_1234_ABCD_EF01; +const PQ_CODEBOOK_SIZE: usize = 256; +const PQ_DSUB: usize = 8; +// (dim, m) pairs at dsub = 8, matching the crate's exercised PQ shape (D=32, M=4). +const PQ_DIM_M: [(usize, usize); 4] = [(128, 16), (384, 48), (768, 96), (1536, 192)]; + +fn random_vector(rng: &mut StdRng, dim: usize) -> Vec { + (0..dim).map(|_| rng.gen_range(-1.0f32..1.0)).collect() +} + +fn bench_l2_squared(c: &mut Criterion) { + let mut group = c.benchmark_group("l2_squared"); + let mut rng = StdRng::seed_from_u64(SEED); + + for &dim in DIMS.iter() { + let a = random_vector(&mut rng, dim); + let b = random_vector(&mut rng, dim); + + group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |bencher, _| { + bencher.iter(|| black_box(l2_squared(black_box(&a), black_box(&b)))); + }); + } + + group.finish(); +} + +fn bench_scalar_l2_squared(c: &mut Criterion) { + let mut group = c.benchmark_group("scalar_l2_squared"); + let mut rng = StdRng::seed_from_u64(SEED); + + for &dim in DIMS.iter() { + let a = random_vector(&mut rng, dim); + let b = random_vector(&mut rng, dim); + + group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |bencher, _| { + bencher.iter(|| black_box(scalar_l2_squared(black_box(&a), black_box(&b)))); + }); + } + + group.finish(); +} + +fn bench_inner_product(c: &mut Criterion) { + let mut group = c.benchmark_group("inner_product"); + let mut rng = StdRng::seed_from_u64(SEED); + + for &dim in DIMS.iter() { + let a = random_vector(&mut rng, dim); + let b = random_vector(&mut rng, dim); + + group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |bencher, _| { + bencher.iter(|| black_box(inner_product(black_box(&a), black_box(&b)))); + }); + } + + group.finish(); +} + +fn bench_pq_asymmetric_distance(c: &mut Criterion) { + let mut group = c.benchmark_group("pq_asymmetric_distance"); + let mut rng = StdRng::seed_from_u64(SEED); + + for &(dim, m) in PQ_DIM_M.iter() { + debug_assert_eq!(dim, m * PQ_DSUB); + let codes: Vec = (0..m) + .map(|_| rng.gen_range(0..PQ_CODEBOOK_SIZE) as u8) + .collect(); + let table: Vec = (0..m * PQ_CODEBOOK_SIZE) + .map(|_| rng.gen_range(-1.0f32..1.0)) + .collect(); + + let label = format!("d{dim}_m{m}"); + group.bench_with_input(BenchmarkId::from_parameter(&label), &m, |bencher, _| { + bencher.iter(|| { + black_box(pq_asymmetric_distance( + black_box(&codes), + black_box(&table), + black_box(PQ_CODEBOOK_SIZE), + )) + }); + }); + } + + group.finish(); +} + +criterion_group!( + benches, + bench_l2_squared, + bench_scalar_l2_squared, + bench_inner_product, + bench_pq_asymmetric_distance +); +criterion_main!(benches);