sweet-search 2.6.10 → 2.6.11
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/core/graph/graph-expansion.js +48 -6
- package/core/incremental-indexing/application/production-reconciler-helpers.mjs +48 -11
- package/core/incremental-indexing/application/production-reconciler.mjs +12 -3
- package/core/infrastructure/hamming-kernel.js +293 -0
- package/core/infrastructure/maxsim.wasm +0 -0
- package/core/infrastructure/simd-distance.js +39 -6
- package/core/infrastructure/simd-distance.wasm +0 -0
- package/core/infrastructure/tombstone-bitmap-reader.js +34 -0
- package/core/ranking/late-interaction-index.js +31 -3
- package/core/search/search-semantic.js +8 -1
- package/core/vector-store/binary-heap.js +12 -3
- package/core/vector-store/binary-hnsw-index.js +354 -107
- package/package.json +8 -8
|
@@ -249,6 +249,18 @@ const TYPE_BOOST = {
|
|
|
249
249
|
struct: 1.2,
|
|
250
250
|
};
|
|
251
251
|
|
|
252
|
+
/**
|
|
253
|
+
* Fetch int8 vectors for many ids with ONE stale-bitmap snapshot when the
|
|
254
|
+
* index supports it (BinaryHNSWIndex.getInt8VectorsForIds); falls back to
|
|
255
|
+
* per-id lookups for index doubles that only expose getInt8Vector.
|
|
256
|
+
*/
|
|
257
|
+
function batchInt8Vectors(hnswIndex, ids) {
|
|
258
|
+
if (typeof hnswIndex.getInt8VectorsForIds === 'function') {
|
|
259
|
+
return hnswIndex.getInt8VectorsForIds(ids);
|
|
260
|
+
}
|
|
261
|
+
return ids.map((id) => hnswIndex.getInt8Vector(id));
|
|
262
|
+
}
|
|
263
|
+
|
|
252
264
|
function clampSemanticWeight(value) {
|
|
253
265
|
if (!Number.isFinite(value)) return 0.4;
|
|
254
266
|
return Math.max(0, Math.min(1, value));
|
|
@@ -614,21 +626,27 @@ export function expandSecondHop(db, seedIds, expanded, edgeTypes, options = {})
|
|
|
614
626
|
}
|
|
615
627
|
|
|
616
628
|
const excluded = new Set([...seedIds, ...expanded.keys()]);
|
|
617
|
-
const
|
|
629
|
+
const eligible = [];
|
|
618
630
|
for (const rel of hop2Forward) {
|
|
619
631
|
if (!edgeTypes.has(rel.type) || excluded.has(rel.target_id)) continue;
|
|
620
632
|
|
|
621
633
|
const hop1Entry = expanded.get(rel.source_id);
|
|
622
634
|
const hop1Score = hop1Entry?.score ?? 1; // identity: preserves old edgePriority × weight
|
|
623
635
|
const graphScore = hop1Score * (EDGE_PRIORITY[rel.type] || 1) * (rel.weight || 1.0);
|
|
636
|
+
eligible.push({ rel, graphScore });
|
|
637
|
+
}
|
|
638
|
+
// Batch int8 lookup: one stale-bitmap snapshot for all targets instead of
|
|
639
|
+
// one stat per target.
|
|
640
|
+
const eligibleInt8 = batchInt8Vectors(hnswIndex, eligible.map((e) => e.rel.target_id));
|
|
641
|
+
const candidates = eligible.map((e, i) => {
|
|
642
|
+
const entityInt8 = eligibleInt8[i];
|
|
624
643
|
let normSim = null;
|
|
625
|
-
const entityInt8 = hnswIndex.getInt8Vector(rel.target_id);
|
|
626
644
|
if (entityInt8) {
|
|
627
645
|
const cosSim = cosineSimilarity(queryInt8, entityInt8);
|
|
628
646
|
normSim = (cosSim + 1) / 2;
|
|
629
647
|
}
|
|
630
|
-
|
|
631
|
-
}
|
|
648
|
+
return { rel: e.rel, graphScore: e.graphScore, normSim };
|
|
649
|
+
});
|
|
632
650
|
|
|
633
651
|
if (candidates.length === 0) return;
|
|
634
652
|
const normalizedGraphScores = normalizeMinMax(candidates.map(c => c.graphScore));
|
|
@@ -729,6 +747,25 @@ export function expandSecondHopAdaptive(db, seedIds, hop1Expanded, edgeTypes, op
|
|
|
729
747
|
const vectorCache = semanticEnabled ? new Map() : null;
|
|
730
748
|
const scoredCandidates = [];
|
|
731
749
|
|
|
750
|
+
// Pre-seed the vector cache with ONE batched int8 fetch (single
|
|
751
|
+
// stale-bitmap snapshot) so the scoring loop below never does per-target
|
|
752
|
+
// lookups.
|
|
753
|
+
if (semanticEnabled) {
|
|
754
|
+
const uniqueTargets = [];
|
|
755
|
+
for (const c of rawCandidates) {
|
|
756
|
+
if (!edgeTypes.has(c.type) || excluded.has(c.target_id)) continue;
|
|
757
|
+
if (vectorCache.has(c.target_id)) continue;
|
|
758
|
+
vectorCache.set(c.target_id, undefined);
|
|
759
|
+
uniqueTargets.push(c.target_id);
|
|
760
|
+
}
|
|
761
|
+
if (uniqueTargets.length > 0) {
|
|
762
|
+
const fetched = batchInt8Vectors(hnswIndex, uniqueTargets);
|
|
763
|
+
for (let i = 0; i < uniqueTargets.length; i++) {
|
|
764
|
+
vectorCache.set(uniqueTargets[i], fetched[i]);
|
|
765
|
+
}
|
|
766
|
+
}
|
|
767
|
+
}
|
|
768
|
+
|
|
732
769
|
for (const c of rawCandidates) {
|
|
733
770
|
if (!edgeTypes.has(c.type) || excluded.has(c.target_id)) continue;
|
|
734
771
|
|
|
@@ -743,6 +780,7 @@ export function expandSecondHopAdaptive(db, seedIds, hop1Expanded, edgeTypes, op
|
|
|
743
780
|
let normSim = null;
|
|
744
781
|
if (semanticEnabled) {
|
|
745
782
|
if (!vectorCache.has(c.target_id)) {
|
|
783
|
+
// Miss fallback only — the pre-seed pass below batch-fills the cache.
|
|
746
784
|
vectorCache.set(c.target_id, hnswIndex.getInt8Vector(c.target_id));
|
|
747
785
|
}
|
|
748
786
|
const entityInt8 = vectorCache.get(c.target_id);
|
|
@@ -929,12 +967,16 @@ export function rerankExpanded(expandedResults, seedResults, options = {}) {
|
|
|
929
967
|
}
|
|
930
968
|
} else {
|
|
931
969
|
const normalizedGraphScores = normalizeMinMax(baseScores);
|
|
970
|
+
// Batch int8 lookup: one stale-bitmap snapshot for the whole result set.
|
|
971
|
+
const entityInt8s = batchInt8Vectors(
|
|
972
|
+
hnswIndex,
|
|
973
|
+
expandedResults.map((er) => er.entity_id || er.id)
|
|
974
|
+
);
|
|
932
975
|
for (let i = 0; i < expandedResults.length; i++) {
|
|
933
976
|
const er = expandedResults[i];
|
|
934
977
|
const normGraph = normalizedGraphScores[i];
|
|
935
978
|
let rerankScore = normGraph;
|
|
936
|
-
const
|
|
937
|
-
const entityInt8 = hnswIndex.getInt8Vector(entityId);
|
|
979
|
+
const entityInt8 = entityInt8s[i];
|
|
938
980
|
if (entityInt8) {
|
|
939
981
|
const cosSim = cosineSimilarity(queryInt8, entityInt8);
|
|
940
982
|
const normSim = (cosSim + 1) / 2;
|
|
@@ -60,19 +60,56 @@ export function insertRelationships(db, relationships, liveIdFor, epoch) {
|
|
|
60
60
|
}
|
|
61
61
|
}
|
|
62
62
|
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
63
|
+
/**
|
|
64
|
+
* Batched stale marking. The per-id `markBinaryStale` loads, mutates, and
|
|
65
|
+
* fsyncs the bitmap file once PER RETIRED ID — O(retires × bitmap size) with
|
|
66
|
+
* an fsync each. A batch loads the bitmap once (lazily, on the first mark),
|
|
67
|
+
* applies every mark in memory with the exact per-op semantics (idToIndex /
|
|
68
|
+
* int8Vectors pruned immediately, reader cache invalidated), and persists
|
|
69
|
+
* once in flush(). End state on disk is identical to N sequential
|
|
70
|
+
* markBinaryStale calls.
|
|
71
|
+
*
|
|
72
|
+
* Crash window: marks staged between flush()es are lost on crash, which the
|
|
73
|
+
* reconciler's persist-before-advance gate already handles — an unflushed
|
|
74
|
+
* tick replays its ops.
|
|
75
|
+
*/
|
|
76
|
+
export function createStaleBatch(index) {
|
|
66
77
|
const stalePath = index.stalePath || `${index.indexPath}.stale.bin`;
|
|
67
78
|
let bitmap = null;
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
79
|
+
let loaded = false;
|
|
80
|
+
let dirty = false;
|
|
81
|
+
return {
|
|
82
|
+
markStale(id) {
|
|
83
|
+
const idx = index.idToIndex.get(id);
|
|
84
|
+
if (idx == null) return false;
|
|
85
|
+
if (!loaded) {
|
|
86
|
+
try { bitmap = loadBitmap(stalePath); } catch {}
|
|
87
|
+
loaded = true;
|
|
88
|
+
}
|
|
89
|
+
bitmap = bitmap
|
|
90
|
+
? resizeBitmap(bitmap, Math.max(idx + 1, index.vectors.length, 1))
|
|
91
|
+
: createBitmap(Math.max(idx + 1, index.vectors.length, 1));
|
|
92
|
+
setBit(bitmap, idx);
|
|
93
|
+
dirty = true;
|
|
94
|
+
index.idToIndex.delete(id);
|
|
95
|
+
index.int8Vectors.delete(id);
|
|
96
|
+
index._staleBitmapCache = null;
|
|
97
|
+
return true;
|
|
98
|
+
},
|
|
99
|
+
flush() {
|
|
100
|
+
if (!dirty) return false;
|
|
101
|
+
saveBitmap(stalePath, bitmap);
|
|
102
|
+
dirty = false;
|
|
103
|
+
return true;
|
|
104
|
+
},
|
|
105
|
+
};
|
|
106
|
+
}
|
|
107
|
+
|
|
108
|
+
export function markBinaryStale(index, id) {
|
|
109
|
+
const batch = createStaleBatch(index);
|
|
110
|
+
const marked = batch.markStale(id);
|
|
111
|
+
batch.flush();
|
|
112
|
+
return marked;
|
|
76
113
|
}
|
|
77
114
|
|
|
78
115
|
/**
|
|
@@ -13,7 +13,7 @@ import { readManifest, writeManifest } from '../infrastructure/manifest.mjs';
|
|
|
13
13
|
import { annotateChunksForDelta, snapshotFileRows, diffChunks, applyDiff } from '../infrastructure/vector-delta-writer.mjs';
|
|
14
14
|
import { appendDeltaRecord, FALLBACK_WEIGHTS_ID, fileIdFor, listDeltaSegments } from '../infrastructure/sparse-gram-delta.mjs';
|
|
15
15
|
import { fts5Merge, fts5MergeBudgetPages } from '../infrastructure/sqlite-fts5.mjs';
|
|
16
|
-
import { insertEntity, insertRelationships, markBinaryStale, maintainFloatStore, flushFloatStore } from './production-reconciler-helpers.mjs';
|
|
16
|
+
import { insertEntity, insertRelationships, markBinaryStale, createStaleBatch, maintainFloatStore, flushFloatStore } from './production-reconciler-helpers.mjs';
|
|
17
17
|
import {
|
|
18
18
|
chunkCutoffEnabled,
|
|
19
19
|
computeCutoffSignature,
|
|
@@ -1110,9 +1110,12 @@ class ProductionReconcileAdapter {
|
|
|
1110
1110
|
if (ctx?.index) {
|
|
1111
1111
|
const index = ctx.index;
|
|
1112
1112
|
let append = 0; let tombstone = 0;
|
|
1113
|
+
// One bitmap load + one save+fsync per file's ops instead of per
|
|
1114
|
+
// retired id; marks apply in op order with identical semantics.
|
|
1115
|
+
const staleBatch = createStaleBatch(index);
|
|
1113
1116
|
for (const op of ops) {
|
|
1114
1117
|
if (op.retireId) {
|
|
1115
|
-
if (
|
|
1118
|
+
if (staleBatch.markStale(op.retireId)) tombstone += 1;
|
|
1116
1119
|
ctx.floatRemoveIds.push(op.retireId);
|
|
1117
1120
|
}
|
|
1118
1121
|
if (op.addId && op.embedding) {
|
|
@@ -1122,6 +1125,7 @@ class ProductionReconcileAdapter {
|
|
|
1122
1125
|
append += 1;
|
|
1123
1126
|
}
|
|
1124
1127
|
}
|
|
1128
|
+
staleBatch.flush();
|
|
1125
1129
|
ctx.tombstone += tombstone;
|
|
1126
1130
|
// append is committed to ctx.append in finalize after the sorted inserts.
|
|
1127
1131
|
ctx.append += append;
|
|
@@ -1137,9 +1141,13 @@ class ProductionReconcileAdapter {
|
|
|
1137
1141
|
let append = 0; let tombstone = 0;
|
|
1138
1142
|
const floatUpserts = [];
|
|
1139
1143
|
const floatRemoveIds = [];
|
|
1144
|
+
// One bitmap load + one save+fsync for the whole op list instead of per
|
|
1145
|
+
// retired id. Marks still apply inline in op order (adds never consult
|
|
1146
|
+
// the bitmap), so retire-then-readd sequences behave exactly as before.
|
|
1147
|
+
const staleBatch = createStaleBatch(index);
|
|
1140
1148
|
for (const op of ops) {
|
|
1141
1149
|
if (op.retireId) {
|
|
1142
|
-
if (
|
|
1150
|
+
if (staleBatch.markStale(op.retireId)) tombstone += 1;
|
|
1143
1151
|
floatRemoveIds.push(op.retireId);
|
|
1144
1152
|
}
|
|
1145
1153
|
if (op.addId && op.embedding) {
|
|
@@ -1150,6 +1158,7 @@ class ProductionReconcileAdapter {
|
|
|
1150
1158
|
}
|
|
1151
1159
|
if ((append + tombstone) > 0 && (append + tombstone) % 100 === 0) this.progress('production:binary-hnsw-loop');
|
|
1152
1160
|
}
|
|
1161
|
+
staleBatch.flush();
|
|
1153
1162
|
await index.save(indexPath);
|
|
1154
1163
|
this.progress('production:binary-hnsw-saved');
|
|
1155
1164
|
await maintainFloatStore(indexPath, { upserts: floatUpserts, removeIds: floatRemoveIds, binaryVectorsBefore, dimension: this.modelInfo.hnswDimension });
|
|
@@ -0,0 +1,293 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* Resident-slab Hamming kernel for the binary HNSW hot path.
|
|
3
|
+
*
|
|
4
|
+
* The legacy wasmHammingDistance() copies BOTH 64-byte operands into WASM
|
|
5
|
+
* linear memory on every call — at M=64/ef=400 a query pays thousands of
|
|
6
|
+
* memcpy+FFI round-trips, construction ~10⁴–10⁵ per insert. This kernel keeps
|
|
7
|
+
* every indexed vector RESIDENT in its own WASM instance (mirrored from a
|
|
8
|
+
* contiguous JS slab), so a distance is a single FFI call with zero copies:
|
|
9
|
+
* measured ~85ns/dist vs ~490ns/dist for the copy-per-call path (M3 Max,
|
|
10
|
+
* 64-byte vectors). When WASM is unavailable it falls back to an unrolled
|
|
11
|
+
* SWAR popcount over the JS slab (~190ns/dist vs ~600ns for the byte-LUT).
|
|
12
|
+
*
|
|
13
|
+
* The JS slab is the source of truth; per-vector Uint8Array subarray views
|
|
14
|
+
* keep the existing `vectors[i].binary` contract intact. The WASM mirror is
|
|
15
|
+
* a private instance of simd-distance.wasm (its exported memory is module-
|
|
16
|
+
* internal), so the shared instance used by wasmInt8BatchDot & friends keeps
|
|
17
|
+
* its scratch-at-offset-0 layout untouched.
|
|
18
|
+
*
|
|
19
|
+
* All arithmetic is integer Hamming — every path returns bit-exact results.
|
|
20
|
+
*/
|
|
21
|
+
|
|
22
|
+
import { readFileSync } from 'fs';
|
|
23
|
+
import { fileURLToPath } from 'url';
|
|
24
|
+
import { dirname, join } from 'path';
|
|
25
|
+
|
|
26
|
+
// WASM memory layout (all offsets fixed so memory.grow never moves them):
|
|
27
|
+
// [0, 256) staged query (256B covers 2048-bit vectors)
|
|
28
|
+
// [256, 256+4*CAP) batch index scratch (u32 node indices)
|
|
29
|
+
// [.., +4*CAP) batch distance output (u32)
|
|
30
|
+
// [SLAB_BASE, ...) resident vector slab
|
|
31
|
+
// BATCH_CAP is far above the max graph degree (2*M0 ≈ 256), so a whole
|
|
32
|
+
// neighbor block always fits in one hamming_batch call.
|
|
33
|
+
const QUERY_SLOT_BYTES = 256;
|
|
34
|
+
const BATCH_CAP = 4096;
|
|
35
|
+
const BATCH_IDX_PTR = QUERY_SLOT_BYTES;
|
|
36
|
+
const BATCH_OUT_PTR = BATCH_IDX_PTR + BATCH_CAP * 4;
|
|
37
|
+
const SLAB_BASE = BATCH_OUT_PTR + BATCH_CAP * 4;
|
|
38
|
+
const WASM_PAGE = 65536;
|
|
39
|
+
|
|
40
|
+
let cachedModule; // WebAssembly.Module | null once resolved
|
|
41
|
+
|
|
42
|
+
function getWasmModule() {
|
|
43
|
+
if (cachedModule !== undefined) return cachedModule;
|
|
44
|
+
try {
|
|
45
|
+
const here = dirname(fileURLToPath(import.meta.url));
|
|
46
|
+
const bytes = readFileSync(join(here, 'simd-distance.wasm'));
|
|
47
|
+
// Sync compile is fine: the module is ~1KB.
|
|
48
|
+
cachedModule = new WebAssembly.Module(bytes);
|
|
49
|
+
} catch {
|
|
50
|
+
cachedModule = null;
|
|
51
|
+
}
|
|
52
|
+
return cachedModule;
|
|
53
|
+
}
|
|
54
|
+
|
|
55
|
+
// SWAR popcount of one u32 (subtract-shift-mask + multiply reduction).
|
|
56
|
+
// Exact for all inputs.
|
|
57
|
+
function popcnt32(x) {
|
|
58
|
+
x -= (x >>> 1) & 0x55555555;
|
|
59
|
+
x = (x & 0x33333333) + ((x >>> 2) & 0x33333333);
|
|
60
|
+
x = (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
61
|
+
return Math.imul(x, 0x01010101) >>> 24;
|
|
62
|
+
}
|
|
63
|
+
|
|
64
|
+
export class HammingSlab {
|
|
65
|
+
/**
|
|
66
|
+
* @param {number} dimension - vector size in BYTES
|
|
67
|
+
*/
|
|
68
|
+
constructor(dimension) {
|
|
69
|
+
this.dim = dimension;
|
|
70
|
+
this.words = dimension >>> 2; // full u32 words
|
|
71
|
+
this.tailBytes = dimension & 3; // non-multiple-of-4 remainder
|
|
72
|
+
this.capacity = 0; // vectors
|
|
73
|
+
this.count = 0;
|
|
74
|
+
this.slab = new Uint8Array(0); // JS source of truth
|
|
75
|
+
this.slabU32 = new Uint32Array(0);
|
|
76
|
+
// Query scratch (JS fallback path): word view + byte view over one buffer
|
|
77
|
+
this.queryU32 = new Uint32Array(Math.ceil(dimension / 4));
|
|
78
|
+
this.queryBytes = new Uint8Array(this.queryU32.buffer, 0, dimension);
|
|
79
|
+
|
|
80
|
+
// Batch scratch: callers fill batchIdx[0..n) with node indices, call
|
|
81
|
+
// batchToQuery(n)/batchToVector(t, n), and read distances from the
|
|
82
|
+
// returned batchOut. One FFI call per block instead of one per neighbor.
|
|
83
|
+
this.batchIdx = new Uint32Array(BATCH_CAP);
|
|
84
|
+
this.batchOut = new Uint32Array(BATCH_CAP);
|
|
85
|
+
this._wasmIdxU32 = null; // cached views over WASM memory (rebuilt on grow)
|
|
86
|
+
this._wasmOutU32 = null;
|
|
87
|
+
|
|
88
|
+
// WASM mirror (optional)
|
|
89
|
+
this.wasm = null;
|
|
90
|
+
this.wasmMem = null;
|
|
91
|
+
const mod = getWasmModule();
|
|
92
|
+
if (mod) {
|
|
93
|
+
try {
|
|
94
|
+
const instance = new WebAssembly.Instance(mod);
|
|
95
|
+
this.wasm = instance.exports;
|
|
96
|
+
this.wasmMem = new Uint8Array(this.wasm.memory.buffer);
|
|
97
|
+
} catch {
|
|
98
|
+
this.wasm = null;
|
|
99
|
+
}
|
|
100
|
+
}
|
|
101
|
+
}
|
|
102
|
+
|
|
103
|
+
_wasmView() {
|
|
104
|
+
if (this.wasm && this.wasmMem.buffer !== this.wasm.memory.buffer) {
|
|
105
|
+
this.wasmMem = new Uint8Array(this.wasm.memory.buffer);
|
|
106
|
+
this._wasmIdxU32 = null;
|
|
107
|
+
this._wasmOutU32 = null;
|
|
108
|
+
}
|
|
109
|
+
return this.wasmMem;
|
|
110
|
+
}
|
|
111
|
+
|
|
112
|
+
/** Grow to hold at least `count` vectors. Existing views stay valid until
|
|
113
|
+
* a grow re-points them — callers must use the views returned by set(). */
|
|
114
|
+
ensure(count) {
|
|
115
|
+
if (count <= this.capacity) return false;
|
|
116
|
+
let newCap = Math.max(1024, this.capacity * 2);
|
|
117
|
+
while (newCap < count) newCap *= 2;
|
|
118
|
+
const newSlab = new Uint8Array(newCap * this.dim);
|
|
119
|
+
newSlab.set(this.slab.subarray(0, this.count * this.dim));
|
|
120
|
+
this.slab = newSlab;
|
|
121
|
+
this.slabU32 = new Uint32Array(newSlab.buffer, 0, (newCap * this.dim) >>> 2);
|
|
122
|
+
this.capacity = newCap;
|
|
123
|
+
if (this.wasm) {
|
|
124
|
+
const needed = SLAB_BASE + newCap * this.dim;
|
|
125
|
+
const have = this.wasm.memory.buffer.byteLength;
|
|
126
|
+
if (needed > have) {
|
|
127
|
+
try {
|
|
128
|
+
this.wasm.memory.grow(Math.ceil((needed - have) / WASM_PAGE));
|
|
129
|
+
} catch {
|
|
130
|
+
// Growth refused (host memory pressure) — drop the mirror, JS
|
|
131
|
+
// fallback keeps everything exact.
|
|
132
|
+
this.wasm = null;
|
|
133
|
+
}
|
|
134
|
+
}
|
|
135
|
+
}
|
|
136
|
+
return true;
|
|
137
|
+
}
|
|
138
|
+
|
|
139
|
+
/**
|
|
140
|
+
* Write vector bytes at index (append or overwrite) into the JS slab and
|
|
141
|
+
* the WASM mirror. Returns the canonical subarray view for storage on the
|
|
142
|
+
* node record.
|
|
143
|
+
*/
|
|
144
|
+
set(idx, bytes) {
|
|
145
|
+
this.ensure(idx + 1);
|
|
146
|
+
const off = idx * this.dim;
|
|
147
|
+
const n = Math.min(bytes.length, this.dim);
|
|
148
|
+
this.slab.set(n === bytes.length ? bytes : bytes.subarray(0, n), off);
|
|
149
|
+
if (n < this.dim) this.slab.fill(0, off + n, off + this.dim);
|
|
150
|
+
if (this.wasm) {
|
|
151
|
+
this._wasmView().set(this.slab.subarray(off, off + this.dim), SLAB_BASE + off);
|
|
152
|
+
}
|
|
153
|
+
if (idx >= this.count) this.count = idx + 1;
|
|
154
|
+
return this.slab.subarray(off, off + this.dim);
|
|
155
|
+
}
|
|
156
|
+
|
|
157
|
+
/** Re-point a node's canonical view after grows (rarely needed by callers
|
|
158
|
+
* that persist views long-term across many inserts). */
|
|
159
|
+
view(idx) {
|
|
160
|
+
const off = idx * this.dim;
|
|
161
|
+
return this.slab.subarray(off, off + this.dim);
|
|
162
|
+
}
|
|
163
|
+
|
|
164
|
+
/** Stage the query for distToQuery(). Shorter inputs are zero-padded. */
|
|
165
|
+
setQuery(bytes) {
|
|
166
|
+
const n = Math.min(bytes.length, this.dim);
|
|
167
|
+
this.queryBytes.set(n === bytes.length ? bytes : bytes.subarray(0, n));
|
|
168
|
+
if (n < this.dim) this.queryBytes.fill(0, n);
|
|
169
|
+
if (this.wasm) {
|
|
170
|
+
this._wasmView().set(this.queryBytes, 0);
|
|
171
|
+
}
|
|
172
|
+
}
|
|
173
|
+
|
|
174
|
+
/** Hamming distance from the staged query to vector idx. */
|
|
175
|
+
distToQuery(idx) {
|
|
176
|
+
if (this.wasm) {
|
|
177
|
+
return this.wasm.hamming_distance(0, SLAB_BASE + idx * this.dim, this.dim);
|
|
178
|
+
}
|
|
179
|
+
if (this.tailBytes === 0) {
|
|
180
|
+
if (this.words === 16) return swar16(this.queryU32, 0, this.slabU32, idx << 4);
|
|
181
|
+
const s = this.slabU32;
|
|
182
|
+
const q = this.queryU32;
|
|
183
|
+
const o = idx * this.words;
|
|
184
|
+
let c = 0;
|
|
185
|
+
for (let w = 0; w < this.words; w++) c += popcnt32(q[w] ^ s[o + w]);
|
|
186
|
+
return c;
|
|
187
|
+
}
|
|
188
|
+
// dim not a multiple of 4: per-vector slab offsets are word-unaligned,
|
|
189
|
+
// so fall back to a byte loop.
|
|
190
|
+
const off = idx * this.dim;
|
|
191
|
+
let c = 0;
|
|
192
|
+
for (let i = 0; i < this.dim; i++) c += popcnt32(this.queryBytes[i] ^ this.slab[off + i]);
|
|
193
|
+
return c;
|
|
194
|
+
}
|
|
195
|
+
|
|
196
|
+
/** Hamming distance between two resident vectors. */
|
|
197
|
+
dist(aIdx, bIdx) {
|
|
198
|
+
if (this.wasm) {
|
|
199
|
+
return this.wasm.hamming_distance(
|
|
200
|
+
SLAB_BASE + aIdx * this.dim,
|
|
201
|
+
SLAB_BASE + bIdx * this.dim,
|
|
202
|
+
this.dim
|
|
203
|
+
);
|
|
204
|
+
}
|
|
205
|
+
if (this.tailBytes === 0) {
|
|
206
|
+
if (this.words === 16) return swar16(this.slabU32, aIdx << 4, this.slabU32, bIdx << 4);
|
|
207
|
+
const s = this.slabU32;
|
|
208
|
+
const ao = aIdx * this.words;
|
|
209
|
+
const bo = bIdx * this.words;
|
|
210
|
+
let c = 0;
|
|
211
|
+
for (let w = 0; w < this.words; w++) c += popcnt32(s[ao + w] ^ s[bo + w]);
|
|
212
|
+
return c;
|
|
213
|
+
}
|
|
214
|
+
const aOff = aIdx * this.dim;
|
|
215
|
+
const bOff = bIdx * this.dim;
|
|
216
|
+
let c = 0;
|
|
217
|
+
for (let i = 0; i < this.dim; i++) c += popcnt32(this.slab[aOff + i] ^ this.slab[bOff + i]);
|
|
218
|
+
return c;
|
|
219
|
+
}
|
|
220
|
+
|
|
221
|
+
/**
|
|
222
|
+
* Batch: distances from the staged query to batchIdx[0..n). Results land
|
|
223
|
+
* in (and are returned as) batchOut. One FFI call per block on the WASM
|
|
224
|
+
* path; order-preserving and bit-exact on every path.
|
|
225
|
+
*/
|
|
226
|
+
batchToQuery(n) {
|
|
227
|
+
return this._batch(0, n);
|
|
228
|
+
}
|
|
229
|
+
|
|
230
|
+
/** Batch: distances from resident vector targetIdx to batchIdx[0..n). */
|
|
231
|
+
batchToVector(targetIdx, n) {
|
|
232
|
+
if (this.wasm) return this._batch(SLAB_BASE + targetIdx * this.dim, n);
|
|
233
|
+
const idx = this.batchIdx;
|
|
234
|
+
const out = this.batchOut;
|
|
235
|
+
for (let i = 0; i < n; i++) out[i] = this.dist(targetIdx, idx[i]);
|
|
236
|
+
return out;
|
|
237
|
+
}
|
|
238
|
+
|
|
239
|
+
_batch(qPtr, n) {
|
|
240
|
+
if (this.wasm && this.wasm.hamming_batch && n <= BATCH_CAP) {
|
|
241
|
+
this._wasmView(); // refresh cached views if memory grew
|
|
242
|
+
if (this._wasmIdxU32 === null) {
|
|
243
|
+
this._wasmIdxU32 = new Uint32Array(this.wasm.memory.buffer, BATCH_IDX_PTR, BATCH_CAP);
|
|
244
|
+
this._wasmOutU32 = new Uint32Array(this.wasm.memory.buffer, BATCH_OUT_PTR, BATCH_CAP);
|
|
245
|
+
}
|
|
246
|
+
const wIdx = this._wasmIdxU32;
|
|
247
|
+
const idx = this.batchIdx;
|
|
248
|
+
for (let i = 0; i < n; i++) wIdx[i] = idx[i];
|
|
249
|
+
this.wasm.hamming_batch(qPtr, SLAB_BASE, this.dim, BATCH_IDX_PTR, n, BATCH_OUT_PTR);
|
|
250
|
+
const wOut = this._wasmOutU32;
|
|
251
|
+
const out = this.batchOut;
|
|
252
|
+
for (let i = 0; i < n; i++) out[i] = wOut[i];
|
|
253
|
+
return out;
|
|
254
|
+
}
|
|
255
|
+
const idx = this.batchIdx;
|
|
256
|
+
const out = this.batchOut;
|
|
257
|
+
for (let i = 0; i < n; i++) out[i] = this.distToQuery(idx[i]);
|
|
258
|
+
return out;
|
|
259
|
+
}
|
|
260
|
+
|
|
261
|
+
reset() {
|
|
262
|
+
this.count = 0;
|
|
263
|
+
// Keep allocation; contents are overwritten by set() before reads.
|
|
264
|
+
}
|
|
265
|
+
}
|
|
266
|
+
|
|
267
|
+
/**
|
|
268
|
+
* Fully-unrolled 16-word (64-byte / 512-bit) SWAR Hamming with deferred
|
|
269
|
+
* reduction in ≤7-word groups — 7 words × 8/byte-lane × 4 lanes = 224 < 256,
|
|
270
|
+
* so the final multiply's top byte can't overflow. Bit-exact.
|
|
271
|
+
*/
|
|
272
|
+
function swar16(A, ao, B, bo) {
|
|
273
|
+
let x, a = 0, b = 0, c = 0;
|
|
274
|
+
x = A[ao] ^ B[bo]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); a += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
275
|
+
x = A[ao + 1] ^ B[bo + 1]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); a += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
276
|
+
x = A[ao + 2] ^ B[bo + 2]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); a += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
277
|
+
x = A[ao + 3] ^ B[bo + 3]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); a += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
278
|
+
x = A[ao + 4] ^ B[bo + 4]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); a += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
279
|
+
x = A[ao + 5] ^ B[bo + 5]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); a += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
280
|
+
x = A[ao + 6] ^ B[bo + 6]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); a += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
281
|
+
x = A[ao + 7] ^ B[bo + 7]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); b += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
282
|
+
x = A[ao + 8] ^ B[bo + 8]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); b += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
283
|
+
x = A[ao + 9] ^ B[bo + 9]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); b += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
284
|
+
x = A[ao + 10] ^ B[bo + 10]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); b += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
285
|
+
x = A[ao + 11] ^ B[bo + 11]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); b += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
286
|
+
x = A[ao + 12] ^ B[bo + 12]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); b += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
287
|
+
x = A[ao + 13] ^ B[bo + 13]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); b += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
288
|
+
x = A[ao + 14] ^ B[bo + 14]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); c += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
289
|
+
x = A[ao + 15] ^ B[bo + 15]; x -= (x >>> 1) & 0x55555555; x = (x & 0x33333333) + ((x >>> 2) & 0x33333333); c += (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
290
|
+
return (Math.imul(a, 0x01010101) >>> 24)
|
|
291
|
+
+ (Math.imul(b, 0x01010101) >>> 24)
|
|
292
|
+
+ (Math.imul(c, 0x01010101) >>> 24);
|
|
293
|
+
}
|
|
Binary file
|
|
@@ -403,9 +403,35 @@ export function wasmMaxSimDequant(queryFlat, docInt8, numQ, numD, dim, min, scal
|
|
|
403
403
|
return maxsimExports.maxsim_dequant(qPtr, dPtr, numQ, numD, dim, min, scale);
|
|
404
404
|
}
|
|
405
405
|
|
|
406
|
+
// Query-copy-once session for per-candidate WASM MaxSim loops.
|
|
407
|
+
//
|
|
408
|
+
// The query lives at the fixed offset DATA_OFFSET, so within one query the
|
|
409
|
+
// caller prepares it once and each per-candidate call skips the Q×dim×4-byte
|
|
410
|
+
// memcpy. The query Float32Array's identity is NOT a safe cache key — callers
|
|
411
|
+
// pool and mutate that buffer across queries — hence an explicit session id
|
|
412
|
+
// that the caller must re-obtain per query.
|
|
413
|
+
let _wasmQuerySession = 0;
|
|
414
|
+
let _wasmQueryPrepared = null; // { session, qBytes, memBuffer }
|
|
415
|
+
|
|
416
|
+
export function wasmMaxSimPrepareQuery(queryFlat, numQ, dim) {
|
|
417
|
+
if (!maxsimExports) return 0;
|
|
418
|
+
const qBytes = numQ * dim * 4;
|
|
419
|
+
if (!maxsimMem || maxsimMem.buffer !== maxsimExports.memory.buffer) {
|
|
420
|
+
maxsimMem = new Uint8Array(maxsimExports.memory.buffer);
|
|
421
|
+
}
|
|
422
|
+
if (qBytes + 1024 > maxsimMem.length) return 0;
|
|
423
|
+
maxsimMem.set(new Uint8Array(queryFlat.buffer, queryFlat.byteOffset, qBytes), DATA_OFFSET);
|
|
424
|
+
_wasmQuerySession += 1;
|
|
425
|
+
_wasmQueryPrepared = { session: _wasmQuerySession, qBytes, memBuffer: maxsimExports.memory.buffer };
|
|
426
|
+
return _wasmQuerySession;
|
|
427
|
+
}
|
|
428
|
+
|
|
406
429
|
// Shared layout: copy query + doc + per-token params into WASM memory, return pointers.
|
|
407
430
|
// Returns null if data doesn't fit in WASM memory.
|
|
408
|
-
|
|
431
|
+
// A live `querySession` from wasmMaxSimPrepareQuery (matching qBytes and
|
|
432
|
+
// memory buffer) means the query bytes already sit at DATA_OFFSET, so the
|
|
433
|
+
// query memcpy is skipped.
|
|
434
|
+
function _wasmPerTokenLayout(queryFlat, docData, minArray, scaleArray, tokenNorms, numQ, numD, dim, dBytes, querySession) {
|
|
409
435
|
const qBytes = numQ * dim * 4;
|
|
410
436
|
const paramBytes = numD * 4 * 3;
|
|
411
437
|
if (!maxsimMem || maxsimMem.buffer !== maxsimExports.memory.buffer) {
|
|
@@ -419,7 +445,14 @@ function _wasmPerTokenLayout(queryFlat, docData, minArray, scaleArray, tokenNorm
|
|
|
419
445
|
const scalePtr = minPtr + numD * 4;
|
|
420
446
|
const normPtr = scalePtr + numD * 4;
|
|
421
447
|
|
|
422
|
-
|
|
448
|
+
const queryAlreadyPlaced = querySession
|
|
449
|
+
&& _wasmQueryPrepared
|
|
450
|
+
&& _wasmQueryPrepared.session === querySession
|
|
451
|
+
&& _wasmQueryPrepared.qBytes === qBytes
|
|
452
|
+
&& _wasmQueryPrepared.memBuffer === maxsimExports.memory.buffer;
|
|
453
|
+
if (!queryAlreadyPlaced) {
|
|
454
|
+
maxsimMem.set(new Uint8Array(queryFlat.buffer, queryFlat.byteOffset, qBytes), qPtr);
|
|
455
|
+
}
|
|
423
456
|
maxsimMem.set(new Uint8Array(docData.buffer, docData.byteOffset, dBytes), dPtr);
|
|
424
457
|
maxsimMem.set(new Uint8Array(minArray.buffer, minArray.byteOffset, numD * 4), minPtr);
|
|
425
458
|
maxsimMem.set(new Uint8Array(scaleArray.buffer, scaleArray.byteOffset, numD * 4), scalePtr);
|
|
@@ -428,16 +461,16 @@ function _wasmPerTokenLayout(queryFlat, docData, minArray, scaleArray, tokenNorm
|
|
|
428
461
|
return { qPtr, dPtr, minPtr, scalePtr, normPtr };
|
|
429
462
|
}
|
|
430
463
|
|
|
431
|
-
export function wasmMaxSimDequantPerToken(queryFlat, docInt8, minArray, scaleArray, tokenNorms, numQ, numD, dim) {
|
|
464
|
+
export function wasmMaxSimDequantPerToken(queryFlat, docInt8, minArray, scaleArray, tokenNorms, numQ, numD, dim, querySession) {
|
|
432
465
|
if (!maxsimExports?.maxsim_dequant_pertoken) return null;
|
|
433
|
-
const ptrs = _wasmPerTokenLayout(queryFlat, docInt8, minArray, scaleArray, tokenNorms, numQ, numD, dim, numD * dim);
|
|
466
|
+
const ptrs = _wasmPerTokenLayout(queryFlat, docInt8, minArray, scaleArray, tokenNorms, numQ, numD, dim, numD * dim, querySession);
|
|
434
467
|
if (!ptrs) return null;
|
|
435
468
|
return maxsimExports.maxsim_dequant_pertoken(ptrs.qPtr, ptrs.dPtr, ptrs.minPtr, ptrs.scalePtr, ptrs.normPtr, numQ, numD, dim);
|
|
436
469
|
}
|
|
437
470
|
|
|
438
|
-
export function wasmMaxSimDequant4Bit(queryFlat, docPacked, minArray, scaleArray, tokenNorms, numQ, numD, dim) {
|
|
471
|
+
export function wasmMaxSimDequant4Bit(queryFlat, docPacked, minArray, scaleArray, tokenNorms, numQ, numD, dim, querySession) {
|
|
439
472
|
if (!maxsimExports?.maxsim_dequant_4bit) return null;
|
|
440
|
-
const ptrs = _wasmPerTokenLayout(queryFlat, docPacked, minArray, scaleArray, tokenNorms, numQ, numD, dim, numD * Math.ceil(dim / 2));
|
|
473
|
+
const ptrs = _wasmPerTokenLayout(queryFlat, docPacked, minArray, scaleArray, tokenNorms, numQ, numD, dim, numD * Math.ceil(dim / 2), querySession);
|
|
441
474
|
if (!ptrs) return null;
|
|
442
475
|
return maxsimExports.maxsim_dequant_4bit(ptrs.qPtr, ptrs.dPtr, ptrs.minPtr, ptrs.scalePtr, ptrs.normPtr, numQ, numD, dim);
|
|
443
476
|
}
|
|
Binary file
|
|
@@ -109,6 +109,40 @@ export function setBit(bitmap, index) {
|
|
|
109
109
|
bitmap.payload[byte] |= mask;
|
|
110
110
|
}
|
|
111
111
|
|
|
112
|
+
/**
|
|
113
|
+
* Count set bits with index < limitBits (clamped to capacity). Matches an
|
|
114
|
+
* isSet() loop over [0, limitBits) exactly — used for O(words) live-vector
|
|
115
|
+
* counting instead of an O(n) per-index probe.
|
|
116
|
+
*/
|
|
117
|
+
export function popcountRange(bitmap, limitBits) {
|
|
118
|
+
if (!bitmap || limitBits <= 0) return 0;
|
|
119
|
+
const bits = Math.min(limitBits, bitmap.capacity);
|
|
120
|
+
const buf = bitmap.payload;
|
|
121
|
+
const fullBytes = bits >>> 3;
|
|
122
|
+
let count = 0;
|
|
123
|
+
let i = 0;
|
|
124
|
+
// SWAR popcount, 4 bytes at a time
|
|
125
|
+
for (; i + 4 <= fullBytes; i += 4) {
|
|
126
|
+
let x = buf.readUInt32LE(i);
|
|
127
|
+
x -= (x >>> 1) & 0x55555555;
|
|
128
|
+
x = (x & 0x33333333) + ((x >>> 2) & 0x33333333);
|
|
129
|
+
x = (x + (x >>> 4)) & 0x0f0f0f0f;
|
|
130
|
+
count += Math.imul(x, 0x01010101) >>> 24;
|
|
131
|
+
}
|
|
132
|
+
for (; i < fullBytes; i++) {
|
|
133
|
+
let b = buf[i];
|
|
134
|
+
while (b !== 0) { b &= b - 1; count += 1; }
|
|
135
|
+
}
|
|
136
|
+
const tailBits = bits & 7;
|
|
137
|
+
if (tailBits > 0) {
|
|
138
|
+
// Bits are LSB-first within a byte (see byteAndMask), so the low mask
|
|
139
|
+
// selects exactly the indices below `bits`.
|
|
140
|
+
let b = buf[fullBytes] & ((1 << tailBits) - 1);
|
|
141
|
+
while (b !== 0) { b &= b - 1; count += 1; }
|
|
142
|
+
}
|
|
143
|
+
return count;
|
|
144
|
+
}
|
|
145
|
+
|
|
112
146
|
export function popcount(bitmap) {
|
|
113
147
|
if (!bitmap) return 0;
|
|
114
148
|
let count = 0;
|