Initial commit

packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/distance_bench.rs

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+use rand::{thread_rng, Rng};
+use vector::{FullPrecisionDistance, Metric};
+
+// make sure the vector is 256-bit (32 bytes) aligned required by _mm256_load_ps
+#[repr(C, align(32))]
+struct Vector32ByteAligned {
+    v: [f32; 256],
+}
+
+fn benchmark_l2_distance_float_rust(c: &mut Criterion) {
+    let (a, b) = prepare_random_aligned_vectors();
+    let mut group = c.benchmark_group("avx-computation");
+    group.sample_size(5000);
+
+    group.bench_function("AVX Rust run", |f| {
+        f.iter(|| {
+            black_box(<[f32; 256]>::distance_compare(
+                black_box(&a.v),
+                black_box(&b.v),
+                Metric::L2,
+            ))
+        })
+    });
+}
+
+// make sure the vector is 256-bit (32 bytes) aligned required by _mm256_load_ps
+fn prepare_random_aligned_vectors() -> (Box<Vector32ByteAligned>, Box<Vector32ByteAligned>) {
+    let a = Box::new(Vector32ByteAligned {
+        v: [(); 256].map(|_| thread_rng().gen_range(0.0..100.0)),
+    });
+
+    let b = Box::new(Vector32ByteAligned {
+        v: [(); 256].map(|_| thread_rng().gen_range(0.0..100.0)),
+    });
+
+    (a, b)
+}
+
+criterion_group!(benches, benchmark_l2_distance_float_rust,);
+criterion_main!(benches);
+

packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/kmeans_bench.rs

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use criterion::{criterion_group, criterion_main, Criterion};
+use diskann::utils::k_means_clustering;
+use rand::Rng;
+
+const NUM_POINTS: usize = 10000;
+const DIM: usize = 100;
+const NUM_CENTERS: usize = 256;
+const MAX_KMEANS_REPS: usize = 12;
+
+fn benchmark_kmeans_rust(c: &mut Criterion) {
+    let mut rng = rand::thread_rng();
+    let data: Vec<f32> = (0..NUM_POINTS * DIM)
+        .map(|_| rng.gen_range(-1.0..1.0))
+        .collect();
+    let centers: Vec<f32> = vec![0.0; NUM_CENTERS * DIM];
+
+    let mut group = c.benchmark_group("kmeans-computation");
+    group.sample_size(500);
+
+    group.bench_function("K-Means Rust run", |f| {
+        f.iter(|| {
+            // let mut centers_copy = centers.clone();
+            let data_copy = data.clone();
+            let mut centers_copy = centers.clone();
+            k_means_clustering(
+                &data_copy,
+                NUM_POINTS,
+                DIM,
+                &mut centers_copy,
+                NUM_CENTERS,
+                MAX_KMEANS_REPS,
+            )
+        })
+    });
+}
+
+fn benchmark_kmeans_c(c: &mut Criterion) {
+    let mut rng = rand::thread_rng();
+    let data: Vec<f32> = (0..NUM_POINTS * DIM)
+        .map(|_| rng.gen_range(-1.0..1.0))
+        .collect();
+    let centers: Vec<f32> = vec![0.0; NUM_CENTERS * DIM];
+
+    let mut group = c.benchmark_group("kmeans-computation");
+    group.sample_size(500);
+
+    group.bench_function("K-Means C++ Run", |f| {
+        f.iter(|| {
+            let data_copy = data.clone();
+            let mut centers_copy = centers.clone();
+            let _ = k_means_clustering(
+                data_copy.as_slice(),
+                NUM_POINTS,
+                DIM,
+                centers_copy.as_mut_slice(),
+                NUM_CENTERS,
+                MAX_KMEANS_REPS,
+            );
+        })
+    });
+}
+
+criterion_group!(benches, benchmark_kmeans_rust, benchmark_kmeans_c);
+
+criterion_main!(benches);
+

packages/leann-backend-diskann/third_party/DiskANN/rust/diskann/benches/neighbor_bench.rs

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use std::time::Duration;
+
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+use diskann::model::{Neighbor, NeighborPriorityQueue};
+use rand::distributions::{Distribution, Uniform};
+use rand::rngs::StdRng;
+use rand::SeedableRng;
+
+fn benchmark_priority_queue_insert(c: &mut Criterion) {
+    let vec = generate_random_floats();
+    let mut group = c.benchmark_group("neighborqueue-insert");
+    group.measurement_time(Duration::from_secs(3)).sample_size(500);
+
+    let mut queue = NeighborPriorityQueue::with_capacity(64_usize);
+    group.bench_function("Neighbor Priority Queue Insert", |f| {
+        f.iter(|| {
+            queue.clear();
+            for n in vec.iter() {
+                queue.insert(*n);
+            }
+
+            black_box(&1)
+        });
+    });
+}
+
+fn generate_random_floats() -> Vec<Neighbor> {
+    let seed: [u8; 32] = [73; 32];
+    let mut rng: StdRng = SeedableRng::from_seed(seed);
+    let range = Uniform::new(0.0, 1.0);
+    let mut random_floats = Vec::with_capacity(100);
+
+    for i in 0..100 {
+        let random_float = range.sample(&mut rng) as f32;
+        let n = Neighbor::new(i, random_float);
+        random_floats.push(n);
+    }
+
+    random_floats
+}
+
+criterion_group!(benches, benchmark_priority_queue_insert);
+criterion_main!(benches);
+