@harperfast/hnsw 0.1.0

package/src/bin/bench.rs

@@ -0,0 +1,246 @@
+//! Standalone cost benchmark: build an N-node graph in the plane file, run queries, report
+//! per-visit cost — the number that decides whether the native plane hits its 0.25–0.4 µs
+//! budget (JS baseline: 4.34 µs/visit at 5M/ef 512).
+//!
+//! Usage: bench [n=100000] [dims=768] [queries=200] [ef=512] [path=/tmp/bench.hnsw] [cap=128] [threads=0]
+//! threads > 0 adds a concurrent-throughput pass: T searcher threads (queries each) + one
+//! background writer inserting throughout, reporting aggregate QPS and per-thread p50/p99.
+
+use hnsw_plane::distance::Query;
+use hnsw_plane::insert::{insert, InsertParams};
+use hnsw_plane::search::{search, SearchScratch};
+use hnsw_plane::{Graph, PlaneFile};
+use std::path::PathBuf;
+use std::time::Instant;
+
+// xorshift for reproducible synthetic vectors without a rand dependency
+struct Rng(u64);
+impl Rng {
+    fn next_unit(&mut self) -> f32 {
+        self.0 ^= self.0 << 13;
+        self.0 ^= self.0 >> 7;
+        self.0 ^= self.0 << 17;
+        (self.0 >> 40) as f32 / (1u64 << 24) as f32
+    }
+    // Box-Muller
+    fn next_gauss(&mut self) -> f32 {
+        let u1 = self.next_unit().max(f32::MIN_POSITIVE);
+        let u2 = self.next_unit();
+        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()
+    }
+}
+
+/// Gaussian-mixture corpus matching benchmarks/hnsw-scale.js: unit centroids, per-dim noise
+/// derived from an intra-cluster cosine target of 0.75 (uniform-random 768-d is a corpus
+/// "no ANN can index" per that benchmark's own calibration notes).
+struct Corpus {
+    centroids: Vec<f32>,
+    n_clusters: usize,
+    dims: usize,
+    noise: f32,
+}
+
+impl Corpus {
+    fn new(n: u64, dims: usize, rng: &mut Rng) -> Self {
+        let intra_cos = 0.75f32;
+        let noise = ((1.0 / (intra_cos * intra_cos) - 1.0) / dims as f32).sqrt();
+        let n_clusters = 8.max((n as f64 / 500.0).round() as usize);
+        let mut centroids = vec![0.0f32; n_clusters * dims];
+        for c in 0..n_clusters {
+            let mut mag = 0.0f32;
+            for d in 0..dims {
+                let x = rng.next_gauss();
+                centroids[c * dims + d] = x;
+                mag += x * x;
+            }
+            let mag = mag.sqrt().max(f32::MIN_POSITIVE);
+            for d in 0..dims {
+                centroids[c * dims + d] /= mag;
+            }
+        }
+        Corpus { centroids, n_clusters, dims, noise }
+    }
+
+    fn row(&self, rng: &mut Rng) -> Vec<f32> {
+        let c = (rng.next_unit() * self.n_clusters as f32) as usize % self.n_clusters;
+        let mut v = vec![0.0f32; self.dims];
+        let mut mag = 0.0f32;
+        for d in 0..self.dims {
+            let x = self.centroids[c * self.dims + d] + rng.next_gauss() * self.noise;
+            v[d] = x;
+            mag += x * x;
+        }
+        let mag = mag.sqrt().max(f32::MIN_POSITIVE);
+        for d in 0..self.dims {
+            v[d] /= mag;
+        }
+        v
+    }
+}
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    let n: u64 = args.get(1).and_then(|a| a.parse().ok()).unwrap_or(100_000);
+    let dims: usize = args.get(2).and_then(|a| a.parse().ok()).unwrap_or(768);
+    let queries: usize = args.get(3).and_then(|a| a.parse().ok()).unwrap_or(200);
+    let ef: usize = args.get(4).and_then(|a| a.parse().ok()).unwrap_or(512);
+    let path: PathBuf = args.get(5).map(Into::into).unwrap_or_else(|| "/tmp/bench.hnsw".into());
+    let layer0_cap: usize = args.get(6).and_then(|a| a.parse().ok()).unwrap_or(128);
+
+    // Reuse an existing plane file when it already holds exactly n nodes at the same cap
+    // (ef sweeps without rebuilding). The corpus RNG below replays identically.
+    let reuse = PlaneFile::open(&path)
+        .ok()
+        .filter(|f| f.id_high_water() == n && f.layer0_cap == layer0_cap)
+        .is_some();
+    let file = if reuse {
+        println!("reusing existing plane at {}", path.display());
+        PlaneFile::open(&path).expect("open")
+    } else {
+        PlaneFile::create(&path, dims, layer0_cap, n + 1024).expect("create")
+    };
+    println!(
+        "plane: {} nodes x {} dims, slot {} B, file {:.1} GB (sparse)",
+        n,
+        dims,
+        file.slot_size,
+        (n * file.slot_size as u64) as f64 / 1e9
+    );
+    let graph = Graph::new(file);
+    let params = InsertParams::default();
+    let mut scratch = SearchScratch::new();
+    let mut rng = Rng(0x1234_5678_9abc_def0);
+    let corpus = Corpus::new(n, dims, &mut rng);
+
+    if reuse {
+        // replay the build's RNG draws so query rows match a fresh run; the upper region
+        // persists inside the plane file
+        for _ in 0..n {
+            let _ = corpus.row(&mut rng);
+        }
+    } else {
+        let build_start = Instant::now();
+        for i in 0..n {
+            let v = corpus.row(&mut rng);
+            insert(&graph, &v, &params, &mut scratch);
+            if (i + 1) % 50_000 == 0 {
+                let rate = (i + 1) as f64 / build_start.elapsed().as_secs_f64();
+                println!("  built {} ({:.0} inserts/s)", i + 1, rate);
+            }
+        }
+        let build = build_start.elapsed();
+        println!("build: {:.1}s ({:.0} inserts/s)", build.as_secs_f64(), n as f64 / build.as_secs_f64());
+        graph.file.msync().expect("msync");
+    }
+
+    // Query with held-out vectors; measure latency and set-recall@10 vs brute-force truth
+    // (same asymmetric metric, so recall isolates graph quality, not quantization).
+    let mut latencies = Vec::with_capacity(queries);
+    let mut total_visits = 0u64;
+    let mut recall_hits = 0usize;
+    let mut recall_total = 0usize;
+    for _ in 0..queries {
+        let q = Query::new(corpus.row(&mut rng));
+        let start = Instant::now();
+        let (results, stats) = search(&graph, &q, 10, ef, &mut scratch);
+        latencies.push(start.elapsed());
+        total_visits += stats.visits;
+        assert!(!results.is_empty());
+
+        let mut truth: Vec<(u32, f32)> = (0..n as u32)
+            .filter_map(|id| graph.distance_to(id, &q).map(|d| (id, d)))
+            .collect();
+        truth.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
+        truth.truncate(10);
+        recall_total += truth.len();
+        recall_hits += truth.iter().filter(|(tid, _)| results.iter().any(|(rid, _)| rid == tid)).count();
+    }
+    latencies.sort();
+    let p50 = latencies[queries / 2];
+    let p95 = latencies[queries * 95 / 100];
+    let p99 = latencies[(queries * 99 / 100).min(queries - 1)];
+    let mean_visits = total_visits as f64 / queries as f64;
+    let us_per_visit = p50.as_micros() as f64 / mean_visits;
+    println!(
+        "search (ef {}): p50 {:.2} ms  p95 {:.2} ms  p99 {:.2} ms  visits/query {:.0}  ->  {:.3} us/visit (JS baseline 4.34)",
+        ef,
+        p50.as_secs_f64() * 1e3,
+        p95.as_secs_f64() * 1e3,
+        p99.as_secs_f64() * 1e3,
+        mean_visits,
+        us_per_visit
+    );
+    println!("recall@10 (set): {:.3}", recall_hits as f64 / recall_total as f64);
+
+    let threads: usize = args.get(7).and_then(|a| a.parse().ok()).unwrap_or(0);
+    if threads > 0 {
+        use std::sync::atomic::{AtomicBool, Ordering};
+        use std::sync::Arc;
+        let graph = Arc::new(graph);
+        let corpus = Arc::new(corpus);
+        let stop = Arc::new(AtomicBool::new(false));
+        let per_thread = queries.max(100);
+        let start = Instant::now();
+        let mut handles = Vec::new();
+        for t in 0..threads {
+            let graph = graph.clone();
+            let corpus = corpus.clone();
+            handles.push(std::thread::spawn(move || {
+                let mut scratch = SearchScratch::new();
+                let mut rng = Rng(0x9e37_79b9 ^ (t as u64 + 1) * 0x1234_5677);
+                let mut lat: Vec<std::time::Duration> = Vec::with_capacity(per_thread);
+                for _ in 0..per_thread {
+                    let q = Query::new(corpus.row(&mut rng));
+                    let s = Instant::now();
+                    let (r, _) = search(&graph, &q, 10, ef, &mut scratch);
+                    lat.push(s.elapsed());
+                    assert!(!r.is_empty());
+                }
+                lat.sort();
+                (lat[per_thread / 2], lat[(per_thread * 99 / 100).min(per_thread - 1)])
+            }));
+        }
+        // background writer: sustained inserts while searchers run
+        let writer = {
+            let graph = graph.clone();
+            let corpus = corpus.clone();
+            let stop = stop.clone();
+            std::thread::spawn(move || {
+                let params = InsertParams::default();
+                let mut scratch = SearchScratch::new();
+                let mut rng = Rng(0xdead_beef_cafe_f00d);
+                let mut count = 0u64;
+                while !stop.load(Ordering::Relaxed) {
+                    let v = corpus.row(&mut rng);
+                    if insert(&graph, &v, &params, &mut scratch).is_err() {
+                        break; // plane full
+                    }
+                    count += 1;
+                }
+                count
+            })
+        };
+        let mut p50s = Vec::new();
+        let mut p99s = Vec::new();
+        for h in handles {
+            let (p50, p99) = h.join().unwrap();
+            p50s.push(p50);
+            p99s.push(p99);
+        }
+        let wall = start.elapsed();
+        stop.store(true, Ordering::Relaxed);
+        let inserted = writer.join().unwrap();
+        let total_q = (threads * per_thread) as f64;
+        p50s.sort();
+        p99s.sort();
+        println!(
+            "concurrent: {} threads x {} queries + writer -> {:.0} QPS aggregate  p50(med) {:.2} ms  p99(worst) {:.2} ms  writer {:.0} inserts/s",
+            threads,
+            per_thread,
+            total_q / wall.as_secs_f64(),
+            p50s[threads / 2].as_secs_f64() * 1e3,
+            p99s[threads - 1].as_secs_f64() * 1e3,
+            inserted as f64 / wall.as_secs_f64()
+        );
+    }
+}

