d803bfe2b1
git-subtree-dir: vendor/ruvector git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
336 lines
11 KiB
Rust
336 lines
11 KiB
Rust
//! SIMD Performance Benchmarks
|
|
//!
|
|
//! This module benchmarks SIMD-optimized distance calculations
|
|
//! across various vector dimensions and operation types.
|
|
|
|
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
|
|
use ruvector_core::simd_intrinsics::*;
|
|
|
|
// ============================================================================
|
|
// Helper Functions
|
|
// ============================================================================
|
|
|
|
fn generate_vectors(dim: usize) -> (Vec<f32>, Vec<f32>) {
|
|
let a: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
|
|
let b: Vec<f32> = (0..dim).map(|i| ((i + 100) as f32) * 0.01).collect();
|
|
(a, b)
|
|
}
|
|
|
|
fn generate_batch_vectors(dim: usize, count: usize) -> (Vec<f32>, Vec<Vec<f32>>) {
|
|
let query: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
|
|
let vectors: Vec<Vec<f32>> = (0..count)
|
|
.map(|j| (0..dim).map(|i| ((i + j * 10) as f32) * 0.01).collect())
|
|
.collect();
|
|
(query, vectors)
|
|
}
|
|
|
|
// ============================================================================
|
|
// Euclidean Distance Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_euclidean_by_dimension(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("euclidean_by_dimension");
|
|
|
|
for dim in [32, 64, 128, 256, 384, 512, 768, 1024, 1536, 2048] {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.throughput(Throughput::Elements(dim as u64));
|
|
|
|
group.bench_with_input(BenchmarkId::new("simd", dim), &dim, |bench, _| {
|
|
bench.iter(|| euclidean_distance_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_euclidean_small_vectors(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("euclidean_small_vectors");
|
|
|
|
// Test small vector sizes that may not benefit from SIMD
|
|
for dim in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 16] {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.bench_with_input(BenchmarkId::new("simd", dim), &dim, |bench, _| {
|
|
bench.iter(|| euclidean_distance_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_euclidean_non_aligned(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("euclidean_non_aligned");
|
|
|
|
// Test non-SIMD-aligned sizes
|
|
for dim in [31, 33, 63, 65, 127, 129, 255, 257, 383, 385, 511, 513] {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.bench_with_input(BenchmarkId::new("simd", dim), &dim, |bench, _| {
|
|
bench.iter(|| euclidean_distance_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Dot Product Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_dot_product_by_dimension(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("dot_product_by_dimension");
|
|
|
|
for dim in [32, 64, 128, 256, 384, 512, 768, 1024, 1536, 2048] {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.throughput(Throughput::Elements(dim as u64));
|
|
|
|
group.bench_with_input(BenchmarkId::new("simd", dim), &dim, |bench, _| {
|
|
bench.iter(|| dot_product_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_dot_product_common_embeddings(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("dot_product_common_embeddings");
|
|
|
|
// Common embedding model dimensions
|
|
let dims = [
|
|
(128, "small"),
|
|
(384, "all-MiniLM-L6"),
|
|
(512, "e5-small"),
|
|
(768, "all-mpnet-base"),
|
|
(1024, "e5-large"),
|
|
(1536, "text-embedding-ada-002"),
|
|
(2048, "llama-7b-hidden"),
|
|
];
|
|
|
|
for (dim, name) in dims {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.bench_with_input(BenchmarkId::new(name, dim), &dim, |bench, _| {
|
|
bench.iter(|| dot_product_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Cosine Similarity Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_cosine_by_dimension(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("cosine_by_dimension");
|
|
|
|
for dim in [32, 64, 128, 256, 384, 512, 768, 1024, 1536, 2048] {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.throughput(Throughput::Elements(dim as u64));
|
|
|
|
group.bench_with_input(BenchmarkId::new("simd", dim), &dim, |bench, _| {
|
|
bench.iter(|| cosine_similarity_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Manhattan Distance Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_manhattan_by_dimension(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("manhattan_by_dimension");
|
|
|
|
for dim in [32, 64, 128, 256, 384, 512, 768, 1024, 1536, 2048] {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.throughput(Throughput::Elements(dim as u64));
|
|
|
|
group.bench_with_input(BenchmarkId::new("simd", dim), &dim, |bench, _| {
|
|
bench.iter(|| manhattan_distance_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Batch Operations Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_batch_euclidean(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("batch_euclidean");
|
|
|
|
for count in [10, 100, 1000, 10000] {
|
|
let (query, vectors) = generate_batch_vectors(384, count);
|
|
|
|
group.throughput(Throughput::Elements(count as u64));
|
|
|
|
group.bench_with_input(BenchmarkId::new("384d", count), &count, |bench, _| {
|
|
bench.iter(|| {
|
|
for v in &vectors {
|
|
euclidean_distance_simd(black_box(&query), black_box(v));
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_batch_dot_product(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("batch_dot_product");
|
|
|
|
for count in [10, 100, 1000, 10000] {
|
|
let (query, vectors) = generate_batch_vectors(768, count);
|
|
|
|
group.throughput(Throughput::Elements(count as u64));
|
|
|
|
group.bench_with_input(BenchmarkId::new("768d", count), &count, |bench, _| {
|
|
bench.iter(|| {
|
|
for v in &vectors {
|
|
dot_product_simd(black_box(&query), black_box(v));
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Comparison Benchmarks (All Metrics)
|
|
// ============================================================================
|
|
|
|
fn bench_all_metrics_comparison(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("metrics_comparison");
|
|
|
|
let dim = 384; // Common embedding dimension
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
group.bench_function("euclidean", |bench| {
|
|
bench.iter(|| euclidean_distance_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
|
|
group.bench_function("dot_product", |bench| {
|
|
bench.iter(|| dot_product_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
|
|
group.bench_function("cosine", |bench| {
|
|
bench.iter(|| cosine_similarity_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
|
|
group.bench_function("manhattan", |bench| {
|
|
bench.iter(|| manhattan_distance_simd(black_box(&a), black_box(&b)));
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Memory Access Pattern Benchmarks
|
|
// ============================================================================
|
|
|
|
fn bench_sequential_vs_random_access(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("access_patterns");
|
|
|
|
let dim = 512;
|
|
let count = 1000;
|
|
|
|
// Generate vectors
|
|
let vectors: Vec<Vec<f32>> = (0..count)
|
|
.map(|j| (0..dim).map(|i| ((i + j * 10) as f32) * 0.01).collect())
|
|
.collect();
|
|
let query: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
|
|
|
|
// Sequential access indices
|
|
let sequential_indices: Vec<usize> = (0..count).collect();
|
|
|
|
// Random-ish access indices
|
|
let random_indices: Vec<usize> = (0..count)
|
|
.map(|i| (i * 37 + 13) % count) // Pseudo-random
|
|
.collect();
|
|
|
|
group.bench_function("sequential", |bench| {
|
|
bench.iter(|| {
|
|
for &idx in &sequential_indices {
|
|
euclidean_distance_simd(black_box(&query), black_box(&vectors[idx]));
|
|
}
|
|
});
|
|
});
|
|
|
|
group.bench_function("random", |bench| {
|
|
bench.iter(|| {
|
|
for &idx in &random_indices {
|
|
euclidean_distance_simd(black_box(&query), black_box(&vectors[idx]));
|
|
}
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Throughput Measurement
|
|
// ============================================================================
|
|
|
|
fn bench_throughput_ops_per_second(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("throughput");
|
|
group.sample_size(50);
|
|
|
|
for dim in [128, 384, 768, 1536] {
|
|
let (a, b) = generate_vectors(dim);
|
|
|
|
// Report throughput in operations/second
|
|
group.bench_with_input(BenchmarkId::new("euclidean_ops", dim), &dim, |bench, _| {
|
|
bench.iter(|| {
|
|
// Perform 100 operations per iteration
|
|
for _ in 0..100 {
|
|
euclidean_distance_simd(black_box(&a), black_box(&b));
|
|
}
|
|
});
|
|
});
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("dot_product_ops", dim),
|
|
&dim,
|
|
|bench, _| {
|
|
bench.iter(|| {
|
|
for _ in 0..100 {
|
|
dot_product_simd(black_box(&a), black_box(&b));
|
|
}
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// ============================================================================
|
|
// Criterion Groups
|
|
// ============================================================================
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_euclidean_by_dimension,
|
|
bench_euclidean_small_vectors,
|
|
bench_euclidean_non_aligned,
|
|
bench_dot_product_by_dimension,
|
|
bench_dot_product_common_embeddings,
|
|
bench_cosine_by_dimension,
|
|
bench_manhattan_by_dimension,
|
|
bench_batch_euclidean,
|
|
bench_batch_dot_product,
|
|
bench_all_metrics_comparison,
|
|
bench_sequential_vs_random_access,
|
|
bench_throughput_ops_per_second,
|
|
);
|
|
|
|
criterion_main!(benches);
|