smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(mnemopi): keep only production-backed native kernels after wrapper-level benchmarking

Browse Source
This commit is contained in:
Wolfgang Schoenberger
2026-07-22 04:45:13 -07:00
parent 3c84e2c410
commit 4ffe1822d4
7 changed files with 154 additions and 514 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pi-natives/src/vectors.rs
+2 -157
View File
@@ -9,7 +9,7 @@
use napi::{
Error, Result, Status,
bindgen_prelude::{Float32Array, Float64Array, Uint8Array, Uint32Array},
bindgen_prelude::{Float32Array, Float64Array, Uint32Array},
};
use napi_derive::napi;
@@ -63,37 +63,6 @@ fn cosine_one(a: &[f64], b: &[f64]) -> f64 {
dot / (norm_a.sqrt() * norm_b.sqrt())
}
/// Cosine similarity of `query` against a batch of candidate vectors.
///
/// `candidates` is `n` vectors flattened row-major at `dim` elements per row
/// (callers zero-pad shorter vectors, which matches the TS `?? 0` missing
/// element semantics). Returns one score per candidate, bit-identical to
/// calling the TS `cosineSimilarity(query, candidate)` per row.
#[napi]
pub fn cosine_similarity_batch(
query: Float64Array,
candidates: Float64Array,
dim: u32,
) -> Result<Float64Array> {
let dim = dim as usize;
let cands: &[f64] = &candidates;
if dim == 0 {
if !cands.is_empty() {
return invalid("candidates must be empty when dim is 0");
}
return Ok(Float64Array::new(Vec::new()));
}
if !cands.len().is_multiple_of(dim) {
return invalid("candidates length must be a multiple of dim");
}
let q: &[f64] = &query;
let scores: Vec<f64> = cands
.chunks_exact(dim)
.map(|row| cosine_one(q, row))
.collect();
Ok(Float64Array::new(scores))
}
/// All pairs `(i, j)` with `i < j` whose cosine similarity meets `threshold`.
///
/// `vectors` is `count` vectors flattened row-major at `dim` `f64` elements
@@ -205,123 +174,6 @@ pub fn vector_index_top_k(
Ok(VectorTopK { indices: Uint32Array::new(indices), scores: Float64Array::new(scores) })
}
#[inline]
fn hamming_one(query: &[u8], row: &[u8]) -> u32 {
let shared = query.len().min(row.len());
let mut distance = 0u32;
for i in 0..shared {
distance += (query[i] ^ row[i]).count_ones();
}
for &byte in &query[shared..] {
distance += byte.count_ones();
}
for &byte in &row[shared..] {
distance += byte.count_ones();
}
distance
}
/// Hamming distance of `query` against a batch of packed binary vectors.
///
/// `candidates` is flattened row-major at `stride` bytes per row.
/// `lengths[i]` gives the meaningful byte length of row `i` (clamped to
/// `stride`); omit it when every row is exactly `stride` bytes. Semantics
/// match the TS `hammingDistance` exactly, including popcounting the
/// unmatched tail of whichever side is longer.
#[napi]
pub fn hamming_distance_batch(
query: Uint8Array,
candidates: Uint8Array,
stride: u32,
lengths: Option<Uint32Array>,
) -> Result<Uint32Array> {
let stride = stride as usize;
let cands: &[u8] = &candidates;
let q: &[u8] = &query;
let count = match &lengths {
Some(lens) => lens.len(),
None if stride == 0 => {
if cands.is_empty() {
return Ok(Uint32Array::new(Vec::new()));
}
return invalid("stride must be positive when lengths is omitted");
},
None => {
if !cands.len().is_multiple_of(stride) {
return invalid("candidates length must be a multiple of stride");
}
cands.len() / stride
},
};
if count * stride > cands.len() {
return invalid("candidates shorter than count * stride");
}
let mut distances: Vec<u32> = Vec::with_capacity(count);
for i in 0..count {
let len = match &lengths {
Some(lens) => (lens[i] as usize).min(stride),
None => stride,
};
let row = &cands[i * stride..i * stride + len];
distances.push(hamming_one(q, row));
}
Ok(Uint32Array::new(distances))
}
/// Dimension-masked Hamming distance of `query` against a batch of packed
/// binary vectors.
///
/// `candidates` is flattened row-major at `stride` bytes per row and
/// `dims[i]` is the bit dimension compared for row `i`. Bytes beyond either
/// side's data read as `0`, and a trailing partial byte is masked to the top
/// `dims[i] % 8` bits — exactly the TS `hammingDistanceForDimension`.
#[napi]
pub fn hamming_distance_for_dim_batch(
query: Uint8Array,
candidates: Uint8Array,
stride: u32,
dims: Uint32Array,
) -> Result<Uint32Array> {
let stride = stride as usize;
let cands: &[u8] = &candidates;
let q: &[u8] = &query;
let count = dims.len();
if count * stride > cands.len() {
return invalid("candidates shorter than dims.length * stride");
}
let mut distances: Vec<u32> = Vec::with_capacity(count);
for i in 0..count {
let row = &cands[i * stride..(i + 1) * stride];
let dim = dims[i] as usize;
let whole_bytes = dim >> 3;
// Bytes beyond either side's data read as 0, so the XOR reduces to a
// plain popcount over whichever side still has data. Sliced loops keep
// the hot path branch-free and autovectorizable (see `hamming_one`).
let q_end = q.len().min(whole_bytes);
let row_end = row.len().min(whole_bytes);
let shared = q_end.min(row_end);
let mut distance = 0u32;
for byte in 0..shared {
distance += (q[byte] ^ row[byte]).count_ones();
}
for &byte in &q[shared..q_end] {
distance += byte.count_ones();
}
for &byte in &row[shared..row_end] {
distance += byte.count_ones();
}
let remaining_bits = dim & 7;
if remaining_bits > 0 {
let mask = (0xffu16 << (8 - remaining_bits)) as u8;
let a = q.get(whole_bytes).copied().unwrap_or(0);
let b = row.get(whole_bytes).copied().unwrap_or(0);
distance += ((a ^ b) & mask).count_ones();
}
distances.push(distance);
}
Ok(Uint32Array::new(distances))
}
/// ECMA-262 `\s` (`WhiteSpace` ∪ `LineTerminator`), which differs from Rust's
/// `char::is_whitespace` (JS additionally includes U+FEFF).
#[inline]
@@ -456,7 +308,7 @@ pub fn mmr_rerank_indices(
#[cfg(test)]
mod tests {
use super::{cosine_one, hamming_one, is_js_whitespace, jaccard_sorted, word_set};
use super::{cosine_one, is_js_whitespace, jaccard_sorted, word_set};
#[test]
fn cosine_matches_reference_semantics() {
@@ -470,13 +322,6 @@ mod tests {
assert!((sim - expect).abs() < 1e-12);
}
#[test]
fn hamming_counts_unmatched_tails() {
assert_eq!(hamming_one(&[0xff, 0x0f], &[0x0f]), 4 + 4);
assert_eq!(hamming_one(&[], &[0xff]), 8);
assert_eq!(hamming_one(&[0b1010], &[0b0101]), 4);
}
#[test]
fn word_set_matches_js_tokenizer() {
let set = word_set("Hello\u{00a0}WORLD hello\u{feff}world");
packages/mnemopi/bench/native-vectors.bench.json
+144 -144
View File
@@ -1,207 +1,207 @@
{
"sha": "4ba69285dfb939e4fc802fa66c3eac76cd075b3b",
"date": "2026-07-22T09:44:06.872Z",
"sha": "3c84e2c410afae163fb1a778a4f7011354cc893c",
"date": "2026-07-22T11:33:29.804Z",
"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",
"rows": [
{
"kernel": "cosineSimilarityBatch",
"count": 10,
"ts_us": 10.61,
"native_us": 3.35,
"speedup": 3.16,
"ts_us": 9.85,
"native_us": 3.2,
"speedup": 3.08,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityBatch",
"count": 100,
"ts_us": 42.48,
"native_us": 25.04,
"speedup": 1.7,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityBatch",
"count": 1000,
"ts_us": 470.41,
"native_us": 248.07,
"speedup": 1.9,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityBatch",
"count": 10000,
"ts_us": 3963.14,
"native_us": 2454.54,
"speedup": 1.61,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "vectorIndexTopK",
"count": 10,
"ts_us": 5.59,
"native_us": 2.97,
"speedup": 1.88,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "vectorIndexTopK",
"count": 100,
"ts_us": 31.47,
"native_us": 17.37,
"speedup": 1.81,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "vectorIndexTopK",
"count": 1000,
"ts_us": 325.12,
"native_us": 166.88,
"speedup": 1.95,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "vectorIndexTopK",
"count": 10000,
"ts_us": 3774.16,
"native_us": 1869.57,
"speedup": 2.02,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch",
"count": 10,
"ts_us": 3.1,
"native_us": 0.59,
"speedup": 5.27,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch",
"count": 100,
"ts_us": 7.61,
"native_us": 0.74,
"speedup": 10.24,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch",
"count": 1000,
"ts_us": 29.27,
"native_us": 3.83,
"speedup": 7.64,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch",
"count": 10000,
"ts_us": 252.75,
"native_us": 39.02,
"speedup": 6.48,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityPairs",
"count": 10,
"ts_us": 18.55,
"native_us": 11.64,
"speedup": 1.59,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityPairs",
"count": 100,
"ts_us": 1955.97,
"native_us": 1201.4,
"ts_us": 41.15,
"native_us": 25.18,
"speedup": 1.63,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityPairs",
"kernel": "cosineSimilarityBatch",
"count": 1000,
"ts_us": 197715.54,
"native_us": 122119.79,
"ts_us": 398.55,
"native_us": 242.88,
"speedup": 1.64,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityBatch",
"count": 10000,
"ts_us": 3936.21,
"native_us": 2431.53,
"speedup": 1.62,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "searchExactVectorIndex (topK wrapper)",
"count": 10,
"ts_us": 5.82,
"native_us": 6.44,
"speedup": 0.9,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "searchExactVectorIndex (topK wrapper)",
"count": 100,
"ts_us": 36.99,
"native_us": 19.48,
"speedup": 1.9,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "searchExactVectorIndex (topK wrapper)",
"count": 1000,
"ts_us": 310.52,
"native_us": 165.81,
"speedup": 1.87,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "searchExactVectorIndex (topK wrapper)",
"count": 10000,
"ts_us": 3513.57,
"native_us": 1884.54,
"speedup": 1.86,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch (incl. packing)",
"count": 10,
"ts_us": 3.07,
"native_us": 1.39,
"speedup": 2.21,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch (incl. packing)",
"count": 100,
"ts_us": 7.09,
"native_us": 3.97,
"speedup": 1.79,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch (incl. packing)",
"count": 1000,
"ts_us": 34.25,
"native_us": 26.11,
"speedup": 1.31,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceBatch (incl. packing)",
"count": 10000,
"ts_us": 221.43,
"native_us": 207.45,
"speedup": 1.07,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityPairs (incl. flatten+adjacency)",
"count": 10,
"ts_us": 246.34,
"native_us": 22.08,
"speedup": 11.16,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityPairs (incl. flatten+adjacency)",
"count": 100,
"ts_us": 5767.19,
"native_us": 1381.02,
"speedup": 4.18,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "cosineSimilarityPairs (incl. flatten+adjacency)",
"count": 1000,
"ts_us": 530658.9,
"native_us": 131050.53,
"speedup": 4.05,
"ts_iterations": 10,
"native_iterations": 10
},
{
"kernel": "hammingDistanceForDimBatch",
"kernel": "hammingDistanceForDimBatch (incl. packing)",
"count": 10,
"ts_us": 2.45,
"native_us": 0.6,
"speedup": 4.12,
"ts_us": 2.52,
"native_us": 1.95,
"speedup": 1.29,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceForDimBatch",
"kernel": "hammingDistanceForDimBatch (incl. packing)",
"count": 100,
"ts_us": 5.83,
"native_us": 0.92,
"speedup": 6.33,
"ts_us": 5.88,
"native_us": 4.16,
"speedup": 1.41,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceForDimBatch",
"kernel": "hammingDistanceForDimBatch (incl. packing)",
"count": 1000,
"ts_us": 21.69,
"native_us": 5.16,
"speedup": 4.2,
"ts_us": 27.69,
"native_us": 28.69,
"speedup": 0.97,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "hammingDistanceForDimBatch",
"kernel": "hammingDistanceForDimBatch (incl. packing)",
"count": 10000,
"ts_us": 202.24,
"native_us": 68.09,
"speedup": 2.97,
"ts_us": 231.52,
"native_us": 235.98,
"speedup": 0.98,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "mmrRerankIndices (via mmrRerank)",
"count": 10,
"ts_us": 302.81,
"native_us": 14.38,
"speedup": 21.05,
"ts_us": 331.56,
"native_us": 15.1,
"speedup": 21.95,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "mmrRerankIndices (via mmrRerank)",
"count": 100,
"ts_us": 7775.68,
"native_us": 197.42,
"speedup": 39.39,
"ts_us": 7965.61,
"native_us": 220.67,
"speedup": 36.1,
"ts_iterations": 200,
"native_iterations": 200
},
{
"kernel": "mmrRerankIndices (via mmrRerank)",
"count": 1000,
"ts_us": 73795.1,
"native_us": 2331.18,
"speedup": 31.66,
"ts_us": 78042.9,
"native_us": 2639.47,
"speedup": 29.57,
"ts_iterations": 200,
"native_iterations": 200
}
],
"sink": 823409556.8853176
"sink": 823377224.8853176
}
packages/mnemopi/bench/native-vectors.bench.ts
-75
View File
@@ -6,12 +6,8 @@
* Run from the repo root: `bun packages/mnemopi/bench/native-vectors.bench.ts`
*/
import {
cosineSimilarityBatch,
cosineSimilarityPairs,
hammingDistanceBatch,
hammingDistanceForDimBatch,
} from "@oh-my-pi/pi-natives";
import { hammingDistance, hammingDistanceForDimension } from "../src/core/binary-vectors";
import { jaccardSimilarity, mmrRerank } from "../src/core/mmr";
import { searchExactVectorIndex } from "../src/core/vector-index";
import { cosineSimilarity } from "../src/core/vector-math";
@@ -64,22 +60,6 @@ function pushRow(kernel: string, count: number, ts: { ns: number; iterations: nu
const rows: Row[] = [];
const rng = makeRng(0xbe4c4);
for (const count of COUNTS) {
const query = Float64Array.from({ length: DIM }, () => rng() * 2 - 1);
const flat = new Float64Array(count * DIM);
for (let i = 0; i < flat.length; i += 1) flat[i] = rng() * 2 - 1;
const ts = timeNs(() => {
for (let row = 0; row < count; row += 1) {
sink += cosineSimilarity(query, flat.subarray(row * DIM, (row + 1) * DIM));
}
});
const native = timeNs(() => {
sink += cosineSimilarityBatch(query, flat, DIM)[0] ?? 0;
});
pushRow("cosineSimilarityBatch", count, ts, native);
}
// 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.
@@ -113,33 +93,6 @@ for (const count of COUNTS) {
pushRow("searchExactVectorIndex (topK wrapper)", count, ts, native);
}
// hammingDistanceBatch: the native side pays FastBinarySearch.search's per-call
// packing (stride scan, packed/lengths allocation, full copy) before crossing.
for (const count of COUNTS) {
const query = Uint8Array.from({ length: STRIDE }, () => Math.floor(rng() * 256));
const backing = new Uint8Array(count * STRIDE);
for (let i = 0; i < backing.length; i += 1) backing[i] = Math.floor(rng() * 256);
const vectors: Uint8Array[] = [];
for (let i = 0; i < count; i += 1) vectors.push(backing.subarray(i * STRIDE, (i + 1) * STRIDE));
const ts = timeNs(() => {
for (let i = 0; i < count; i += 1) sink += hammingDistance(query, vectors[i] ?? new Uint8Array());
});
const native = timeNs(() => {
let stride = 0;
for (const vector of vectors) if (vector.length > stride) stride = vector.length;
const packed = new Uint8Array(vectors.length * stride);
const lengths = new Uint32Array(vectors.length);
for (let i = 0; i < vectors.length; i += 1) {
const vector = vectors[i] ?? new Uint8Array();
lengths[i] = vector.length;
packed.set(vector, i * stride);
}
sink += hammingDistanceBatch(query, packed, stride, lengths)[0] ?? 0;
});
pushRow("hammingDistanceBatch (incl. packing)", 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
@@ -192,34 +145,6 @@ for (const count of COUNTS.filter(n => n <= 1000)) {
pushRow("cosineSimilarityPairs (incl. flatten+adjacency)", count, ts, native);
}
// hammingDistanceForDimBatch: ragged/dim-masked variant used by
// BinaryVectorStore.search; the native side pays the production per-call
// packed/dims allocation and copy before crossing.
for (const count of COUNTS) {
const query = Uint8Array.from({ length: STRIDE }, () => Math.floor(rng() * 256));
const packed = new Uint8Array(count * STRIDE);
for (let i = 0; i < packed.length; i += 1) packed[i] = Math.floor(rng() * 256);
const dims = Uint32Array.from({ length: count }, () => (rng() < 0.5 ? DIM : DIM / 2));
const vectors: Uint8Array[] = [];
for (let i = 0; i < count; i += 1) vectors.push(packed.subarray(i * STRIDE, (i + 1) * STRIDE));
const ts = timeNs(() => {
for (let i = 0; i < count; i += 1) {
sink += hammingDistanceForDimension(query, vectors[i] ?? new Uint8Array(), dims[i] ?? DIM);
}
});
const native = timeNs(() => {
const packed2 = new Uint8Array(vectors.length * STRIDE);
const dims2 = new Uint32Array(vectors.length);
for (let i = 0; i < vectors.length; i += 1) {
const vector = vectors[i] ?? new Uint8Array();
dims2[i] = dims[i] ?? DIM;
packed2.set(vector, i * STRIDE);
}
sink += hammingDistanceForDimBatch(query, packed2, STRIDE, dims2)[0] ?? 0;
});
pushRow("hammingDistanceForDimBatch (incl. packing)", 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
packages/mnemopi/src/core/binary-vectors.ts
+7 -37
View File
@@ -1,5 +1,4 @@
import type { Database } from "bun:sqlite";
import { hammingDistanceBatch, hammingDistanceForDimBatch } from "@oh-my-pi/pi-natives";
import { embeddingDim, type VecType } from "../config";
import { closeQuietly, type DatabasePath, openDatabase } from "../db";
@@ -148,7 +147,7 @@ export function hammingDistance(binaryA: Uint8Array | ArrayBuffer, binaryB: Uint
return distance;
}
export function hammingDistanceForDimension(
function hammingDistanceForDimension(
binaryA: Uint8Array | ArrayBuffer,
binaryB: Uint8Array | ArrayBuffer,
dim: number,
@@ -227,31 +226,15 @@ export class BinaryVectorStore {
const rows = this.conn
.query(`SELECT memory_id, binary_vector, original_dim, magnitude FROM ${this.tableName}`)
.all() as VectorRow[];
// Native batch kernel: pack all rows at a fixed stride (zero-padded,
// matching the TS `?? 0` out-of-range reads) and compute every
// dimension-masked distance in one N-API crossing.
const stride = BYTES_PER_VECTOR;
const packed = new Uint8Array(rows.length * stride);
const dims = new Uint32Array(rows.length);
const comparedDims: number[] = new Array(rows.length);
for (let i = 0; i < rows.length; i += 1) {
const row = rows[i];
if (row === undefined) continue;
const results: BinaryVectorSearchResult[] = [];
for (const row of rows) {
const storedDim = Math.max(0, Math.min(EMBEDDING_DIM, Math.trunc(toFiniteNumber(row.original_dim))));
const comparedDim = Math.min(queryDim, storedDim);
comparedDims[i] = comparedDim;
dims[i] = comparedDim;
const bytes = bytesFromBlob(row.binary_vector);
packed.set(bytes.length > stride ? bytes.subarray(0, stride) : bytes, i * stride);
}
const distances = hammingDistanceForDimBatch(queryBinary, packed, stride, dims);
const results: BinaryVectorSearchResult[] = [];
for (let i = 0; i < rows.length; i += 1) {
const distance = distances[i] ?? 0;
const distance = hammingDistanceForDimension(queryBinary, bytesFromBlob(row.binary_vector), comparedDim);
results.push({
memory_id: rows[i]?.memory_id ?? "",
memory_id: row.memory_id,
distance,
score: informationTheoreticScore(distance, comparedDims[i] ?? 0),
score: informationTheoreticScore(distance, comparedDim),
});
}
results.sort((a, b) => b.score - a.score || a.memory_id.localeCompare(b.memory_id));
@@ -319,22 +302,9 @@ export class FastBinarySearch {
search(queryBinary: Uint8Array | ArrayBuffer, topK = 10): BinaryVectorSearchResult[] {
const query = queryBinary instanceof Uint8Array ? queryBinary : new Uint8Array(queryBinary);
// Native batch kernel: pack all vectors at the max byte length and
// compute every Hamming distance in one N-API crossing. Per-row
// lengths preserve the TS unmatched-tail popcount semantics.
let stride = 0;
for (const vector of this.vectors) if (vector.length > stride) stride = vector.length;
const packed = new Uint8Array(this.vectors.length * stride);
const lengths = new Uint32Array(this.vectors.length);
for (let i = 0; i < this.vectors.length; i += 1) {
const vector = this.vectors[i] ?? new Uint8Array();
lengths[i] = vector.length;
packed.set(vector, i * stride);
}
const distances = hammingDistanceBatch(query, packed, stride, lengths);
const results: BinaryVectorSearchResult[] = [];
for (let i = 0; i < this.vectors.length; i += 1) {
const distance = distances[i] ?? 0;
const distance = hammingDistance(query, this.vectors[i] ?? new Uint8Array());
results.push({
memory_id: this.memoryIds[i] ?? "",
distance,
packages/mnemopi/test/native-vector-parity.test.ts
+1 -66
View File
@@ -1,13 +1,5 @@
import { describe, expect, test } from "bun:test";
import {
cosineSimilarityBatch,
cosineSimilarityPairs,
hammingDistanceBatch,
hammingDistanceForDimBatch,
mmrRerankIndices,
vectorIndexTopK,
} from "@oh-my-pi/pi-natives";
import { hammingDistance, hammingDistanceForDimension } from "../src/core/binary-vectors";
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";
@@ -32,24 +24,6 @@ function expectClose(actual: number, expected: number): void {
}
describe("native vector kernel parity", () => {
test("cosineSimilarityBatch matches TS cosineSimilarity per row", () => {
const rng = makeRng(0xc051e);
const dim = 384;
const count = 200;
const query = Float64Array.from({ length: dim }, () => rng() * 2 - 1);
const candidates = new Float64Array(count * dim);
for (let i = 0; i < candidates.length; i += 1) candidates[i] = rng() * 2 - 1;
// Sprinkle non-finite values to exercise the finite_or_zero path.
candidates[3] = Number.NaN;
candidates[dim + 7] = Number.POSITIVE_INFINITY;
const scores = cosineSimilarityBatch(query, candidates, dim);
expect(scores.length).toBe(count);
for (let row = 0; row < count; row += 1) {
const expected = cosineSimilarity(query, candidates.subarray(row * dim, (row + 1) * dim));
expectClose(scores[row] ?? Number.NaN, expected);
}
});
test("cosineSimilarityPairs matches the TS pairwise threshold loop", () => {
const rng = makeRng(0x9a175);
const dim = 64;
@@ -98,45 +72,6 @@ describe("native vector kernel parity", () => {
}
});
test("hammingDistanceBatch is exactly equal to TS hammingDistance", () => {
const rng = makeRng(0xba7c4);
const stride = 48; // 384-dim binarized
const count = 128;
const query = Uint8Array.from({ length: stride }, () => Math.floor(rng() * 256));
const packed = new Uint8Array(count * stride);
const lengths = new Uint32Array(count);
const vectors: Uint8Array[] = [];
for (let i = 0; i < count; i += 1) {
const len = i % 7 === 0 ? Math.floor(rng() * stride) : stride; // ragged rows
const vector = Uint8Array.from({ length: len }, () => Math.floor(rng() * 256));
vectors.push(vector);
lengths[i] = len;
packed.set(vector, i * stride);
}
const distances = hammingDistanceBatch(query, packed, stride, lengths);
for (let i = 0; i < count; i += 1) {
expect(distances[i]).toBe(hammingDistance(query, vectors[i] ?? new Uint8Array()));
}
});
test("hammingDistanceForDimBatch is exactly equal to TS hammingDistanceForDimension", () => {
const rng = makeRng(0xd1235);
const stride = 48;
const count = 96;
const query = Uint8Array.from({ length: stride }, () => Math.floor(rng() * 256));
const packed = new Uint8Array(count * stride);
const dims = new Uint32Array(count);
for (let i = 0; i < count; i += 1) {
dims[i] = Math.floor(rng() * (stride * 8 + 1)); // includes partial-byte tails and 0
for (let b = 0; b < stride; b += 1) packed[i * stride + b] = Math.floor(rng() * 256);
}
const distances = hammingDistanceForDimBatch(query, packed, stride, dims);
for (let i = 0; i < count; i += 1) {
const row = packed.subarray(i * stride, (i + 1) * stride);
expect(distances[i]).toBe(hammingDistanceForDimension(query, row, dims[i] ?? 0));
}
});
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"];
packages/natives/native/index.d.ts
Vendored
-32
View File
@@ -470,16 +470,6 @@ export interface ContextLine {
*/
export declare function copyToClipboard(text: string): void
/**
* Cosine similarity of `query` against a batch of candidate vectors.
*
* `candidates` is `n` vectors flattened row-major at `dim` elements per row
* (callers zero-pad shorter vectors, which matches the TS `?? 0` missing
* element semantics). Returns one score per candidate, bit-identical to
* calling the TS `cosineSimilarity(query, candidate)` per row.
*/
export declare function cosineSimilarityBatch(query: Float64Array, candidates: Float64Array, dim: number): Float64Array
/**
* All pairs `(i, j)` with `i < j` whose cosine similarity meets `threshold`.
*
@@ -830,28 +820,6 @@ export interface GrepResult {
skippedOversized?: number
}
/**
* Hamming distance of `query` against a batch of packed binary vectors.
*
* `candidates` is flattened row-major at `stride` bytes per row.
* `lengths[i]` gives the meaningful byte length of row `i` (clamped to
* `stride`); omit it when every row is exactly `stride` bytes. Semantics
* match the TS `hammingDistance` exactly, including popcounting the
* unmatched tail of whichever side is longer.
*/
export declare function hammingDistanceBatch(query: Uint8Array, candidates: Uint8Array, stride: number, lengths?: Uint32Array | undefined | null): Uint32Array
/**
* Dimension-masked Hamming distance of `query` against a batch of packed
* binary vectors.
*
* `candidates` is flattened row-major at `stride` bytes per row and
* `dims[i]` is the bit dimension compared for row `i`. Bytes beyond either
* side's data read as `0`, and a trailing partial byte is masked to the top
* `dims[i] % 8` bits — exactly the TS `hammingDistanceForDimension`.
*/
export declare function hammingDistanceForDimBatch(query: Uint8Array, candidates: Uint8Array, stride: number, dims: Uint32Array): Uint32Array
/**
* Quick check if content matches a pattern.
*
packages/natives/native/index.js
-3
View File
@@ -30,7 +30,6 @@ export const astGrep = nativeBindings.astGrep;
export const astMatch = nativeBindings.astMatch;
export const blockRangeAt = nativeBindings.blockRangeAt;
export const copyToClipboard = nativeBindings.copyToClipboard;
export const cosineSimilarityBatch = nativeBindings.cosineSimilarityBatch;
export const cosineSimilarityPairs = nativeBindings.cosineSimilarityPairs;
export const countTokens = nativeBindings.countTokens;
export const detectMacOSAppearance = nativeBindings.detectMacOSAppearance;
@@ -43,8 +42,6 @@ export const getSupportedLanguages = nativeBindings.getSupportedLanguages;
export const getWorkProfile = nativeBindings.getWorkProfile;
export const glob = nativeBindings.glob;
export const grep = nativeBindings.grep;
export const hammingDistanceBatch = nativeBindings.hammingDistanceBatch;
export const hammingDistanceForDimBatch = nativeBindings.hammingDistanceForDimBatch;
export const hasMatch = nativeBindings.hasMatch;
export const highlightCode = nativeBindings.highlightCode;
export const htmlToMarkdown = nativeBindings.htmlToMarkdown;