diff --git a/crates/pi-natives/src/lib.rs b/crates/pi-natives/src/lib.rs index 0eae423a4..06a0826f4 100644 --- a/crates/pi-natives/src/lib.rs +++ b/crates/pi-natives/src/lib.rs @@ -54,6 +54,7 @@ pub(crate) mod testing; pub mod text; pub mod tokens; pub(crate) mod utils; +pub mod vectors; pub mod workspace; #[cfg(target_os = "windows")] diff --git a/crates/pi-natives/src/vectors.rs b/crates/pi-natives/src/vectors.rs new file mode 100644 index 000000000..5dbbc127c --- /dev/null +++ b/crates/pi-natives/src/vectors.rs @@ -0,0 +1,342 @@ +//! Batch numeric vector kernels for mnemopi recall paths. +//! +//! Every export processes an entire candidate batch per N-API crossing so the +//! crossing cost is amortized over the whole recall operation. Semantics +//! mirror the TypeScript reference implementations in +//! `packages/mnemopi/src/core` exactly — same accumulation order, same +//! non-finite handling, same tie-breaking — so float scores are +//! bit-identical to the TS versions and integer results are exactly equal. + +use napi::{ + Error, Result, Status, + bindgen_prelude::{Float32Array, Float64Array, Uint32Array}, +}; +use napi_derive::napi; + +fn invalid(message: &str) -> Result { + Err(Error::new(Status::InvalidArg, message)) +} + +#[inline] +const fn finite_or_zero(value: f64) -> f64 { + if value.is_finite() { value } else { 0.0 } +} + +/// Cosine similarity with the exact semantics of mnemopi's TS +/// `cosineSimilarity`: iterate `max(len_a, len_b)` elements, treat missing +/// and non-finite entries as `0`, return `0` when either norm is zero. +/// +/// Splitting the shared prefix from the tails preserves bit-exactness: tail +/// terms of the shorter side only ever add `±0.0` to `dot` and `+0.0` to its +/// own norm, in the same index order as the TS loop. +#[inline] +#[allow( + clippy::suboptimal_flops, + reason = "mul_add rounds differently; bit-exact with the TS loops is the contract" +)] +fn cosine_one(a: &[f64], b: &[f64]) -> f64 { + if a.is_empty() && b.is_empty() { + return 0.0; + } + let shared = a.len().min(b.len()); + let mut dot = 0.0f64; + let mut norm_a = 0.0f64; + let mut norm_b = 0.0f64; + for i in 0..shared { + let av = finite_or_zero(a[i]); + let bv = finite_or_zero(b[i]); + dot += av * bv; + norm_a += av * av; + norm_b += bv * bv; + } + for &raw in &a[shared..] { + let av = finite_or_zero(raw); + norm_a += av * av; + } + for &raw in &b[shared..] { + let bv = finite_or_zero(raw); + norm_b += bv * bv; + } + if norm_a == 0.0 || norm_b == 0.0 { + return 0.0; + } + dot / (norm_a.sqrt() * norm_b.sqrt()) +} + +/// All pairs `(i, j)` with `i < j` whose cosine similarity meets `threshold`. +/// +/// `vectors` is `count` vectors flattened row-major at `dim` `f64` elements +/// per row (zero-padded, which matches the TS `?? 0` missing-element +/// semantics), so the similarity is bit-identical to the TS pairwise loop in +/// `clusterBySimilarity`. Returns pairs flattened as `[i0, j0, i1, j1, ...]` +/// in the same `(i, j)` visit order as the TS nested loop. +#[napi] +pub fn cosine_similarity_pairs( + vectors: Float64Array, + count: u32, + dim: u32, + threshold: f64, +) -> Result { + let count = count as usize; + let dim = dim as usize; + let data: &[f64] = &vectors; + if data.len() != count * dim { + return invalid("vectors length must equal count * dim"); + } + let widened: &[f64] = data; + let mut pairs: Vec = Vec::new(); + for i in 0..count { + let left = &widened[i * dim..(i + 1) * dim]; + for j in (i + 1)..count { + let right = &widened[j * dim..(j + 1) * dim]; + if cosine_one(left, right) >= threshold { + pairs.push(i as u32); + pairs.push(j as u32); + } + } + } + Ok(Uint32Array::new(pairs)) +} + +/// Top-k rows of a normalized vector matrix ranked by dot product with a +/// normalized query. +#[napi(object)] +pub struct VectorTopK { + /// Row indices of the selected hits, best score first. + pub indices: Uint32Array, + /// Scores aligned with `indices`. + pub scores: Float64Array, +} + +/// Score every row of a normalized `f32` matrix against `query` and return +/// the top `limit` rows. +/// +/// Mirrors the TS `searchExactVectorIndex` loop bit-exactly: the query is +/// normalized by the L2 norm of its *full* length, each row score sums +/// `matrix[row][col] * (query[col] / norm)` over +/// `min(query.len, dimensions)` columns in column order. Ranking matches the +/// TS stable sort: score descending, lower row index first on exact ties +/// (`-0.0` and `+0.0` compare equal). Callers are expected to enforce the TS +/// guards first (finite query with a positive norm, non-empty matrix). +#[napi] +#[allow( + clippy::suboptimal_flops, + reason = "mul_add rounds differently; bit-exact with the TS loops is the contract" +)] +pub fn vector_index_top_k( + matrix: Float32Array, + dimensions: u32, + query: Float64Array, + limit: u32, +) -> Result { + let dims = dimensions as usize; + let data: &[f32] = &matrix; + if dims == 0 || !data.len().is_multiple_of(dims) { + return invalid("matrix length must be a positive multiple of dimensions"); + } + let count = data.len() / dims; + let q: &[f64] = &query; + let mut norm_sq = 0.0f64; + for &value in q { + norm_sq += value * value; + } + let norm = norm_sq.sqrt(); + // Hoisting the per-column division out of the row loop is bitwise + // identical to the TS per-row `query[col] / queryNorm`. + let query_dims = q.len().min(dims); + let normalized: Vec = q[..query_dims].iter().map(|&v| v / norm).collect(); + + let mut order: Vec<(f64, u32)> = Vec::with_capacity(count); + for row in 0..count { + let base = row * dims; + let mut score = 0.0f64; + for (col, &qv) in normalized.iter().enumerate() { + score += f64::from(data[base + col]) * qv; + } + order.push((score, row as u32)); + } + // JS comparator `(a, b) => b.score - a.score` under a stable sort: strict + // score ordering, otherwise (equal, including ±0.0) original row order. + order.sort_by(|a, b| { + let diff = b.0 - a.0; + if diff > 0.0 { + core::cmp::Ordering::Greater + } else if diff < 0.0 { + core::cmp::Ordering::Less + } else { + a.1.cmp(&b.1) + } + }); + let take = (limit as usize).min(order.len()); + order.truncate(take); + let indices: Vec = order.iter().map(|&(_, row)| row).collect(); + let scores: Vec = order.iter().map(|&(score, _)| score).collect(); + Ok(VectorTopK { indices: Uint32Array::new(indices), scores: Float64Array::new(scores) }) +} + +/// ECMA-262 `\s` (`WhiteSpace` ∪ `LineTerminator`), which differs from Rust's +/// `char::is_whitespace` (JS additionally includes U+FEFF). +#[inline] +const fn is_js_whitespace(c: char) -> bool { + matches!( + c, + '\u{0009}' + | '\u{000a}' + | '\u{000b}' + | '\u{000c}' + | '\u{000d}' + | '\u{0020}' + | '\u{00a0}' + | '\u{1680}' + | '\u{2000}' + ..='\u{200a}' + | '\u{2028}' + | '\u{2029}' + | '\u{202f}' + | '\u{205f}' + | '\u{3000}' + | '\u{feff}' + ) +} + +/// Lowercased word set per the TS `jaccardSimilarity` tokenizer: +/// `text.toLowerCase().split(/\s+/).filter(Boolean)` into a `Set`. +/// Returned sorted and deduplicated for merge-based intersection counting. +fn word_set(text: &str) -> Vec> { + let lower = text.to_lowercase(); + let mut words: Vec> = lower + .split(is_js_whitespace) + .filter(|w| !w.is_empty()) + .map(Box::from) + .collect(); + words.sort_unstable(); + words.dedup(); + words +} + +/// Jaccard similarity of two sorted, deduplicated word sets. Matches the TS +/// `jaccardSimilarity`: `0` when either set is empty, otherwise +/// `|A ∩ B| / (|A| + |B| - |A ∩ B|)` with exact integer counts. +fn jaccard_sorted(a: &[Box], b: &[Box]) -> f64 { + if a.is_empty() || b.is_empty() { + return 0.0; + } + let mut intersection = 0usize; + let (mut i, mut j) = (0usize, 0usize); + while i < a.len() && j < b.len() { + match a[i].cmp(&b[j]) { + core::cmp::Ordering::Less => i += 1, + core::cmp::Ordering::Greater => j += 1, + core::cmp::Ordering::Equal => { + intersection += 1; + i += 1; + j += 1; + }, + } + } + intersection as f64 / (a.len() + b.len() - intersection) as f64 +} + +/// MMR selection over pre-sorted candidates using Jaccard word similarity. +/// +/// `contents[i]` and `scores[i]` describe candidate `i`, already sorted by +/// relevance exactly as the TS `mmrRerank` sorts them (the JS stable sort +/// stays on the TS side so its tie and NaN semantics are preserved). +/// Replicates the TS selection loop exactly: candidate `0` is always taken +/// first; each round picks the remaining candidate maximizing +/// `lambda * score - (1 - lambda) * maxSimilarity(selected)` with strict +/// `>` comparisons, so ties keep the earliest remaining candidate — and a +/// round where every score is `NaN` picks the first remaining candidate, +/// matching the TS `bestIdx = 0` initialisation. Returns the selected +/// indices into the input order. +/// +/// Word tokenization matches `text.toLowerCase().split(/\s+/)` (ECMA `\s`, +/// Unicode default full case conversion). Known divergence: unpaired +/// surrogates arrive here as U+FFFD, while JS keeps the lone surrogate; both +/// tokenize to a single non-whitespace word so Jaccard counts still agree +/// unless a text mixes U+FFFD words with lone-surrogate words. +#[napi] +#[allow( + clippy::suboptimal_flops, + reason = "mul_add rounds differently; bit-exact with the TS loops is the contract" +)] +pub fn mmr_rerank_indices( + contents: Vec, + scores: Float64Array, + lambda_param: f64, + top_k: u32, +) -> Result { + if scores.len() != contents.len() { + return invalid("scores length must equal contents length"); + } + let limit = top_k as usize; + let count = contents.len(); + if limit == 0 || count == 0 { + return Ok(Uint32Array::new(Vec::new())); + } + let sets: Vec>> = contents.iter().map(|text| word_set(text)).collect(); + let mut selected: Vec = Vec::with_capacity(limit.min(count)); + selected.push(0); + let mut remaining: Vec = (1..count as u32).collect(); + + while !remaining.is_empty() && selected.len() < limit { + let mut best_idx = 0usize; + let mut best_score = f64::NEG_INFINITY; + for (idx, &candidate) in remaining.iter().enumerate() { + let mut max_similarity = 0.0f64; + for &picked in &selected { + let similarity = jaccard_sorted(&sets[candidate as usize], &sets[picked as usize]); + if similarity > max_similarity { + max_similarity = similarity; + } + } + let relevance = scores[candidate as usize]; + let mmr_score = lambda_param * relevance - (1.0 - lambda_param) * max_similarity; + if mmr_score > best_score { + best_score = mmr_score; + best_idx = idx; + } + } + selected.push(remaining.remove(best_idx)); + } + if selected.len() < limit { + selected.extend(remaining); + selected.truncate(limit); + } + Ok(Uint32Array::new(selected)) +} + +#[cfg(test)] +mod tests { + use super::{cosine_one, is_js_whitespace, jaccard_sorted, word_set}; + + #[test] + fn cosine_matches_reference_semantics() { + assert_eq!(cosine_one(&[], &[]), 0.0); + assert_eq!(cosine_one(&[1.0, 0.0], &[0.0, 0.0]), 0.0); + let same = cosine_one(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0]); + assert!((same - 1.0).abs() < 1e-12); + // Non-finite entries are zeroed, mismatched lengths pad with zero. + let sim = cosine_one(&[f64::NAN, 1.0], &[0.5, 1.0, 2.0]); + let expect = 1.0 / (1.0f64.sqrt() * (0.25f64 + 1.0 + 4.0).sqrt()); + assert!((sim - expect).abs() < 1e-12); + } + + #[test] + fn word_set_matches_js_tokenizer() { + let set = word_set("Hello\u{00a0}WORLD hello\u{feff}world"); + assert_eq!(set, vec![Box::from("hello"), Box::from("world")]); + assert!(word_set("").is_empty()); + assert!(word_set(" \t\n").is_empty()); + assert!(!is_js_whitespace('\u{200b}')); // ZWSP is not JS \s + } + + #[test] + fn jaccard_matches_reference() { + let a = word_set("the quick brown fox"); + let b = word_set("the lazy brown dog"); + let sim = jaccard_sorted(&a, &b); + assert!((sim - 2.0 / 6.0).abs() < 1e-12); + assert_eq!(jaccard_sorted(&a, &word_set("")), 0.0); + } +} diff --git a/packages/mnemopi/CHANGELOG.md b/packages/mnemopi/CHANGELOG.md index dee554a9e..7963c2055 100644 --- a/packages/mnemopi/CHANGELOG.md +++ b/packages/mnemopi/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- Recall hot loops (exact vector-index search, SHMR similarity clustering, and default-similarity MMR rerank) now run on native batch kernels with one N-API crossing per operation. Wrapper-inclusive speedups at dim=384 on Apple M1 (see `bench/native-vectors.bench.json`): 1.6-1.8x top-K search, 2.4x pairwise clustering scans at the measured 100- and 1000-candidate points (0.86x at the measured 10-candidate point, where the whole scan costs ~20µs either way), and 22-36x MMR rerank. Binary vector search stays on the TypeScript loops (per-call packing measured as a wash), as do custom `similarityFn` MMR reranks and incremental per-row scoring ([#6280](https://github.com/can1357/oh-my-pi/pull/6280) by [@wolfiesch](https://github.com/wolfiesch)). + ## [17.0.4] - 2026-07-18 ### Fixed diff --git a/packages/mnemopi/bench/native-vectors.bench.json b/packages/mnemopi/bench/native-vectors.bench.json new file mode 100644 index 000000000..dc3ab8513 --- /dev/null +++ b/packages/mnemopi/bench/native-vectors.bench.json @@ -0,0 +1,100 @@ +{ + "sha": "8047bedaa46276710c9544d08f12bb13996da858", + "date": "2026-07-22T12:17:28.977Z", + "scenario": "dim=384, stride=48B, warmup=20, iterations=200 (adaptive for O(n²) rows, see per-row fields), crossing-inclusive", + "runtime": "bun 1.3.14", + "host": "Apple M1, darwin-arm64", + "rows": [ + { + "kernel": "searchExactVectorIndex (topK wrapper)", + "count": 10, + "ts_us": 17.34, + "native_us": 9.49, + "speedup": 1.83, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "searchExactVectorIndex (topK wrapper)", + "count": 100, + "ts_us": 56, + "native_us": 34.65, + "speedup": 1.62, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "searchExactVectorIndex (topK wrapper)", + "count": 1000, + "ts_us": 508.91, + "native_us": 300.93, + "speedup": 1.69, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "searchExactVectorIndex (topK wrapper)", + "count": 10000, + "ts_us": 5562.41, + "native_us": 3437.82, + "speedup": 1.62, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "cosineSimilarityPairs (incl. flatten+adjacency)", + "count": 10, + "ts_us": 26.42, + "native_us": 30.67, + "speedup": 0.86, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "cosineSimilarityPairs (incl. flatten+adjacency)", + "count": 100, + "ts_us": 4418.75, + "native_us": 1866.25, + "speedup": 2.37, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "cosineSimilarityPairs (incl. flatten+adjacency)", + "count": 1000, + "ts_us": 445625.44, + "native_us": 179685.28, + "speedup": 2.48, + "ts_iterations": 10, + "native_iterations": 10 + }, + { + "kernel": "mmrRerankIndices (via mmrRerank)", + "count": 10, + "ts_us": 395.5, + "native_us": 17.62, + "speedup": 22.44, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "mmrRerankIndices (via mmrRerank)", + "count": 100, + "ts_us": 10439.19, + "native_us": 287.99, + "speedup": 36.25, + "ts_iterations": 200, + "native_iterations": 200 + }, + { + "kernel": "mmrRerankIndices (via mmrRerank)", + "count": 1000, + "ts_us": 103521.56, + "native_us": 3430.07, + "speedup": 30.18, + "ts_iterations": 200, + "native_iterations": 200 + } + ], + "sink": 15782.734317403312 +} diff --git a/packages/mnemopi/bench/native-vectors.bench.ts b/packages/mnemopi/bench/native-vectors.bench.ts new file mode 100644 index 000000000..b721989be --- /dev/null +++ b/packages/mnemopi/bench/native-vectors.bench.ts @@ -0,0 +1,195 @@ +/** + * Crossing-inclusive benchmark of native batch vector kernels vs the TS + * reference loops, at realistic mnemopi recall shapes (fastembed + * bge-small-en-v1.5, dim=384; binarized stride=48 bytes). + * + * Run from the repo root: `bun packages/mnemopi/bench/native-vectors.bench.ts` + */ +import * as os from "node:os"; +import { + cosineSimilarityPairs, +} from "@oh-my-pi/pi-natives"; +import { jaccardSimilarity, mmrRerank } from "../src/core/mmr"; +import { searchExactVectorIndex } from "../src/core/vector-index"; +import { cosineSimilarity } from "../src/core/vector-math"; + +const DIM = 384; +const STRIDE = DIM / 8; +const COUNTS = [10, 100, 1000, 10000]; +const WARMUP = 20; +const ITERATIONS = 200; + +function makeRng(seed: number): () => number { + let state = seed >>> 0; + return () => { + state = (state * 1664525 + 1013904223) >>> 0; + return state / 4294967296; + }; +} + +let sink = 0; + +function timeNs(fn: () => void, iterations = ITERATIONS, warmup = WARMUP): { ns: number; iterations: number } { + for (let i = 0; i < warmup; i += 1) fn(); + const start = Bun.nanoseconds(); + for (let i = 0; i < iterations; i += 1) fn(); + return { ns: (Bun.nanoseconds() - start) / iterations, iterations }; +} + +interface Row { + kernel: string; + count: number; + tsNs: number; + nativeNs: number; + speedup: number; + tsIterations: number; + nativeIterations: number; +} + +function pushRow(kernel: string, count: number, ts: { ns: number; iterations: number }, native: { ns: number; iterations: number }): void { + rows.push({ + kernel, + count, + tsNs: ts.ns, + nativeNs: native.ns, + speedup: ts.ns / native.ns, + tsIterations: ts.iterations, + nativeIterations: native.iterations, + }); +} + +const rows: Row[] = []; +const rng = makeRng(0xbe4c4); + +// vectorIndexTopK is measured through the public wrapper searchExactVectorIndex, +// so the native side pays the production per-call query conversion, guards, and +// hit-array construction. The TS side replicates the pre-native wrapper body. +for (const count of COUNTS) { + const matrix = new Float32Array(count * DIM); + for (let i = 0; i < matrix.length; i += 1) matrix[i] = rng() * 2 - 1; + const queryArr: number[] = Array.from({ length: DIM }, () => rng() * 2 - 1); + const ids: number[] = Array.from({ length: count }, (_v, i) => i); + const index = { ids, matrix, dimensions: DIM, count }; + const limit = 10; + + const ts = timeNs(() => { + let queryNormSq = 0; + for (const value of queryArr) queryNormSq += value * value; + const queryNorm = Math.sqrt(queryNormSq); + const hits: Array<{ id: number; score: number }> = []; + for (let row = 0; row < count; row += 1) { + const offset = row * DIM; + let score = 0; + for (let col = 0; col < DIM; col += 1) { + score += (matrix[offset + col] ?? 0) * ((queryArr[col] ?? 0) / queryNorm); + } + hits.push({ id: ids[row] ?? 0, score }); + } + hits.sort((a, b) => b.score - a.score); + sink += hits.slice(0, limit)[0]?.score ?? 0; + }); + const native = timeNs(() => { + sink += searchExactVectorIndex(index, queryArr, limit)[0]?.score ?? 0; + }); + pushRow("searchExactVectorIndex (topK wrapper)", count, ts, native); +} + +// cosineSimilarityPairs: O(n²) pair scan. The TS baseline is the pre-native +// clusterBySimilarity loop over per-item vectors; the native side pays the +// production flatten (dim scan + Float64Array fill) before crossing. Capped at +// 1k candidates with adaptive iterations (the TS side at n=1000 is ~500k pair +// cosines × 384 dims per run). +for (const count of COUNTS.filter(n => n <= 1000)) { + const vectors: number[][] = Array.from({ length: count }, () => Array.from({ length: DIM }, () => rng() * 2 - 1)); + const threshold = 0.15; + + const pairIterations = count >= 1000 ? 10 : ITERATIONS; + const pairWarmup = count >= 1000 ? 2 : WARMUP; + const ts = timeNs( + () => { + // Pre-native clusterBySimilarity: adjacency unions built inline + // during the pair scan. + const adjacency: number[][] = Array.from({ length: count }, () => []); + for (let i = 0; i < count; i += 1) { + const a = vectors[i] ?? []; + for (let j = i + 1; j < count; j += 1) { + if (cosineSimilarity(a, vectors[j] ?? []) >= threshold) { + adjacency[i]?.push(j); + adjacency[j]?.push(i); + } + } + } + sink += adjacency[0]?.length ?? 0; + }, + pairIterations, + pairWarmup, + ); + const native = timeNs(() => { + // Current clusterBySimilarity: flatten, one crossing, then adjacency + // built from the materialized pair list. + let dim = 0; + for (const vector of vectors) if (vector.length > dim) dim = vector.length; + const flat = new Float64Array(vectors.length * dim); + for (let i = 0; i < vectors.length; i += 1) { + const vector = vectors[i]; + if (vector === undefined) continue; + for (let col = 0; col < vector.length; col += 1) flat[i * dim + col] = vector[col] ?? 0; + } + const pairs = cosineSimilarityPairs(flat, vectors.length, dim, threshold); + const adjacency: number[][] = Array.from({ length: count }, () => []); + for (let p = 0; p < pairs.length; p += 2) { + adjacency[pairs[p] ?? 0]?.push(pairs[p + 1] ?? 0); + adjacency[pairs[p + 1] ?? 0]?.push(pairs[p] ?? 0); + } + sink += adjacency[0]?.length ?? 0; + }, pairIterations, pairWarmup); + pushRow("cosineSimilarityPairs (incl. flatten+adjacency)", count, ts, native); +} + + +// mmrRerank production paths: the TS side wraps jaccardSimilarity in a lambda, +// defeating the identity check so the exact pre-native selection loop runs; the +// native side calls mmrRerank with the default similarity, exercising the real +// fast path including its sort and wrapper overhead. +const tsJaccard = (a: string, b: string): number => jaccardSimilarity(a, b); +for (const count of COUNTS.filter(n => n <= 1000)) { + const words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta", "iota", "kappa"]; + const results: Array<{ content: string; score: number }> = []; + for (let i = 0; i < count; i += 1) { + const n = 5 + Math.floor(rng() * 20); + results.push({ + content: Array.from({ length: n }, () => words[Math.floor(rng() * words.length)]).join(" "), + score: rng(), + }); + } + const topK = 10; + + const ts = timeNs(() => { + sink += mmrRerank(results, 0.7, topK, tsJaccard).length; + }); + const native = timeNs(() => { + sink += mmrRerank(results, 0.7, topK).length; + }); + pushRow("mmrRerankIndices (via mmrRerank)", count, ts, native); +} + +const sha = Bun.env.BENCH_SHA ?? Bun.spawnSync(["git", "rev-parse", "HEAD"]).stdout.toString().trim(); +const report = { + sha, + date: new Date().toISOString(), + scenario: `dim=${DIM}, stride=${STRIDE}B, warmup=${WARMUP}, iterations=${ITERATIONS} (adaptive for O(n²) rows, see per-row fields), crossing-inclusive`, + runtime: `bun ${Bun.version}`, + host: `${os.cpus()[0]?.model ?? "unknown"}, ${os.platform()}-${os.arch()}`, + rows: rows.map(r => ({ + kernel: r.kernel, + count: r.count, + ts_us: +(r.tsNs / 1000).toFixed(2), + native_us: +(r.nativeNs / 1000).toFixed(2), + speedup: +r.speedup.toFixed(2), + ts_iterations: r.tsIterations, + native_iterations: r.nativeIterations, + })), + sink, +}; + +console.log(JSON.stringify(report, null, 2)); diff --git a/packages/mnemopi/package.json b/packages/mnemopi/package.json index 4d5b505af..c38357420 100644 --- a/packages/mnemopi/package.json +++ b/packages/mnemopi/package.json @@ -41,6 +41,7 @@ "dependencies": { "@oh-my-pi/pi-ai": "catalog:", "@oh-my-pi/pi-catalog": "catalog:", + "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-utils": "catalog:", "lru-cache": "catalog:" }, diff --git a/packages/mnemopi/src/core/mmr.ts b/packages/mnemopi/src/core/mmr.ts index 6cc87322b..aae5662cf 100644 --- a/packages/mnemopi/src/core/mmr.ts +++ b/packages/mnemopi/src/core/mmr.ts @@ -1,3 +1,5 @@ +import { mmrRerankIndices } from "@oh-my-pi/pi-natives"; + export interface MmrResult { readonly content?: string; readonly score?: number; @@ -33,6 +35,34 @@ export function mmrRerank( const first = sortedResults[0]; if (first === undefined) return []; + // Native batch kernel: one N-API crossing selects all indices. Only valid + // for the default Jaccard similarity; custom similarity functions stay in + // TS. Lone-surrogate contents also stay in TS: N-API converts them to + // U+FFFD, which would merge distinct tokens. NaN topK stays in TS so the + // pre-native contract (loop guard is false, first result still returned) + // is preserved. + if ( + similarityFn === jaccardSimilarity && + !Number.isNaN(limit) && + sortedResults.every(result => (result.content ?? "").isWellFormed()) + ) { + // Pre-lowercase with JS semantics so contextual mappings (Final_Sigma: + // "ΟΣ" -> "ος") are applied before the context-insensitive native + // lowercase, which is idempotent on already-lowercased text. + const contents = sortedResults.map(result => (result.content ?? "").toLowerCase()); + const scores = Float64Array.from(sortedResults, result => result.score ?? 0); + // Clamp before the u32 N-API boundary: Infinity or >= 2**32 would + // otherwise wrap (ToUint32) and silently return nothing. + const nativeLimit = Math.min(limit, sortedResults.length); + const picked = mmrRerankIndices(contents, scores, lambdaParam, nativeLimit); + const out: T[] = []; + for (const index of picked) { + const item = sortedResults[index]; + if (item !== undefined) out.push(item); + } + return out; + } + const selected: T[] = [first]; const remaining = sortedResults.slice(1); diff --git a/packages/mnemopi/src/core/shmr.ts b/packages/mnemopi/src/core/shmr.ts index dca63ba24..e6497ed9b 100644 --- a/packages/mnemopi/src/core/shmr.ts +++ b/packages/mnemopi/src/core/shmr.ts @@ -1,5 +1,6 @@ import type { Database } from "bun:sqlite"; import { createHash } from "node:crypto"; +import { cosineSimilarityPairs } from "@oh-my-pi/pi-natives"; import { logger } from "@oh-my-pi/pi-utils"; import * as embeddings from "./embeddings"; import { cosineSimilarity } from "./vector-math"; @@ -166,17 +167,23 @@ export async function clusterBySimilarity(items: readonly ShmrItem[], threshold: if (items.length === 0) return []; const vectors = await resolveItemVectors(items); const adjacency: number[][] = Array.from({ length: items.length }, () => []); - for (let i = 0; i < items.length; i++) { - const leftEmbedding = vectors[i]; - if (leftEmbedding === undefined) continue; - for (let j = i + 1; j < items.length; j++) { - const rightEmbedding = vectors[j]; - if (rightEmbedding === undefined) continue; - if (cosineSimilarity(leftEmbedding, rightEmbedding) >= threshold) { - adjacency[i]?.push(j); - adjacency[j]?.push(i); - } - } + // Native batch kernel: one N-API crossing evaluates all O(n^2) pairs. + // Vectors are zero-padded to a shared dim, matching the TS `?? 0` + // missing-element semantics of cosineSimilarity. + let dim = 0; + for (const vector of vectors) if (vector.length > dim) dim = vector.length; + const flat = new Float64Array(vectors.length * dim); + for (let i = 0; i < vectors.length; i++) { + const vector = vectors[i]; + if (vector === undefined) continue; + for (let col = 0; col < vector.length; col++) flat[i * dim + col] = vector[col] ?? 0; + } + const pairs = cosineSimilarityPairs(flat, vectors.length, dim, threshold); + for (let p = 0; p < pairs.length; p += 2) { + const i = pairs[p] ?? 0; + const j = pairs[p + 1] ?? 0; + adjacency[i]?.push(j); + adjacency[j]?.push(i); } const visited = new Set(); const clusters: ShmrItem[][] = []; diff --git a/packages/mnemopi/src/core/vector-index.ts b/packages/mnemopi/src/core/vector-index.ts index 3092cba52..3c43607bd 100644 --- a/packages/mnemopi/src/core/vector-index.ts +++ b/packages/mnemopi/src/core/vector-index.ts @@ -1,3 +1,5 @@ +import { vectorIndexTopK } from "@oh-my-pi/pi-natives"; + export interface ExactVectorSearchHit { id: TId; score: number; @@ -66,19 +68,16 @@ export function searchExactVectorIndex( queryNormSq += value * value; } if (queryNormSq <= 0) return []; - const queryNorm = Math.sqrt(queryNormSq); - const queryDimensions = Math.min(query.length, index.dimensions); + // Native batch kernel: one N-API crossing scores every row and ranks the + // top k with the same stable ordering as the TS sort. TS guards above + // (finite query, positive norm, non-empty index) are preserved. Clamp k to + // the row count before the u32 boundary: Infinity or >= 2**32 would + // otherwise wrap (ToUint32) and return no hits. + const topK = vectorIndexTopK(index.matrix, index.dimensions, Float64Array.from(query), Math.min(k, index.count)); const hits: ExactVectorSearchHit[] = []; - - for (let row = 0; row < index.count; row += 1) { - const offset = row * index.dimensions; - let score = 0; - for (let col = 0; col < queryDimensions; col += 1) { - score += index.matrix[offset + col] * ((query[col] ?? 0) / queryNorm); - } - hits.push({ id: index.ids[row] as TId, score }); + for (let i = 0; i < topK.indices.length; i += 1) { + const row = topK.indices[i] ?? 0; + hits.push({ id: index.ids[row] as TId, score: topK.scores[i] ?? 0 }); } - - hits.sort((a, b) => b.score - a.score); - return hits.slice(0, Math.min(k, hits.length)); + return hits; } diff --git a/packages/mnemopi/test/native-vector-parity.test.ts b/packages/mnemopi/test/native-vector-parity.test.ts new file mode 100644 index 000000000..2a7340675 --- /dev/null +++ b/packages/mnemopi/test/native-vector-parity.test.ts @@ -0,0 +1,175 @@ +import { describe, expect, test } from "bun:test"; +import { cosineSimilarityPairs, mmrRerankIndices, vectorIndexTopK } from "@oh-my-pi/pi-natives"; +import { jaccardSimilarity, mmrRerank } from "../src/core/mmr"; +import { buildExactVectorIndex, searchExactVectorIndex } from "../src/core/vector-index"; +import { cosineSimilarity } from "../src/core/vector-math"; + +/** Deterministic LCG so parity failures reproduce exactly. */ +function makeRng(seed: number): () => number { + let state = seed >>> 0; + return () => { + state = (state * 1664525 + 1013904223) >>> 0; + return state / 4294967296; + }; +} + +const REL_TOL = 1e-9; + +function expectClose(actual: number, expected: number): void { + if (expected === 0) { + expect(actual).toBe(0); + return; + } + expect(Math.abs(actual - expected) / Math.abs(expected)).toBeLessThanOrEqual(REL_TOL); +} + +describe("native vector kernel parity", () => { + test("cosineSimilarityPairs matches the TS pairwise threshold loop", () => { + const rng = makeRng(0x9a175); + const dim = 64; + const count = 40; + const flat = new Float64Array(count * dim); + for (let i = 0; i < flat.length; i += 1) flat[i] = rng() * 2 - 1; + const threshold = 0.02; + const expected: number[] = []; + for (let i = 0; i < count; i += 1) { + for (let j = i + 1; j < count; j += 1) { + if ( + cosineSimilarity(flat.subarray(i * dim, (i + 1) * dim), flat.subarray(j * dim, (j + 1) * dim)) >= + threshold + ) { + expected.push(i, j); + } + } + } + expect(Array.from(cosineSimilarityPairs(flat, count, dim, threshold))).toEqual(expected); + }); + + test("vectorIndexTopK matches the TS scoring loop and stable sort", () => { + const rng = makeRng(0x70b1); + const dims = 384; + const count = 300; + const matrix = new Float32Array(count * dims); + for (let i = 0; i < matrix.length; i += 1) matrix[i] = rng() * 2 - 1; + const query = Float64Array.from({ length: dims }, () => rng() * 2 - 1); + let normSq = 0; + for (const v of query) normSq += v * v; + const norm = Math.sqrt(normSq); + const hits: Array<{ row: number; score: number }> = []; + for (let row = 0; row < count; row += 1) { + let score = 0; + for (let col = 0; col < dims; col += 1) { + score += (matrix[row * dims + col] ?? 0) * ((query[col] ?? 0) / norm); + } + hits.push({ row, score }); + } + hits.sort((a, b) => b.score - a.score); + const limit = 25; + const result = vectorIndexTopK(matrix, dims, query, limit); + expect(Array.from(result.indices)).toEqual(hits.slice(0, limit).map(h => h.row)); + for (let i = 0; i < limit; i += 1) { + expectClose(result.scores[i] ?? Number.NaN, hits[i]?.score ?? Number.NaN); + } + }); + + test("mmrRerankIndices selects identical index sequences to the TS loop", () => { + const rng = makeRng(0x33a11); + const words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta", "iota", "kappa"]; + const count = 60; + const contents: string[] = []; + const scores = new Float64Array(count); + for (let i = 0; i < count; i += 1) { + const n = 1 + Math.floor(rng() * 8); + contents.push(Array.from({ length: n }, () => words[Math.floor(rng() * words.length)]).join(" ")); + scores[i] = rng(); + } + scores[5] = scores[6] = 0.5; // exercise strict-> tie keeping the earlier candidate + for (const lambda of [0.0, 0.3, 0.7, 1.0]) { + for (const topK of [1, 10, count, count + 5]) { + // TS reference selection over pre-sorted candidates (mirrors mmrRerank + // after its sort step). + const order = contents.map((_, i) => i).sort((a, b) => (scores[b] ?? 0) - (scores[a] ?? 0)); + const sortedContents = order.map(i => contents[i] ?? ""); + const sortedScores = order.map(i => scores[i] ?? 0); + const selected: number[] = [0]; + const remaining = sortedContents.map((_, i) => i).slice(1); + while (remaining.length > 0 && selected.length < topK) { + let bestIdx = 0; + let bestScore = Number.NEGATIVE_INFINITY; + for (let idx = 0; idx < remaining.length; idx += 1) { + const candidate = remaining[idx] ?? 0; + let maxSimilarity = 0; + for (const picked of selected) { + const sim = jaccardSimilarity(sortedContents[candidate] ?? "", sortedContents[picked] ?? ""); + if (sim > maxSimilarity) maxSimilarity = sim; + } + const mmrScore = lambda * (sortedScores[candidate] ?? 0) - (1 - lambda) * maxSimilarity; + if (mmrScore > bestScore) { + bestScore = mmrScore; + bestIdx = idx; + } + } + selected.push(remaining.splice(bestIdx, 1)[0] ?? 0); + } + if (selected.length < topK) selected.push(...remaining.slice(0, topK - selected.length)); + const native = mmrRerankIndices(sortedContents, Float64Array.from(sortedScores), lambda, topK); + expect(Array.from(native)).toEqual(selected.slice(0, topK)); + } + } + }); + + test("mmrRerank wrapper preserves the pre-native limit contract at u32 boundaries", () => { + const results = Array.from({ length: 8 }, (_v, i) => ({ + content: `item ${i} alpha beta`, + score: (8 - i) / 10, + })); + const all = mmrRerank(results, 0.7, results.length); + // Infinity and >= 2**32 previously walked every candidate; ToUint32 + // would have collapsed them to 0/1 without the TS-side clamp. + expect(mmrRerank(results, 0.7, Number.POSITIVE_INFINITY)).toEqual(all); + expect(mmrRerank(results, 0.7, 2 ** 32)).toEqual(all); + expect(mmrRerank(results, 0.7, 2 ** 32 + 1)).toEqual(all); + // NaN: loop guard is false but the first result is already selected. + expect(mmrRerank(results, 0.7, Number.NaN)).toEqual([results[0] as (typeof results)[number]]); + expect(mmrRerank(results, 0.7, 0)).toEqual([]); + expect(mmrRerank(results, 0.7, -3)).toEqual([]); + }); + + test("mmrRerank matches the TS path on contextual-lowercase and lone-surrogate content", () => { + // Force the TS selection loop by defeating the identity check. + const tsJaccard = (a: string, b: string): number => jaccardSimilarity(a, b); + // Final_Sigma: JS lowercases "ΟΣ" to "ος"; a context-insensitive + // lowercase would produce "οσ" and score these words as distinct. + const sigma = [ + { content: "ΟΣ", score: 0.9 }, + { content: "ος", score: 0.8 }, + { content: "other words entirely", score: 0.7 }, + ]; + for (const lambda of [0, 0.3, 0.7]) { + expect(mmrRerank(sigma, lambda, 2)).toEqual(mmrRerank(sigma, lambda, 2, tsJaccard)); + } + // Lone surrogates route to the TS path: N-API would convert them to + // U+FFFD and merge the first two tokens. + const surrogate = [ + { content: "\ud800", score: 0.9 }, + { content: "\ufffd", score: 0.8 }, + { content: "other", score: 0.7 }, + ]; + expect(mmrRerank(surrogate, 0, 2)).toEqual(mmrRerank(surrogate, 0, 2, tsJaccard)); + }); + + test("searchExactVectorIndex preserves the pre-native limit contract at u32 boundaries", () => { + const rows = Array.from({ length: 6 }, (_v, i) => ({ + id: i, + vector: Array.from({ length: 8 }, (_x, j) => Math.sin(i * 8 + j)), + })); + const index = buildExactVectorIndex(rows); + const query = Array.from({ length: 8 }, (_x, j) => Math.cos(j)); + const all = searchExactVectorIndex(index, query, index.count); + expect(all.length).toBe(index.count); + expect(searchExactVectorIndex(index, query, Number.POSITIVE_INFINITY)).toEqual(all); + expect(searchExactVectorIndex(index, query, 2 ** 32)).toEqual(all); + expect(searchExactVectorIndex(index, query, Number.NaN)).toEqual([]); + expect(searchExactVectorIndex(index, query, 0)).toEqual([]); + }); +}); diff --git a/packages/natives/CHANGELOG.md b/packages/natives/CHANGELOG.md index 7eee7f035..370e04e12 100644 --- a/packages/natives/CHANGELOG.md +++ b/packages/natives/CHANGELOG.md @@ -8,6 +8,9 @@ ### Added - Added jsdiff-compatible native diff exports: `diffLines`, `diffWords`, `diffLineRuns`, and `structuredPatchHunks` ([#6279](https://github.com/can1357/oh-my-pi/pull/6279) by [@wolfiesch](https://github.com/wolfiesch)). +### Added + +- Added batch vector kernels for mnemopi recall paths: `cosineSimilarityPairs`, `vectorIndexTopK`, and `mmrRerankIndices` ([#6280](https://github.com/can1357/oh-my-pi/pull/6280) by [@wolfiesch](https://github.com/wolfiesch)). ## [17.0.5] - 2026-07-18 diff --git a/packages/natives/native/index.d.ts b/packages/natives/native/index.d.ts index 950901af7..b22104082 100644 --- a/packages/natives/native/index.d.ts +++ b/packages/natives/native/index.d.ts @@ -470,6 +470,17 @@ export interface ContextLine { */ export declare function copyToClipboard(text: string): void +/** + * All pairs `(i, j)` with `i < j` whose cosine similarity meets `threshold`. + * + * `vectors` is `count` vectors flattened row-major at `dim` `f64` elements + * per row (zero-padded, which matches the TS `?? 0` missing-element + * semantics), so the similarity is bit-identical to the TS pairwise loop in + * `clusterBySimilarity`. Returns pairs flattened as `[i0, j0, i1, j1, ...]` + * in the same `(i, j)` visit order as the TS nested loop. + */ +export declare function cosineSimilarityPairs(vectors: Float64Array, count: number, dim: number, threshold: number): Uint32Array + /** * Count tokens in `input`. * @@ -1250,6 +1261,28 @@ export interface MinimizerResult { outputBytes: number } +/** + * MMR selection over pre-sorted candidates using Jaccard word similarity. + * + * `contents[i]` and `scores[i]` describe candidate `i`, already sorted by + * relevance exactly as the TS `mmrRerank` sorts them (the JS stable sort + * stays on the TS side so its tie and NaN semantics are preserved). + * Replicates the TS selection loop exactly: candidate `0` is always taken + * first; each round picks the remaining candidate maximizing + * `lambda * score - (1 - lambda) * maxSimilarity(selected)` with strict + * `>` comparisons, so ties keep the earliest remaining candidate — and a + * round where every score is `NaN` picks the first remaining candidate, + * matching the TS `bestIdx = 0` initialisation. Returns the selected + * indices into the input order. + * + * Word tokenization matches `text.toLowerCase().split(/\s+/)` (ECMA `\s`, + * Unicode default full case conversion). Known divergence: unpaired + * surrogates arrive here as U+FFFD, while JS keeps the lone surrogate; both + * tokenize to a single non-whitespace word so Jaccard counts still agree + * unless a text mixes U+FFFD words with lone-surrogate words. + */ +export declare function mmrRerankIndices(contents: Array, scores: Float64Array, lambdaParam: number, topK: number): Uint32Array + /** Parsed Kitty keyboard protocol sequence result for a Kitty input sequence. */ export interface ParsedKittyResult { /** Primary codepoint associated with the key. */ @@ -1668,6 +1701,31 @@ export declare function supportsLanguage(lang: string): boolean */ export declare function truncateToWidth(text: string, maxWidth: number, ellipsisKind: Ellipsis | undefined | null, pad: boolean | undefined | null, tabWidth: number): string +/** + * Score every row of a normalized `f32` matrix against `query` and return + * the top `limit` rows. + * + * Mirrors the TS `searchExactVectorIndex` loop bit-exactly: the query is + * normalized by the L2 norm of its *full* length, each row score sums + * `matrix[row][col] * (query[col] / norm)` over + * `min(query.len, dimensions)` columns in column order. Ranking matches the + * TS stable sort: score descending, lower row index first on exact ties + * (`-0.0` and `+0.0` compare equal). Callers are expected to enforce the TS + * guards first (finite query with a positive norm, non-empty matrix). + */ +export declare function vectorIndexTopK(matrix: Float32Array, dimensions: number, query: Float64Array, limit: number): VectorTopK + +/** + * Top-k rows of a normalized vector matrix ranked by dot product with a + * normalized query. + */ +export interface VectorTopK { + /** Row indices of the selected hits, best score first. */ + indices: Uint32Array + /** Scores aligned with `indices`. */ + scores: Float64Array +} + /** * Calculate visible width of text, excluding ANSI escape sequences. * diff --git a/packages/natives/native/index.js b/packages/natives/native/index.js index 27f70ec15..487dbb643 100644 --- a/packages/natives/native/index.js +++ b/packages/natives/native/index.js @@ -30,6 +30,7 @@ export const astGrep = nativeBindings.astGrep; export const astMatch = nativeBindings.astMatch; export const blockRangeAt = nativeBindings.blockRangeAt; export const copyToClipboard = nativeBindings.copyToClipboard; +export const cosineSimilarityPairs = nativeBindings.cosineSimilarityPairs; export const countTokens = nativeBindings.countTokens; export const detectMacOSAppearance = nativeBindings.detectMacOSAppearance; export const diffLineRuns = nativeBindings.diffLineRuns; @@ -59,6 +60,7 @@ export const listWorkspace = nativeBindings.listWorkspace; export const matchesKey = nativeBindings.matchesKey; export const matchesKittySequence = nativeBindings.matchesKittySequence; export const matchesLegacySequence = nativeBindings.matchesLegacySequence; +export const mmrRerankIndices = nativeBindings.mmrRerankIndices; export const parseKey = nativeBindings.parseKey; export const parseKittySequence = nativeBindings.parseKittySequence; export const readImageFromClipboard = nativeBindings.readImageFromClipboard; @@ -71,6 +73,7 @@ export const structuredPatchHunks = nativeBindings.structuredPatchHunks; export const summarizeCode = nativeBindings.summarizeCode; export const supportsLanguage = nativeBindings.supportsLanguage; export const truncateToWidth = nativeBindings.truncateToWidth; +export const vectorIndexTopK = nativeBindings.vectorIndexTopK; export const visibleWidth = nativeBindings.visibleWidth; export const wrapTextWithAnsi = nativeBindings.wrapTextWithAnsi;