Generalize SIMD distance implementation to n-length vectors

.github/workflows/rust.yml

@@ -53,7 +53,7 @@ jobs:
       - name: Test Python bindings
         run: |
           sudo apt-get install -y wamerican
-          cp target/release/libinstant_distance.so instant-distance-py/test/instant_distance.so
+          cp target/release/libinstant_distance_py.so instant-distance-py/test/instant_distance.so
           PYTHONPATH=instant-distance-py/test/ python3 -m test
 
   lint:

Makefile

@@ -1,6 +1,9 @@
-test-python:
-	cargo build --release
-	cp target/release/libinstant_distance.dylib instant-distance-py/test/instant_distance.so
+instant-distance-py/test/instant_distance.so: instant-distance-py/src/lib.rs
+	RUSTFLAGS="-C target-cpu=native" cargo build --release
+	([ -f target/release/libinstant_distance_py.dylib ] && cp target/release/libinstant_distance_py.dylib instant-distance-py/test/instant_distance.so) || \
+	([ -f target/release/libinstant_distance_py.so ] && cp target/release/libinstant_distance_py.so instant-distance-py/test/instant_distance.so)
+
+test-python: instant-distance-py/test/instant_distance.so
 	PYTHONPATH=instant-distance-py/test/ python3 -m test
 
 clean:

instant-distance-py/Cargo.toml

@@ -11,12 +11,23 @@ repository = "https://github.com/InstantDomain/instant-distance"
 readme = "../README.md"
 
 [lib]
+name = "instant_distance_py"
+crate-type = ["cdylib", "lib"]
+
+[package.metadata.maturin]
 name = "instant_distance"
-crate-type = ["cdylib"]
 
 [dependencies]
+aligned-vec = { version = "0.5.0", features = ["serde"] }
 bincode = "1.3.1"
 instant-distance = { version = "0.6", path = "../instant-distance", features = ["with-serde"] }
 pyo3 = { version = "0.18.0", features = ["extension-module"] }
 serde = { version = "1", features = ["derive"] }
-serde-big-array = "0.5.0"
+
+[dev-dependencies]
+bencher = "0.1.5"
+rand = { version = "0.8", features = ["small_rng"] }
+
+[[bench]]
+name = "all"
+harness = false

instant-distance-py/benches/all.rs

@@ -0,0 +1,48 @@
+use aligned_vec::avec;
+use bencher::{benchmark_group, benchmark_main, Bencher};
+
+use instant_distance::{Builder, Metric, Search};
+use instant_distance_py::{EuclidMetric, PointStorage};
+use rand::{rngs::StdRng, Rng, SeedableRng};
+
+benchmark_main!(benches);
+benchmark_group!(benches, distance, build, query);
+
+fn distance(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let point_a = avec![rng.gen(); 304];
+    let point_b = avec![rng.gen(); 304];
+
+    bench.iter(|| EuclidMetric::distance(&point_a, &point_b));
+}
+
+fn build(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let points = (0..1024)
+        .map(|_| vec![rng.gen(); 304])
+        .collect::<Vec<_>>();
+
+    bench.iter(|| {
+        Builder::default()
+            .seed(SEED)
+            .build_hnsw::<Vec<f32>, [f32], EuclidMetric, PointStorage>(points.clone())
+    });
+}
+
+fn query(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let points = (0..1024)
+        .map(|_| vec![rng.gen(); 304])
+        .collect::<Vec<_>>();
+    let (hnsw, _) = Builder::default()
+        .seed(SEED)
+        .build_hnsw::<Vec<f32>, [f32], EuclidMetric, PointStorage>(points);
+    let point = avec![rng.gen(); 304];
+
+    bench.iter(|| {
+        let mut search = Search::default();
+        let _ = hnsw.search(&point, &mut search);
+    });
+}
+
+const SEED: u64 = 123456789;

instant-distance-py/pyproject.toml

@@ -4,3 +4,6 @@ name = "instant-distance"
 [build-system]
 requires = ["maturin >= 0.14, < 0.15"]
 build-backend = "maturin"
+
+[tool.maturin]
+python-source = "python"

instant-distance-py/python/instant_distance/__init__.py

@@ -0,0 +1,5 @@
+from .instant_distance import *
+
+__doc__ = instant_distance.__doc__
+if hasattr(instant_distance, "__all__"):
+    __all__ = instant_distance.__all__