From 3a5b174c75590861e6804e23c3fb068ae9c2e79f Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:56:03 -0400 Subject: [PATCH 1/2] bench(diskann): add distance benchmarks Criterion harness for the four public distance kernels (l2_squared, scalar_l2_squared, inner_product, pq_asymmetric_distance), parameterized over dimensions 128/384/768/1536 with seeded-RNG operands. --- Cargo.lock | 1 + crates/ruvector-diskann/Cargo.toml | 5 ++ crates/ruvector-diskann/benches/distance.rs | 96 +++++++++++++++++++++ 3 files changed, 102 insertions(+) create mode 100644 crates/ruvector-diskann/benches/distance.rs 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..92781a233a --- /dev/null +++ b/crates/ruvector-diskann/benches/distance.rs @@ -0,0 +1,96 @@ +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; + +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 in DIMS.iter() { + let codes: Vec = (0..dim) + .map(|_| rng.gen_range(0..PQ_CODEBOOK_SIZE) as u8) + .collect(); + let table: Vec = (0..dim * PQ_CODEBOOK_SIZE) + .map(|_| rng.gen_range(-1.0f32..1.0)) + .collect(); + + group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |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); From 4ede6fbc4eed5688f95ec5e1798acc389a804635 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:55:26 -0400 Subject: [PATCH 2/2] bench(diskann): size the PQ fixture by subspace count, not embedding dimension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pq_asymmetric_distance fixture built dim codes and dim*256 lookup entries, but the kernel's unit of work is the PQ subspace count M (src/distance.rs:415), not the embedding dimension. ProductQuantizer::encode emits exactly m codes and build_distance_table allocates m*256 entries (src/pq.rs:173,206). Labelling M as dim inflated the loop count and lookup-table working set ~8x versus the realistic dsub=8 configuration matching the crate's exercised shape (D=32, M=4 — src/index.rs:847,852). Parametrize by (D, M) pairs at dsub=8: D=128->M=16, D=384->M=48, D=768->M=96, D=1536->M=192, with case labels d{dim}_m{m} making the measured unit of work explicit in Criterion output. Co-Authored-By: Claude Sonnet 5 --- crates/ruvector-diskann/benches/distance.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/crates/ruvector-diskann/benches/distance.rs b/crates/ruvector-diskann/benches/distance.rs index 92781a233a..915790c275 100644 --- a/crates/ruvector-diskann/benches/distance.rs +++ b/crates/ruvector-diskann/benches/distance.rs @@ -7,6 +7,9 @@ use ruvector_diskann::distance::{ 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() @@ -64,15 +67,17 @@ 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 in DIMS.iter() { - let codes: Vec = (0..dim) + 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..dim * PQ_CODEBOOK_SIZE) + let table: Vec = (0..m * PQ_CODEBOOK_SIZE) .map(|_| rng.gen_range(-1.0f32..1.0)) .collect(); - group.bench_with_input(BenchmarkId::from_parameter(dim), &dim, |bencher, _| { + 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),