package/src/distance.rs

@@ -0,0 +1,138 @@
+//! Distance kernels. Asymmetric: full-precision f32 query × int8-stored vector (matches the JS
+//! quantizeInt8 scale + cached 1/|v| model). Symmetric int8×int8 for construction-time
+//! neighbor↔neighbor checks (stored per-edge distances were dropped from the format; recompute).
+//! AVX2 with scalar fallback; Linux x86_64 is the performance target, other platforms take the
+//! scalar path (fine for dev).
+
+/// Precomputed query state, built once per search.
+pub struct Query {
+    pub vector: Vec<f32>,
+    pub inv_mag: f32,
+}
+
+impl Query {
+    pub fn new(vector: Vec<f32>) -> Self {
+        let mag_sq: f32 = vector.iter().map(|v| v * v).sum();
+        let inv_mag = 1.0 / mag_sq.sqrt().max(f32::MIN_POSITIVE);
+        Query { vector, inv_mag }
+    }
+}
+
+#[inline]
+fn dot_f32_i8_scalar(q: &[f32], v: *const i8) -> f32 {
+    let mut acc = [0.0f32; 8];
+    let chunks = q.len() / 8;
+    for c in 0..chunks {
+        let base = c * 8;
+        for lane in 0..8 {
+            acc[lane] += q[base + lane] * unsafe { *v.add(base + lane) } as f32;
+        }
+    }
+    let mut dot: f32 = acc.iter().sum();
+    for i in chunks * 8..q.len() {
+        dot += q[i] * unsafe { *v.add(i) } as f32;
+    }
+    dot
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2", enable = "fma")]
+unsafe fn dot_f32_i8_avx2(q: &[f32], v: *const i8) -> f32 {
+    use std::arch::x86_64::*;
+    let mut acc0 = _mm256_setzero_ps();
+    let mut acc1 = _mm256_setzero_ps();
+    let chunks = q.len() / 16;
+    for c in 0..chunks {
+        let base = c * 16;
+        let v16 = _mm_loadu_si128(v.add(base) as *const __m128i);
+        let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(v16));
+        let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128(v16, 8)));
+        acc0 = _mm256_fmadd_ps(_mm256_loadu_ps(q.as_ptr().add(base)), lo, acc0);
+        acc1 = _mm256_fmadd_ps(_mm256_loadu_ps(q.as_ptr().add(base + 8)), hi, acc1);
+    }
+    let acc = _mm256_add_ps(acc0, acc1);
+    let s = _mm_add_ps(_mm256_extractf128_ps(acc, 1), _mm256_castps256_ps128(acc));
+    let s = _mm_hadd_ps(s, s);
+    let s = _mm_hadd_ps(s, s);
+    let mut dot = _mm_cvtss_f32(s);
+    for i in chunks * 16..q.len() {
+        dot += q[i] * *v.add(i) as f32;
+    }
+    dot
+}
+
+#[inline]
+fn dot_f32_i8(q: &[f32], v: *const i8) -> f32 {
+    #[cfg(target_arch = "x86_64")]
+    {
+        if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") {
+            return unsafe { dot_f32_i8_avx2(q, v) };
+        }
+    }
+    dot_f32_i8_scalar(q, v)
+}
+
+/// Cosine distance: f32 query × raw int8 vector at `stored` (dims = query.vector.len()).
+/// Zero-copy: `stored` points into the mmap; the caller's seqlock read discards torn results.
+#[inline]
+pub fn cosine_int8_raw(query: &Query, stored: *const i8, scale: f32, stored_inv_mag: f32) -> f32 {
+    let dot = dot_f32_i8(&query.vector, stored);
+    1.0 - dot * scale * stored_inv_mag * query.inv_mag
+}
+
+#[inline]
+fn dot_i8_i8_scalar(a: *const i8, b: *const i8, len: usize) -> i32 {
+    let mut dot = 0i32;
+    for i in 0..len {
+        dot += unsafe { *a.add(i) as i32 * *b.add(i) as i32 };
+    }
+    dot
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+unsafe fn dot_i8_i8_avx2(a: *const i8, b: *const i8, len: usize) -> i32 {
+    use std::arch::x86_64::*;
+    let mut acc = _mm256_setzero_si256();
+    let chunks = len / 16;
+    for c in 0..chunks {
+        let av = _mm256_cvtepi8_epi16(_mm_loadu_si128(a.add(c * 16) as *const __m128i));
+        let bv = _mm256_cvtepi8_epi16(_mm_loadu_si128(b.add(c * 16) as *const __m128i));
+        acc = _mm256_add_epi32(acc, _mm256_madd_epi16(av, bv));
+    }
+    let lo = _mm256_castsi256_si128(acc);
+    let hi = _mm256_extracti128_si256(acc, 1);
+    let s = _mm_add_epi32(lo, hi);
+    let s = _mm_add_epi32(s, _mm_srli_si128(s, 8));
+    let s = _mm_add_epi32(s, _mm_srli_si128(s, 4));
+    let mut dot = _mm_cvtsi128_si32(s);
+    for i in chunks * 16..len {
+        dot += *a.add(i) as i32 * *b.add(i) as i32;
+    }
+    dot
+}
+
+/// Cosine distance between two int8-stored vectors (construction-time neighbor checks).
+#[inline]
+pub fn cosine_i8_i8_raw(a: *const i8, scale_a: f32, inv_mag_a: f32, b: *const i8, scale_b: f32, inv_mag_b: f32, len: usize) -> f32 {
+    #[cfg(target_arch = "x86_64")]
+    let dot = if std::arch::is_x86_feature_detected!("avx2") {
+        unsafe { dot_i8_i8_avx2(a, b, len) }
+    } else {
+        dot_i8_i8_scalar(a, b, len)
+    };
+    #[cfg(not(target_arch = "x86_64"))]
+    let dot = dot_i8_i8_scalar(a, b, len);
+    1.0 - dot as f32 * scale_a * scale_b * inv_mag_a * inv_mag_b
+}
+
+/// Symmetric int8 quantization matching the JS quantizeInt8: scale maps max |component| to 127.
+pub fn quantize_int8(vector: &[f32]) -> (Vec<i8>, f32, f32) {
+    let max_abs = vector.iter().fold(0.0f32, |m, v| m.max(v.abs()));
+    let scale = if max_abs == 0.0 { 1.0 } else { max_abs / 127.0 };
+    let inv_scale = 1.0 / scale;
+    let bytes: Vec<i8> = vector.iter().map(|v| (v * inv_scale).round().clamp(-127.0, 127.0) as i8).collect();
+    let mag_sq: f32 = vector.iter().map(|v| v * v).sum();
+    let inv_mag = 1.0 / mag_sq.sqrt().max(f32::MIN_POSITIVE);
+    (bytes, scale, inv_mag)
+}