129 lines
3.6 KiB
Rust
129 lines
3.6 KiB
Rust
//! Criterion benchmarks for ruv-neural-memory.
|
|
//!
|
|
//! Benchmarks the performance-critical vector search operations:
|
|
//! - HNSW insert (building the index)
|
|
//! - HNSW search (approximate nearest neighbor queries)
|
|
//! - Brute-force nearest neighbor (baseline comparison)
|
|
|
|
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
|
|
use rand::Rng;
|
|
|
|
use ruv_neural_memory::HnswIndex;
|
|
|
|
const DIM: usize = 64;
|
|
|
|
/// Generate a set of random embeddings.
|
|
fn generate_embeddings(count: usize, dim: usize) -> Vec<Vec<f64>> {
|
|
let mut rng = rand::thread_rng();
|
|
(0..count)
|
|
.map(|_| (0..dim).map(|_| rng.gen_range(-1.0..1.0)).collect())
|
|
.collect()
|
|
}
|
|
|
|
/// Build an HNSW index from a set of embeddings.
|
|
fn build_hnsw(embeddings: &[Vec<f64>]) -> HnswIndex {
|
|
let mut index = HnswIndex::new(16, 200);
|
|
for emb in embeddings {
|
|
index.insert(emb);
|
|
}
|
|
index
|
|
}
|
|
|
|
/// Euclidean distance between two vectors.
|
|
fn euclidean_distance(a: &[f64], b: &[f64]) -> f64 {
|
|
a.iter()
|
|
.zip(b.iter())
|
|
.map(|(x, y)| (x - y) * (x - y))
|
|
.sum::<f64>()
|
|
.sqrt()
|
|
}
|
|
|
|
/// Brute-force k-nearest-neighbor search.
|
|
fn brute_force_knn(
|
|
embeddings: &[Vec<f64>],
|
|
query: &[f64],
|
|
k: usize,
|
|
) -> Vec<(usize, f64)> {
|
|
let mut distances: Vec<(usize, f64)> = embeddings
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, v)| (i, euclidean_distance(query, v)))
|
|
.collect();
|
|
distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
distances.truncate(k);
|
|
distances
|
|
}
|
|
|
|
fn bench_hnsw_insert(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_insert");
|
|
group.sample_size(10);
|
|
|
|
for &count in &[1_000, 10_000] {
|
|
let embeddings = generate_embeddings(count, DIM);
|
|
group.bench_with_input(
|
|
BenchmarkId::new("embeddings", count),
|
|
&embeddings,
|
|
|b, embeddings| {
|
|
b.iter(|| {
|
|
let mut index = HnswIndex::new(16, 200);
|
|
for emb in embeddings.iter() {
|
|
index.insert(black_box(emb));
|
|
}
|
|
index
|
|
})
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_hnsw_search(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("hnsw_search");
|
|
|
|
for &count in &[1_000, 10_000] {
|
|
let embeddings = generate_embeddings(count, DIM);
|
|
let index = build_hnsw(&embeddings);
|
|
let mut rng = rand::thread_rng();
|
|
let query: Vec<f64> = (0..DIM).map(|_| rng.gen_range(-1.0..1.0)).collect();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("k10_embeddings", count),
|
|
&(index, query),
|
|
|b, (index, query)| {
|
|
b.iter(|| index.search(black_box(query), black_box(10), black_box(50)))
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_brute_force_nn(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("brute_force_nn");
|
|
|
|
for &count in &[1_000, 10_000] {
|
|
let embeddings = generate_embeddings(count, DIM);
|
|
let mut rng = rand::thread_rng();
|
|
let query: Vec<f64> = (0..DIM).map(|_| rng.gen_range(-1.0..1.0)).collect();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("k10_embeddings", count),
|
|
&(embeddings, query),
|
|
|b, (embeddings, query)| {
|
|
b.iter(|| brute_force_knn(black_box(embeddings), black_box(query), black_box(10)))
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_hnsw_insert,
|
|
bench_hnsw_search,
|
|
bench_brute_force_nn,
|
|
);
|
|
criterion_main!(benches);
|