wehub-resource-sync
102 lines
3.2 KiB
Rust
102 lines
3.2 KiB
Rust
//! Benchmark comparing SIMD vs Scalar L2 distance calculations.
|
|
//!
|
|
//! Run with: `cargo run --example simd_benchmark --features simd --release`
|
|
|
|
use std::hint::black_box;
|
|
use std::time::Instant;
|
|
|
|
fn main() {
|
|
let num_vectors = 10_000;
|
|
let dims = 384; // Standard embedding dimension
|
|
let num_iterations = 100;
|
|
|
|
println!("SIMD vs Scalar L2 Distance Benchmark");
|
|
println!("=====================================");
|
|
println!("Vectors: {}", num_vectors);
|
|
println!("Dimensions: {}", dims);
|
|
println!("Iterations: {}", num_iterations);
|
|
println!();
|
|
|
|
// Generate random vectors
|
|
let query: Vec<f32> = (0..dims).map(|i| (i as f32 * 0.001) % 1.0).collect();
|
|
let vectors: Vec<Vec<f32>> = (0..num_vectors)
|
|
.map(|v| (0..dims).map(|i| ((v + i) as f32 * 0.0017) % 1.0).collect())
|
|
.collect();
|
|
|
|
// Benchmark SIMD version
|
|
let simd_start = Instant::now();
|
|
let mut simd_sum = 0.0f32;
|
|
for _ in 0..num_iterations {
|
|
for vec in &vectors {
|
|
simd_sum += black_box(l2_distance_simd(black_box(&query), black_box(vec)));
|
|
}
|
|
}
|
|
let simd_elapsed = simd_start.elapsed();
|
|
black_box(simd_sum);
|
|
|
|
// Benchmark Scalar version
|
|
let scalar_start = Instant::now();
|
|
let mut scalar_sum = 0.0f32;
|
|
for _ in 0..num_iterations {
|
|
for vec in &vectors {
|
|
scalar_sum += black_box(l2_distance_scalar(black_box(&query), black_box(vec)));
|
|
}
|
|
}
|
|
let scalar_elapsed = scalar_start.elapsed();
|
|
black_box(scalar_sum);
|
|
|
|
// Results
|
|
let total_ops = num_vectors * num_iterations;
|
|
let simd_per_op_ns = simd_elapsed.as_nanos() as f64 / total_ops as f64;
|
|
let scalar_per_op_ns = scalar_elapsed.as_nanos() as f64 / total_ops as f64;
|
|
let speedup = scalar_elapsed.as_nanos() as f64 / simd_elapsed.as_nanos() as f64;
|
|
|
|
println!("Results:");
|
|
println!("--------");
|
|
println!(
|
|
"SIMD: {:>8.2}ms total, {:>6.1}ns per distance",
|
|
simd_elapsed.as_secs_f64() * 1000.0,
|
|
simd_per_op_ns
|
|
);
|
|
println!(
|
|
"Scalar: {:>8.2}ms total, {:>6.1}ns per distance",
|
|
scalar_elapsed.as_secs_f64() * 1000.0,
|
|
scalar_per_op_ns
|
|
);
|
|
println!();
|
|
println!("Speedup: {:.2}x", speedup);
|
|
|
|
// Verify correctness
|
|
let simd_result = l2_distance_simd(&query, &vectors[0]);
|
|
let scalar_result = l2_distance_scalar(&query, &vectors[0]);
|
|
let diff = (simd_result - scalar_result).abs();
|
|
println!();
|
|
println!("Correctness check:");
|
|
println!(" SIMD result: {:.8}", simd_result);
|
|
println!(" Scalar result: {:.8}", scalar_result);
|
|
println!(" Difference: {:.2e} (should be < 1e-5)", diff);
|
|
assert!(diff < 1e-4, "Results differ too much!");
|
|
println!(" ✓ Results match!");
|
|
}
|
|
|
|
/// SIMD L2 distance using the wide crate
|
|
#[cfg(feature = "simd")]
|
|
fn l2_distance_simd(a: &[f32], b: &[f32]) -> f32 {
|
|
memvid_core::simd::l2_distance_simd(a, b)
|
|
}
|
|
|
|
#[cfg(not(feature = "simd"))]
|
|
fn l2_distance_simd(a: &[f32], b: &[f32]) -> f32 {
|
|
l2_distance_scalar(a, b)
|
|
}
|
|
|
|
/// Scalar L2 distance (the OLD implementation)
|
|
#[inline(never)]
|
|
fn l2_distance_scalar(a: &[f32], b: &[f32]) -> f32 {
|
|
a.iter()
|
|
.zip(b.iter())
|
|
.map(|(x, y)| (x - y).powi(2))
|
|
.sum::<f32>()
|
|
.sqrt()
|
|
}
|