difflib-fast 0.4.0__tar.gz → 0.5.0__tar.gz
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.
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/.github/workflows/CI.yml +3 -4
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/Cargo.lock +1 -1
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/Cargo.toml +1 -1
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/PKG-INFO +11 -14
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/README.md +10 -13
- difflib_fast-0.5.0/examples/simjoin_double.rs +110 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/gestalt.rs +1 -1
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/simjoin.rs +590 -105
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/simjoin_gpu.rs +8 -1
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/.gitignore +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/benchmarks.md +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/examples/simjoin_bench.rs +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/examples/simjoin_gpu_bench.rs +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/examples/simjoin_pypi.rs +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/pyproject.toml +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/python/difflib_fast/__init__.py +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/python/difflib_fast/__init__.pyi +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/python/difflib_fast/py.typed +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/bin/bench.rs +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/gpu.rs +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/lib.rs +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/rationer.rs +0 -0
- {difflib_fast-0.4.0 → difflib_fast-0.5.0}/tests/approx_ro.rs +0 -0
|
@@ -1,8 +1,7 @@
|
|
|
1
1
|
# Build + test, then on a version tag (`vX.Y.Z`) attach wheels + sdist to a GitHub Release.
|
|
2
|
-
# crates.io is published manually (`cargo publish`)
|
|
3
|
-
#
|
|
4
|
-
#
|
|
5
|
-
# No secrets needed — the GitHub Release uses the built-in GITHUB_TOKEN.
|
|
2
|
+
# crates.io is published manually (`cargo publish`), and so is PyPI: download the release's
|
|
3
|
+
# wheels + sdist and `uv publish` them. Nothing here needs a secret — the GitHub Release uses the
|
|
4
|
+
# built-in GITHUB_TOKEN.
|
|
6
5
|
#
|
|
7
6
|
# abi3 (pyo3 abi3-py39) → one wheel per platform/arch works on CPython 3.9+, so no per-Python matrix.
|
|
8
7
|
name: CI
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[package]
|
|
2
2
|
name = "difflib-fast"
|
|
3
|
-
version = "0.
|
|
3
|
+
version = "0.5.0"
|
|
4
4
|
description = "Fast, byte-for-byte exact difflib Ratcliff–Obershelp (gestalt) similarity ratio + single-linkage clustering (suffix automaton), plus an exact all-pairs weighted-cosine similarity join (L2AP, CPU+GPU)."
|
|
5
5
|
keywords = ["difflib", "similarity", "ratcliff-obershelp", "suffix-automaton", "fuzzy"]
|
|
6
6
|
categories = ["algorithms", "text-processing"]
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: difflib-fast
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.5.0
|
|
4
4
|
Classifier: Programming Language :: Rust
|
|
5
5
|
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
6
6
|
Classifier: Topic :: Text Processing :: General
|
|
@@ -58,7 +58,7 @@ assert_eq!(ratio("the quick brown fox", "the quick brown dog"), 0.89473684210526
|
|
|
58
58
|
```
|
|
59
59
|
|
|
60
60
|
```toml
|
|
61
|
-
difflib-fast = "0.
|
|
61
|
+
difflib-fast = "0.5"
|
|
62
62
|
```
|
|
63
63
|
|
|
64
64
|
---
|
|
@@ -269,27 +269,24 @@ Swap `difflib.SequenceMatcher(...).ratio()` for `difflib_fast.ratio(...)` and th
|
|
|
269
269
|
back — hundreds of times faster per call, and **thousands of times** faster when you score a whole
|
|
270
270
|
corpus. That's the entire change: `import difflib_fast`.
|
|
271
271
|
|
|
272
|
-
**No PyPI** — `pip install difflib-fast` won't work. Two ways to install:
|
|
273
|
-
|
|
274
|
-
**1. Prebuilt wheel — no Rust toolchain needed.** The wheels are `cp39-abi3`, so one wheel per platform
|
|
275
|
-
works on **every CPython ≥ 3.9, including 3.14**. GitHub Releases isn't a package index, so pip can't
|
|
276
|
-
pick the wheel for you — grab the one for your platform from the
|
|
277
|
-
[Releases page](https://github.com/prostomarkeloff/difflib-fast/releases/latest) and install it by URL:
|
|
278
|
-
|
|
279
272
|
```bash
|
|
280
|
-
|
|
281
|
-
pip install https://github.com/prostomarkeloff/difflib-fast/releases/download/v0.3.5/difflib_fast-0.3.5-cp39-abi3-macosx_11_0_arm64.whl
|
|
273
|
+
pip install difflib-fast
|
|
282
274
|
```
|
|
283
275
|
|
|
284
|
-
|
|
276
|
+
The wheels on [PyPI](https://pypi.org/project/difflib-fast/) are `cp39-abi3`, so one wheel per
|
|
277
|
+
platform works on **every CPython ≥ 3.9, including 3.14** — no Rust toolchain needed:
|
|
278
|
+
|
|
279
|
+
| platform | wheel |
|
|
285
280
|
|---|---|
|
|
286
|
-
| macOS Apple Silicon | `…-macosx_11_0_arm64.whl` |
|
|
281
|
+
| macOS Apple Silicon | `…-macosx_11_0_arm64.whl` (ships the Metal GPU path) |
|
|
287
282
|
| macOS Intel | `…-macosx_10_12_x86_64.whl` |
|
|
288
283
|
| Linux x86_64 | `…-manylinux_2_17_x86_64.manylinux2014_x86_64.whl` |
|
|
289
284
|
| Linux aarch64 | `…-manylinux_2_17_aarch64.manylinux2014_aarch64.whl` |
|
|
290
285
|
| Windows x64 | `…-win_amd64.whl` |
|
|
291
286
|
|
|
292
|
-
|
|
287
|
+
The same wheels are attached to every
|
|
288
|
+
[GitHub release](https://github.com/prostomarkeloff/difflib-fast/releases/latest). Any other
|
|
289
|
+
platform builds from source — pip drives maturin, so it only needs a Rust toolchain:
|
|
293
290
|
|
|
294
291
|
```bash
|
|
295
292
|
pip install git+https://github.com/prostomarkeloff/difflib-fast
|
|
@@ -42,7 +42,7 @@ assert_eq!(ratio("the quick brown fox", "the quick brown dog"), 0.89473684210526
|
|
|
42
42
|
```
|
|
43
43
|
|
|
44
44
|
```toml
|
|
45
|
-
difflib-fast = "0.
|
|
45
|
+
difflib-fast = "0.5"
|
|
46
46
|
```
|
|
47
47
|
|
|
48
48
|
---
|
|
@@ -253,27 +253,24 @@ Swap `difflib.SequenceMatcher(...).ratio()` for `difflib_fast.ratio(...)` and th
|
|
|
253
253
|
back — hundreds of times faster per call, and **thousands of times** faster when you score a whole
|
|
254
254
|
corpus. That's the entire change: `import difflib_fast`.
|
|
255
255
|
|
|
256
|
-
**No PyPI** — `pip install difflib-fast` won't work. Two ways to install:
|
|
257
|
-
|
|
258
|
-
**1. Prebuilt wheel — no Rust toolchain needed.** The wheels are `cp39-abi3`, so one wheel per platform
|
|
259
|
-
works on **every CPython ≥ 3.9, including 3.14**. GitHub Releases isn't a package index, so pip can't
|
|
260
|
-
pick the wheel for you — grab the one for your platform from the
|
|
261
|
-
[Releases page](https://github.com/prostomarkeloff/difflib-fast/releases/latest) and install it by URL:
|
|
262
|
-
|
|
263
256
|
```bash
|
|
264
|
-
|
|
265
|
-
pip install https://github.com/prostomarkeloff/difflib-fast/releases/download/v0.3.5/difflib_fast-0.3.5-cp39-abi3-macosx_11_0_arm64.whl
|
|
257
|
+
pip install difflib-fast
|
|
266
258
|
```
|
|
267
259
|
|
|
268
|
-
|
|
260
|
+
The wheels on [PyPI](https://pypi.org/project/difflib-fast/) are `cp39-abi3`, so one wheel per
|
|
261
|
+
platform works on **every CPython ≥ 3.9, including 3.14** — no Rust toolchain needed:
|
|
262
|
+
|
|
263
|
+
| platform | wheel |
|
|
269
264
|
|---|---|
|
|
270
|
-
| macOS Apple Silicon | `…-macosx_11_0_arm64.whl` |
|
|
265
|
+
| macOS Apple Silicon | `…-macosx_11_0_arm64.whl` (ships the Metal GPU path) |
|
|
271
266
|
| macOS Intel | `…-macosx_10_12_x86_64.whl` |
|
|
272
267
|
| Linux x86_64 | `…-manylinux_2_17_x86_64.manylinux2014_x86_64.whl` |
|
|
273
268
|
| Linux aarch64 | `…-manylinux_2_17_aarch64.manylinux2014_aarch64.whl` |
|
|
274
269
|
| Windows x64 | `…-win_amd64.whl` |
|
|
275
270
|
|
|
276
|
-
|
|
271
|
+
The same wheels are attached to every
|
|
272
|
+
[GitHub release](https://github.com/prostomarkeloff/difflib-fast/releases/latest). Any other
|
|
273
|
+
platform builds from source — pip drives maturin, so it only needs a Rust toolchain:
|
|
277
274
|
|
|
278
275
|
```bash
|
|
279
276
|
pip install git+https://github.com/prostomarkeloff/difflib-fast
|
|
@@ -0,0 +1,110 @@
|
|
|
1
|
+
//! Bench for the fused two-corpus join [`difflib_fast::simjoin::cosine_join_double`] — the typical
|
|
2
|
+
//! "two views of the same objects" shape (e.g. line-vectors and part-vectors of the same functions:
|
|
3
|
+
//! two corpora with the same row count). Times one fused join against two separate
|
|
4
|
+
//! [`cosine_join`]s and checks each half matches the separate result.
|
|
5
|
+
//!
|
|
6
|
+
//! `cargo run --release --example simjoin_double -- <a.bin> [b.bin] [t] [reps]`
|
|
7
|
+
//! defaults: b.bin = a.bin, t=0.8, reps=3. `SJ_NSUB=<n>` truncates both corpora; thread count via
|
|
8
|
+
//! `RAYON_NUM_THREADS`.
|
|
9
|
+
|
|
10
|
+
#![allow(
|
|
11
|
+
clippy::cast_possible_truncation,
|
|
12
|
+
clippy::cast_precision_loss,
|
|
13
|
+
clippy::many_single_char_names,
|
|
14
|
+
clippy::doc_markdown
|
|
15
|
+
)]
|
|
16
|
+
|
|
17
|
+
use std::time::Instant;
|
|
18
|
+
|
|
19
|
+
use difflib_fast::simjoin::{cosine_join, cosine_join_double, Corpus};
|
|
20
|
+
|
|
21
|
+
fn le_u32(b: &[u8], p: usize) -> u32 {
|
|
22
|
+
u32::from_le_bytes(b[p..p + 4].try_into().unwrap())
|
|
23
|
+
}
|
|
24
|
+
|
|
25
|
+
fn arg<T: std::str::FromStr>(i: usize, def: T) -> T {
|
|
26
|
+
std::env::args().nth(i).and_then(|s| s.parse().ok()).unwrap_or(def)
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
/// Parse the `SIMJOIN1` binary into raw `(dim, weight)` rows (same format as `examples/simjoin_pypi`).
|
|
30
|
+
fn load(path: &str) -> Vec<Vec<(u32, f64)>> {
|
|
31
|
+
let b = std::fs::read(path).expect("read corpus");
|
|
32
|
+
assert!(b.len() >= 16, "corpus too small");
|
|
33
|
+
let magic = u64::from_le_bytes(b[0..8].try_into().unwrap());
|
|
34
|
+
assert_eq!(magic, 0x5349_4D4A_4F49_4E31, "bad magic");
|
|
35
|
+
let n = le_u32(&b, 8) as usize;
|
|
36
|
+
let mut p = 16usize;
|
|
37
|
+
let mut rows = Vec::with_capacity(n);
|
|
38
|
+
for _ in 0..n {
|
|
39
|
+
let nnz = le_u32(&b, p) as usize;
|
|
40
|
+
p += 4;
|
|
41
|
+
let mut row = Vec::with_capacity(nnz);
|
|
42
|
+
for _ in 0..nnz {
|
|
43
|
+
let d = le_u32(&b, p);
|
|
44
|
+
p += 4;
|
|
45
|
+
let w = f64::from(f32::from_le_bytes(b[p..p + 4].try_into().unwrap()));
|
|
46
|
+
p += 4;
|
|
47
|
+
row.push((d, w));
|
|
48
|
+
}
|
|
49
|
+
rows.push(row);
|
|
50
|
+
}
|
|
51
|
+
rows
|
|
52
|
+
}
|
|
53
|
+
|
|
54
|
+
/// Canonical order for the parity check (join output order is arbitrary).
|
|
55
|
+
fn sort_pairs(mut v: Vec<(usize, usize, f64)>) -> Vec<(usize, usize, u64)> {
|
|
56
|
+
v.sort_by_key(|a| (a.0, a.1));
|
|
57
|
+
v.into_iter().map(|(a, b, s)| (a, b, s.to_bits())).collect()
|
|
58
|
+
}
|
|
59
|
+
|
|
60
|
+
fn main() {
|
|
61
|
+
let path_a: String = arg(1, "perf-local/pypi-type3.simjoin.bin".to_string());
|
|
62
|
+
let path_b: String = arg(2, path_a.clone());
|
|
63
|
+
let t: f64 = arg(3, 0.8);
|
|
64
|
+
let reps: usize = arg(4, 3);
|
|
65
|
+
|
|
66
|
+
let mut rows_a = load(&path_a);
|
|
67
|
+
let mut rows_b = load(&path_b);
|
|
68
|
+
if let Ok(nsub) = std::env::var("SJ_NSUB") {
|
|
69
|
+
let nsub: usize = nsub.parse().expect("SJ_NSUB");
|
|
70
|
+
rows_a.truncate(nsub);
|
|
71
|
+
rows_b.truncate(nsub);
|
|
72
|
+
}
|
|
73
|
+
assert_eq!(rows_a.len(), rows_b.len(), "corpora must share the row count");
|
|
74
|
+
let n = rows_a.len();
|
|
75
|
+
|
|
76
|
+
let corpus_a = Corpus::from_rows(&rows_a);
|
|
77
|
+
let corpus_b = Corpus::from_rows(&rows_b);
|
|
78
|
+
|
|
79
|
+
let mut fused: Vec<f64> = Vec::with_capacity(reps);
|
|
80
|
+
let mut separate: Vec<f64> = Vec::with_capacity(reps);
|
|
81
|
+
let (mut pa, mut pb) = (0usize, 0usize);
|
|
82
|
+
for r in 0..reps {
|
|
83
|
+
let (ga, gb) = {
|
|
84
|
+
let t0 = Instant::now();
|
|
85
|
+
let r = cosine_join_double(&corpus_a, &corpus_b, t);
|
|
86
|
+
fused.push(t0.elapsed().as_secs_f64() * 1000.0);
|
|
87
|
+
r
|
|
88
|
+
};
|
|
89
|
+
let (sa, sb) = {
|
|
90
|
+
let t0 = Instant::now();
|
|
91
|
+
let r = (cosine_join(&corpus_a, t), cosine_join(&corpus_b, t));
|
|
92
|
+
separate.push(t0.elapsed().as_secs_f64() * 1000.0);
|
|
93
|
+
r
|
|
94
|
+
};
|
|
95
|
+
if r == 0 {
|
|
96
|
+
pa = ga.len();
|
|
97
|
+
pb = gb.len();
|
|
98
|
+
// The fused halves must be bit-identical to the separate joins (real-data smoke).
|
|
99
|
+
assert_eq!(sort_pairs(ga), sort_pairs(sa), "half a differs from the separate join");
|
|
100
|
+
assert_eq!(sort_pairs(gb), sort_pairs(sb), "half b differs from the separate join");
|
|
101
|
+
}
|
|
102
|
+
}
|
|
103
|
+
fused.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
|
104
|
+
separate.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
|
105
|
+
let (f, s) = (fused[fused.len() / 2], separate[separate.len() / 2]);
|
|
106
|
+
eprintln!(
|
|
107
|
+
"double n={n} t={t} | pairs=({pa},{pb}) | fused={f:.1}ms separate={s:.1}ms speedup={:.2}x",
|
|
108
|
+
s / f
|
|
109
|
+
);
|
|
110
|
+
}
|
|
@@ -894,7 +894,7 @@ impl Builder {
|
|
|
894
894
|
fn sort_big(&mut self, out: &mut Sam, nstates: usize, blen: usize) {
|
|
895
895
|
// sorted copies for the big endpos sets (a handful of states: short, frequent substrings);
|
|
896
896
|
// positions are below `blen`, so a radix sort needs one pass per byte of that
|
|
897
|
-
let passes = (
|
|
897
|
+
let passes = (blen.bit_width() as usize).div_ceil(8).max(1);
|
|
898
898
|
let scratch = &mut self.child_arr;
|
|
899
899
|
for s in 0..nstates {
|
|
900
900
|
let cnt = out.chain_slot[s][7] as usize;
|
|
@@ -43,6 +43,9 @@ use rayon::prelude::*;
|
|
|
43
43
|
|
|
44
44
|
use crate::Concurrency;
|
|
45
45
|
|
|
46
|
+
/// One join's result: `(row_a, row_b, cosine)` triples with `row_a < row_b`, arbitrary order.
|
|
47
|
+
pub type JoinPairs = Vec<(usize, usize, f64)>;
|
|
48
|
+
|
|
46
49
|
/// A corpus of L2-normalised sparse non-negative vectors in CSR form, with dimensions relabelled to
|
|
47
50
|
/// a global rank (decreasing max weight). Built once; joinable at any threshold.
|
|
48
51
|
pub struct Corpus {
|
|
@@ -58,6 +61,41 @@ pub struct Corpus {
|
|
|
58
61
|
maxw: Vec<f64>,
|
|
59
62
|
}
|
|
60
63
|
|
|
64
|
+
/// Merge duplicate dims of one row (weights summed in input order within each dim group, dims in
|
|
65
|
+
/// ascending dim order — deterministic, unlike the per-call-random `HashMap` iteration order of the
|
|
66
|
+
/// previous implementation) and L2-normalise. A row whose norm is 0 (all weights 0) becomes empty.
|
|
67
|
+
#[must_use]
|
|
68
|
+
fn norm_row(r: &[(u32, f64)]) -> Vec<(u32, f64)> {
|
|
69
|
+
let mut v = r.to_vec();
|
|
70
|
+
v.sort_by_key(|&(d, _)| d); // stable: preserves input order within each duplicate group
|
|
71
|
+
let mut out: Vec<(u32, f64)> = Vec::with_capacity(v.len());
|
|
72
|
+
let mut sq = 0.0;
|
|
73
|
+
let mut iter = v.into_iter();
|
|
74
|
+
if let Some((mut d, mut w)) = iter.next() {
|
|
75
|
+
for (d2, w2) in iter {
|
|
76
|
+
if d2 == d {
|
|
77
|
+
w += w2;
|
|
78
|
+
} else {
|
|
79
|
+
sq += w * w;
|
|
80
|
+
out.push((d, w));
|
|
81
|
+
d = d2;
|
|
82
|
+
w = w2;
|
|
83
|
+
}
|
|
84
|
+
}
|
|
85
|
+
sq += w * w;
|
|
86
|
+
out.push((d, w));
|
|
87
|
+
}
|
|
88
|
+
let norm = sq.sqrt();
|
|
89
|
+
if norm > 0.0 {
|
|
90
|
+
for (_, w) in &mut out {
|
|
91
|
+
*w /= norm;
|
|
92
|
+
}
|
|
93
|
+
out
|
|
94
|
+
} else {
|
|
95
|
+
Vec::new()
|
|
96
|
+
}
|
|
97
|
+
}
|
|
98
|
+
|
|
61
99
|
impl Corpus {
|
|
62
100
|
/// Number of vectors.
|
|
63
101
|
#[must_use]
|
|
@@ -131,61 +169,131 @@ impl Corpus {
|
|
|
131
169
|
/// Build a corpus from raw `(dim, weight)` rows. Duplicate dims within a row are summed; each row
|
|
132
170
|
/// is L2-normalised; dims are relabelled to a global rank by decreasing max weight. Weights are
|
|
133
171
|
/// expected non-negative (the prefix-filter bound requires it — IDF weights satisfy this).
|
|
172
|
+
///
|
|
173
|
+
/// The build is rayon-parallel (per-row normalise, global max-weight reduction, per-row relabel
|
|
174
|
+
/// into the CSR) and **deterministic**: rows are merged in dim order, the rank is a function of
|
|
175
|
+
/// the exact per-dim max weight (ties by dim id), and parallel reductions only combine
|
|
176
|
+
/// order-independent operations — the resulting corpus is a pure function of `rows`.
|
|
177
|
+
///
|
|
178
|
+
/// Dims that occur in exactly one row (`df = 1`) are dropped from the rank and the rows **after**
|
|
179
|
+
/// normalisation: no pair shares them, so they contribute exactly `0.0` to every pairwise dot and
|
|
180
|
+
/// their postings are never walked — the join output is bit-identical with or without them, while
|
|
181
|
+
/// the index and the rows get smaller and the pruning bounds tighter. Normalising over the full
|
|
182
|
+
/// row first is what keeps the surviving weights — and every emitted score — unchanged.
|
|
134
183
|
#[must_use]
|
|
135
184
|
pub fn from_rows(rows: &[Vec<(u32, f64)>]) -> Corpus {
|
|
185
|
+
Self::from_rows_with(rows, true)
|
|
186
|
+
}
|
|
187
|
+
|
|
188
|
+
/// Build a corpus from normalized rows with optional **hapax stripping**. Hapax dims (present in
|
|
189
|
+
/// exactly one row) are dropped from the rank and the rows **after** L2 normalisation — no pair
|
|
190
|
+
/// shares them, so they contribute exactly `0.0` to every pairwise dot, their postings are never
|
|
191
|
+
/// walked (probes visit only postings of their own dims), and dropping their mass only tightens
|
|
192
|
+
/// the Cauchy–Schwarz bounds (which are filters). The join output is therefore bit-identical
|
|
193
|
+
/// with stripping on or off; the strip just shrinks the index/rows and prunes harder. Normalising
|
|
194
|
+
/// over the full row first keeps the surviving weights — and every emitted score — unchanged
|
|
195
|
+
/// (removing the mass *before* normalisation would rescale the weights and change the indexing).
|
|
196
|
+
/// `strip_hapax = false` is the un-stripped reference build, used by tests to pin the equality.
|
|
197
|
+
fn from_rows_with(rows: &[Vec<(u32, f64)>], strip_hapax: bool) -> Corpus {
|
|
136
198
|
let n = rows.len();
|
|
137
|
-
// 1. Merge duplicate dims + L2-normalise each row (kept as (orig_dim, weight)).
|
|
138
|
-
let normed: Vec<Vec<(u32, f64)>> = rows
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
199
|
+
// 1. Merge duplicate dims + L2-normalise each row (kept as (orig_dim, weight)), in parallel.
|
|
200
|
+
let normed: Vec<Vec<(u32, f64)>> = rows.par_iter().map(|r| norm_row(r)).collect();
|
|
201
|
+
// 2. Document frequency per original dim (over the normalised rows) — a parallel
|
|
202
|
+
// fold+reduce of exact counts, so the chunking cannot change the result.
|
|
203
|
+
let df: HashMap<u32, u32> = normed
|
|
204
|
+
.par_iter()
|
|
205
|
+
.fold(HashMap::new, |mut m, v| {
|
|
206
|
+
for &(d, _) in v {
|
|
207
|
+
*m.entry(d).or_insert(0) += 1;
|
|
144
208
|
}
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
209
|
+
m
|
|
210
|
+
})
|
|
211
|
+
.reduce(HashMap::new, |mut a, mut b| {
|
|
212
|
+
if a.len() < b.len() {
|
|
213
|
+
std::mem::swap(&mut a, &mut b);
|
|
214
|
+
}
|
|
215
|
+
for (d, c) in b {
|
|
216
|
+
*a.entry(d).or_insert(0) += c;
|
|
150
217
|
}
|
|
218
|
+
a
|
|
219
|
+
});
|
|
220
|
+
// A dim earns a rank only if some pair could share it. `strip_hapax = false` keeps every
|
|
221
|
+
// dim (test-only reference).
|
|
222
|
+
let shared = |d: u32| !strip_hapax || df.get(&d).is_some_and(|&c| c >= 2);
|
|
223
|
+
// 3. Max normalised weight per original dim — a parallel fold+reduce. `max` is exact and
|
|
224
|
+
// order-independent, so the chunking cannot change any value.
|
|
225
|
+
let maxw_orig: HashMap<u32, f64> = normed
|
|
226
|
+
.par_iter()
|
|
227
|
+
.fold(HashMap::new, |mut m, v| {
|
|
228
|
+
for &(d, w) in v {
|
|
229
|
+
if shared(d) {
|
|
230
|
+
let e = m.entry(d).or_insert(0.0);
|
|
231
|
+
if w > *e {
|
|
232
|
+
*e = w;
|
|
233
|
+
}
|
|
234
|
+
}
|
|
235
|
+
}
|
|
236
|
+
m
|
|
151
237
|
})
|
|
152
|
-
.
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
for v in &normed {
|
|
156
|
-
for &(d, w) in v {
|
|
157
|
-
let e = maxw_orig.entry(d).or_insert(0.0);
|
|
158
|
-
if w > *e {
|
|
159
|
-
*e = w;
|
|
238
|
+
.reduce(HashMap::new, |mut a, mut b| {
|
|
239
|
+
if a.len() < b.len() {
|
|
240
|
+
std::mem::swap(&mut a, &mut b);
|
|
160
241
|
}
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
242
|
+
for (d, w) in b {
|
|
243
|
+
let e = a.entry(d).or_insert(0.0);
|
|
244
|
+
if w > *e {
|
|
245
|
+
*e = w;
|
|
246
|
+
}
|
|
247
|
+
}
|
|
248
|
+
a
|
|
249
|
+
});
|
|
250
|
+
// 4. Rank dims by (max weight ASC, dim asc) → dense global order. Ascending so the common,
|
|
164
251
|
// low-weight dims land at the FRONT: they fill the un-indexed prefix (their tiny
|
|
165
252
|
// `w·maxw` keeps the cumulative bound under `t` for many of them), and only the rare,
|
|
166
253
|
// high-weight tail is indexed — short postings. (Reversing this indexes the common dims
|
|
167
254
|
// and their huge postings, which is correct but orders of magnitude slower.)
|
|
255
|
+
// The comparator is a total order on distinct dims, so the unstable parallel sort yields
|
|
256
|
+
// the same unique sorted sequence the serial stable sort it replaces would.
|
|
168
257
|
let mut dims_sorted: Vec<(u32, f64)> = maxw_orig.into_iter().collect();
|
|
169
|
-
dims_sorted.
|
|
258
|
+
dims_sorted.par_sort_unstable_by(|a, b| {
|
|
170
259
|
a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal).then(a.0.cmp(&b.0))
|
|
171
260
|
});
|
|
172
261
|
let rank: HashMap<u32, u32> =
|
|
173
262
|
dims_sorted.iter().enumerate().map(|(i, &(d, _))| (d, i as u32)).collect();
|
|
174
263
|
let ndims = dims_sorted.len();
|
|
175
264
|
let maxw: Vec<f64> = dims_sorted.iter().map(|&(_, w)| w).collect();
|
|
176
|
-
//
|
|
265
|
+
// 5. CSR with relabelled dims, ascending within each row. Relabel + sort each row in
|
|
266
|
+
// parallel (rank lookups are the cost), then a serial flatten into the flat CSR arrays —
|
|
267
|
+
// a pure memcpy; a parallel fill into disjoint slices would need unsafe borrows for no
|
|
268
|
+
// measurable win. Hapax dims have no rank, so the `rank.get` filter drops them.
|
|
177
269
|
let mut indptr = Vec::with_capacity(n + 1);
|
|
178
270
|
indptr.push(0);
|
|
179
|
-
let mut
|
|
180
|
-
let mut wts = Vec::new();
|
|
271
|
+
let mut off = 0usize;
|
|
181
272
|
for v in &normed {
|
|
182
|
-
let
|
|
183
|
-
|
|
273
|
+
let kept = if strip_hapax {
|
|
274
|
+
v.iter().filter(|&&(d, _)| rank.contains_key(&d)).count()
|
|
275
|
+
} else {
|
|
276
|
+
v.len()
|
|
277
|
+
};
|
|
278
|
+
off += kept;
|
|
279
|
+
indptr.push(off);
|
|
280
|
+
}
|
|
281
|
+
let relabelled: Vec<Vec<(u32, f64)>> = normed
|
|
282
|
+
.par_iter()
|
|
283
|
+
.map(|v| {
|
|
284
|
+
let mut rv: Vec<(u32, f64)> =
|
|
285
|
+
v.iter().filter_map(|&(d, w)| rank.get(&d).map(|&r| (r, w))).collect();
|
|
286
|
+
rv.sort_unstable_by_key(|&(d, _)| d);
|
|
287
|
+
rv
|
|
288
|
+
})
|
|
289
|
+
.collect();
|
|
290
|
+
let mut dims = Vec::with_capacity(off);
|
|
291
|
+
let mut wts = Vec::with_capacity(off);
|
|
292
|
+
for rv in relabelled {
|
|
184
293
|
for (d, w) in rv {
|
|
185
294
|
dims.push(d);
|
|
186
295
|
wts.push(w);
|
|
187
296
|
}
|
|
188
|
-
indptr.push(dims.len());
|
|
189
297
|
}
|
|
190
298
|
Corpus { n, ndims, indptr, dims, wts, maxw }
|
|
191
299
|
}
|
|
@@ -248,8 +356,53 @@ struct Scratch {
|
|
|
248
356
|
touched: Vec<u32>,
|
|
249
357
|
/// Probe prefix L2 norms for this probe: `xpn[k] = ‖wi[..k]‖₂`, length `nnz+1`.
|
|
250
358
|
xpn: Vec<f64>,
|
|
359
|
+
/// Sparse-mode accumulator: `(candidate, partial dot)` pairs sorted by candidate id — the
|
|
360
|
+
/// touched list and the accumulator merged (classic L2AP). Used instead of `acc`/`touched` when
|
|
361
|
+
/// the corpus's average per-probe candidate count is low; the probe falls back to the dense pair
|
|
362
|
+
/// if the list overflows [`LIST_ACC_MAX`].
|
|
363
|
+
list: Vec<(u32, f64)>,
|
|
364
|
+
}
|
|
365
|
+
|
|
366
|
+
/// Accumulator strategy for the probe phase, chosen per join from the corpus's shape (and forced
|
|
367
|
+
/// per-mode in tests — the two must agree bit-for-bit).
|
|
368
|
+
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
|
369
|
+
enum AccMode {
|
|
370
|
+
/// Dense `acc[n]` array: O(1) scatter per posting entry. Wins when probes touch many candidates.
|
|
371
|
+
Dense,
|
|
372
|
+
/// Sorted `(y, acc)` list: few candidates per probe ⇒ the list stays tiny and L1-resident, and
|
|
373
|
+
/// the n-wide dense array's footprint never competes with the postings for L2.
|
|
374
|
+
Sparse,
|
|
251
375
|
}
|
|
252
376
|
|
|
377
|
+
/// Average per-probe posting touches at/below which [`cosine_join`] picks the [`AccMode::Sparse`]
|
|
378
|
+
/// accumulator. The touch volume is estimated from the built index as `Σ len²/2` (posting `d` is
|
|
379
|
+
/// walked by about half the probes, each walk visiting `len` entries) divided by `n` — the per-probe
|
|
380
|
+
/// mean walk volume, a corpus-shape signal. Per-probe outliers are caught by [`LIST_ACC_MAX`]'s
|
|
381
|
+
/// dense fallback.
|
|
382
|
+
const LIST_ACC_MAX_TOUCHES: usize = 32;
|
|
383
|
+
|
|
384
|
+
/// Pick the accumulator mode from a built inverted index: sparse list when the average probe walks
|
|
385
|
+
/// few posting entries, dense scatter otherwise.
|
|
386
|
+
#[must_use]
|
|
387
|
+
fn acc_mode_for(index: &[Vec<(u32, f64)>], n: usize) -> AccMode {
|
|
388
|
+
let mut walks = 0u64;
|
|
389
|
+
for p in index {
|
|
390
|
+
let len = p.len() as u64;
|
|
391
|
+
walks += len * len;
|
|
392
|
+
}
|
|
393
|
+
if walks / 2 / (n.max(1) as u64) < LIST_ACC_MAX_TOUCHES as u64 {
|
|
394
|
+
AccMode::Sparse
|
|
395
|
+
} else {
|
|
396
|
+
AccMode::Dense
|
|
397
|
+
}
|
|
398
|
+
}
|
|
399
|
+
|
|
400
|
+
/// Hard cap on one probe's sparse list: beyond this the per-insert memmove dominates the O(1)
|
|
401
|
+
/// dense scatter it replaces, so [`accumulate_list`] bails out and the probe is redone densely
|
|
402
|
+
/// (re-walking its postings — rare by construction; the additions repeat in the same order, so the
|
|
403
|
+
/// values are unchanged).
|
|
404
|
+
const LIST_ACC_MAX: usize = 512;
|
|
405
|
+
|
|
253
406
|
/// Exact all-pairs cosine join via inverted index + **L2AP** prefix filtering and accumulation-time
|
|
254
407
|
/// pruning. Returns `(j, i, cos)` with `j < i` for every pair with `cos ≥ t`. Bit-identical pair set
|
|
255
408
|
/// to [`cosine_join_bruteforce`].
|
|
@@ -276,19 +429,32 @@ pub fn cosine_join(c: &Corpus, t: f64) -> Vec<(usize, usize, f64)> {
|
|
|
276
429
|
let (di, wi) = c.row(i);
|
|
277
430
|
index_suffix(c, i, (di, wi), t, &mut index, &mut cached);
|
|
278
431
|
}
|
|
279
|
-
//
|
|
280
|
-
//
|
|
281
|
-
|
|
432
|
+
// Accumulator strategy from the corpus shape: few candidates per probe on average (sparse /
|
|
433
|
+
// duplicate-heavy corpora) ⇒ the sorted list wins; many (dense rows) ⇒ the O(1) dense scatter.
|
|
434
|
+
let mode = acc_mode_for(&index, n);
|
|
435
|
+
cosine_join_mode(c, t, &index, &cached, mode)
|
|
436
|
+
}
|
|
437
|
+
|
|
438
|
+
/// The join body under an explicit [`AccMode`]: probe vectors in parallel; each worker keeps one
|
|
439
|
+
/// reusable `Scratch`. `with_min_len` batches many probes per rayon task so the per-probe work
|
|
440
|
+
/// (tiny) isn't dwarfed by task-splitting / scheduling overhead (`swtch_pri` in the profile).
|
|
441
|
+
#[must_use]
|
|
442
|
+
fn cosine_join_mode(
|
|
443
|
+
c: &Corpus,
|
|
444
|
+
t: f64,
|
|
445
|
+
index: &[Vec<(u32, f64)>],
|
|
446
|
+
cached: &Cached,
|
|
447
|
+
mode: AccMode,
|
|
448
|
+
) -> Vec<(usize, usize, f64)> {
|
|
449
|
+
let n = c.n;
|
|
282
450
|
(0..n)
|
|
283
451
|
.into_par_iter()
|
|
284
452
|
.with_min_len(256)
|
|
285
453
|
.map_init(
|
|
286
|
-
|| Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new() },
|
|
454
|
+
|| Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() },
|
|
287
455
|
|scratch, i| {
|
|
288
|
-
let (di, wi) = c.row(i);
|
|
289
|
-
accumulate(&index, (di, wi), i as u32, scratch);
|
|
290
456
|
let mut out = Vec::new();
|
|
291
|
-
|
|
457
|
+
probe_half(c, i, t, index, cached, mode, scratch, &mut out);
|
|
292
458
|
out
|
|
293
459
|
},
|
|
294
460
|
)
|
|
@@ -296,6 +462,118 @@ pub fn cosine_join(c: &Corpus, t: f64) -> Vec<(usize, usize, f64)> {
|
|
|
296
462
|
.collect()
|
|
297
463
|
}
|
|
298
464
|
|
|
465
|
+
/// One probe's accumulate + prune + score against `c`'s index, under `mode`, appending to `out`.
|
|
466
|
+
/// Shared by the single join and the fused two-corpus join so the per-half work is structurally
|
|
467
|
+
/// identical (the fused join relies on that for its bit-identical halves).
|
|
468
|
+
#[allow(clippy::too_many_arguments, clippy::inline_always)]
|
|
469
|
+
#[inline(always)] // per-probe dispatch only — must fold into the rayon closure, not add a call
|
|
470
|
+
fn probe_half(
|
|
471
|
+
c: &Corpus,
|
|
472
|
+
i: usize,
|
|
473
|
+
t: f64,
|
|
474
|
+
index: &[Vec<(u32, f64)>],
|
|
475
|
+
cached: &Cached,
|
|
476
|
+
mode: AccMode,
|
|
477
|
+
s: &mut Scratch,
|
|
478
|
+
out: &mut Vec<(usize, usize, f64)>,
|
|
479
|
+
) {
|
|
480
|
+
let (di, wi) = c.row(i);
|
|
481
|
+
match mode {
|
|
482
|
+
AccMode::Dense => {
|
|
483
|
+
accumulate(index, (di, wi), i as u32, s);
|
|
484
|
+
verify_pruned(c, i, t, s, cached, out);
|
|
485
|
+
}
|
|
486
|
+
AccMode::Sparse => {
|
|
487
|
+
if accumulate_list(index, (di, wi), i as u32, s) {
|
|
488
|
+
verify_list(c, i, t, s, cached, out);
|
|
489
|
+
} else {
|
|
490
|
+
// The sparse list overflowed — redo this probe with the dense pass (identical
|
|
491
|
+
// additions, identical order, so the values are unchanged).
|
|
492
|
+
accumulate(index, (di, wi), i as u32, s);
|
|
493
|
+
verify_pruned(c, i, t, s, cached, out);
|
|
494
|
+
}
|
|
495
|
+
}
|
|
496
|
+
}
|
|
497
|
+
}
|
|
498
|
+
|
|
499
|
+
/// [`cosine_join`] under an explicitly forced [`AccMode`]. Test-only: the fuzz gates must cover both
|
|
500
|
+
/// modes, and the shape-based pick exercises only one of them per corpus.
|
|
501
|
+
#[cfg(test)]
|
|
502
|
+
#[must_use]
|
|
503
|
+
fn cosine_join_for_mode(c: &Corpus, t: f64, mode: AccMode) -> Vec<(usize, usize, f64)> {
|
|
504
|
+
let n = c.n;
|
|
505
|
+
let mut index: Vec<Vec<(u32, f64)>> = vec![Vec::new(); c.ndims];
|
|
506
|
+
let mut cached = Cached { bound: vec![Bound { pnorm: 0.0, split: u32::MAX }; n] };
|
|
507
|
+
for i in 0..n {
|
|
508
|
+
let (di, wi) = c.row(i);
|
|
509
|
+
index_suffix(c, i, (di, wi), t, &mut index, &mut cached);
|
|
510
|
+
}
|
|
511
|
+
cosine_join_mode(c, t, &index, &cached, mode)
|
|
512
|
+
}
|
|
513
|
+
|
|
514
|
+
/// Fused two-corpus cosine join: one probe pass over `a` and `b` — two corpora with the **same row
|
|
515
|
+
/// count**, the typical "two views of the same objects" shape (e.g. line-vectors and part-vectors of
|
|
516
|
+
/// the same functions). Returns `(pairs_a, pairs_b)`; each half is **bit-identical** to a separate
|
|
517
|
+
/// [`cosine_join`] on its corpus — the per-half phases ([`accumulate`]/[`accumulate_list`],
|
|
518
|
+
/// [`verify_pruned`]/[`verify_list`]) run on the same inputs in the same order, only the rayon
|
|
519
|
+
/// schedule and the per-worker [`Scratch`] are shared (one `acc` buffer serves both halves: half A's
|
|
520
|
+
/// entries are reset by its verify pass before half B accumulates).
|
|
521
|
+
///
|
|
522
|
+
/// Wins over two separate joins by paying the per-probe fixed costs once: one parallel schedule, one
|
|
523
|
+
/// `Scratch` per worker, and the probe iteration bookkeeping is shared. Pair order within each half
|
|
524
|
+
/// is arbitrary (as in [`cosine_join`]). The accumulator mode is chosen per corpus from its shape
|
|
525
|
+
/// (see [`cosine_join`]), so the halves may use different modes.
|
|
526
|
+
///
|
|
527
|
+
/// # Panics
|
|
528
|
+
///
|
|
529
|
+
/// Panics if `a.len() != b.len()`.
|
|
530
|
+
#[must_use]
|
|
531
|
+
pub fn cosine_join_double(a: &Corpus, b: &Corpus, t: f64) -> (JoinPairs, JoinPairs) {
|
|
532
|
+
assert_eq!(a.n, b.n, "cosine_join_double: corpora must share the row count");
|
|
533
|
+
let n = a.n;
|
|
534
|
+
let mut index_a: Vec<Vec<(u32, f64)>> = vec![Vec::new(); a.ndims];
|
|
535
|
+
let mut index_b: Vec<Vec<(u32, f64)>> = vec![Vec::new(); b.ndims];
|
|
536
|
+
let mut cached_a = Cached { bound: vec![Bound { pnorm: 0.0, split: u32::MAX }; n] };
|
|
537
|
+
let mut cached_b = Cached { bound: vec![Bound { pnorm: 0.0, split: u32::MAX }; n] };
|
|
538
|
+
for i in 0..n {
|
|
539
|
+
let (di, wi) = a.row(i);
|
|
540
|
+
index_suffix(a, i, (di, wi), t, &mut index_a, &mut cached_a);
|
|
541
|
+
let (di, wi) = b.row(i);
|
|
542
|
+
index_suffix(b, i, (di, wi), t, &mut index_b, &mut cached_b);
|
|
543
|
+
}
|
|
544
|
+
let mode_a = acc_mode_for(&index_a, n);
|
|
545
|
+
let mode_b = acc_mode_for(&index_b, n);
|
|
546
|
+
(0..n)
|
|
547
|
+
.into_par_iter()
|
|
548
|
+
.with_min_len(256)
|
|
549
|
+
.map_init(
|
|
550
|
+
|| Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() },
|
|
551
|
+
|scratch, i| {
|
|
552
|
+
let mut out_a = Vec::new();
|
|
553
|
+
probe_half(a, i, t, &index_a, &cached_a, mode_a, scratch, &mut out_a);
|
|
554
|
+
let mut out_b = Vec::new();
|
|
555
|
+
probe_half(b, i, t, &index_b, &cached_b, mode_b, scratch, &mut out_b);
|
|
556
|
+
(out_a, out_b)
|
|
557
|
+
},
|
|
558
|
+
)
|
|
559
|
+
.fold(
|
|
560
|
+
|| (Vec::new(), Vec::new()),
|
|
561
|
+
|(mut va, mut vb), (oa, ob)| {
|
|
562
|
+
va.extend(oa);
|
|
563
|
+
vb.extend(ob);
|
|
564
|
+
(va, vb)
|
|
565
|
+
},
|
|
566
|
+
)
|
|
567
|
+
.reduce(
|
|
568
|
+
|| (Vec::new(), Vec::new()),
|
|
569
|
+
|(mut a1, mut b1), (mut a2, mut b2)| {
|
|
570
|
+
a1.append(&mut a2);
|
|
571
|
+
b1.append(&mut b2);
|
|
572
|
+
(a1, b1)
|
|
573
|
+
},
|
|
574
|
+
)
|
|
575
|
+
}
|
|
576
|
+
|
|
299
577
|
/// Run the cosine join under a chosen [`Concurrency`] backend. Returns `(j, i, cos)` pairs with
|
|
300
578
|
/// `j < i` and `cos ≥ t`, scores as `f64` (the `Gpu` mode's f32 cosines are widened losslessly).
|
|
301
579
|
///
|
|
@@ -478,9 +756,60 @@ fn accumulate(index: &[Vec<(u32, f64)>], (di, wi): (&[u32], &[f64]), cutoff: u32
|
|
|
478
756
|
unsafe { touched.set_len(tlen) };
|
|
479
757
|
}
|
|
480
758
|
|
|
481
|
-
///
|
|
482
|
-
///
|
|
483
|
-
///
|
|
759
|
+
/// Fill `xpn` with the probe's running prefix L2 norms (`xpn[k] = ‖wi[..k]‖₂`) and return the full
|
|
760
|
+
/// probe norm plus whether the cheap monotone pre-bound is worth it for this row length. One sqrt
|
|
761
|
+
/// per probe dim, reused across all its candidates.
|
|
762
|
+
///
|
|
763
|
+
/// Since `xpn` is monotonic, `xpn[kstar] ≤ xnorm` for every candidate, so `a + xnorm·pnorm ≥ a +
|
|
764
|
+
/// xpn[kstar]·pnorm` (the exact bound). A candidate failing the cheap `a + xnorm·pnorm < need` test
|
|
765
|
+
/// therefore also fails the exact bound — prune it without the per-candidate `partition_point`. The
|
|
766
|
+
/// survivor set is bit-identical; only the binary search is skipped for the ~99% of touched
|
|
767
|
+
/// candidates that can't reach `t` on dense corpora.
|
|
768
|
+
#[must_use]
|
|
769
|
+
fn probe_prefix_norms(wi: &[f64], xpn: &mut Vec<f64>) -> (f64, bool) {
|
|
770
|
+
xpn.clear();
|
|
771
|
+
xpn.push(0.0);
|
|
772
|
+
let mut sq = 0.0f64;
|
|
773
|
+
for &w in wi {
|
|
774
|
+
sq += w * w;
|
|
775
|
+
xpn.push(sq.sqrt());
|
|
776
|
+
}
|
|
777
|
+
(sq.sqrt(), wi.len() >= PREBOUND_MIN_DIMS)
|
|
778
|
+
}
|
|
779
|
+
|
|
780
|
+
/// Per-candidate Cauchy–Schwarz prune test, shared by the dense verify pass, the sparse verify pass
|
|
781
|
+
/// and the GPU survivor collector: `true` iff candidate `y` (accumulated partial dot `a`) survives —
|
|
782
|
+
/// i.e. its exact [`cos_full`] score must be computed. The bound is a filter only; the emitted value
|
|
783
|
+
/// is always the exact dot, so the pair set is bit-identical to brute force whatever this decides.
|
|
784
|
+
#[must_use]
|
|
785
|
+
#[inline(always)]
|
|
786
|
+
#[allow(clippy::inline_always)] // this runs once per touched candidate — call overhead is the cost
|
|
787
|
+
fn candidate_survives(
|
|
788
|
+
di: &[u32],
|
|
789
|
+
xpn: &[f64],
|
|
790
|
+
xnorm: f64,
|
|
791
|
+
prebound: bool,
|
|
792
|
+
need: f64,
|
|
793
|
+
a: f64,
|
|
794
|
+
bd: Bound,
|
|
795
|
+
) -> bool {
|
|
796
|
+
// Cheap monotone pre-bound (dense rows only) — skip the binary search when even `xnorm` can't
|
|
797
|
+
// clear `need`. `prebound` is loop-invariant, so sparse rows pay nothing.
|
|
798
|
+
if prebound && a + xnorm * bd.pnorm < need {
|
|
799
|
+
return false;
|
|
800
|
+
}
|
|
801
|
+
// Number of probe dims with rank < split[y] → index into xpn (di sorted ascending).
|
|
802
|
+
let kstar = di.partition_point(|&d| d < bd.split);
|
|
803
|
+
// SAFETY: kstar ≤ di.len() == wi.len() = xpn.len()-1.
|
|
804
|
+
// (An added maxweight cap `min(…, Σ wx·maxw)` was tried + reverted: never tighter than the
|
|
805
|
+
// L2 cap on either synthetic or real data — it sums over all probe-prefix dims, not just the
|
|
806
|
+
// shared ones — so it pruned nothing and cost ~16% on the real PyPI corpus.)
|
|
807
|
+
a + unsafe { xpn.get_unchecked(kstar) } * bd.pnorm >= need
|
|
808
|
+
}
|
|
809
|
+
|
|
810
|
+
/// Phase 2 — prune + verify (dense mode). For each touched candidate `y`, reset its accumulator and
|
|
811
|
+
/// test the **L2AP `l2` bound**: the dot mass missing from `acc[y]` (the dims in `prefix(y)`) is at
|
|
812
|
+
/// most `‖x_{rank<split[y]}‖ · ‖prefix(y)‖` by Cauchy–Schwarz, where `x_{rank<split[y]}` is the probe
|
|
484
813
|
/// restricted to the rank range `prefix(y)` lives in. Since the probe's mass sits in its rare
|
|
485
814
|
/// (high-rank) dims, that restricted norm is tiny — a far tighter cap than the whole-probe `‖x‖=1`.
|
|
486
815
|
/// Only if `acc[y] + that bound ≥ t` (minus FP slack) do we score exactly with [`cos_full`]. Filter
|
|
@@ -495,24 +824,9 @@ fn verify_pruned(
|
|
|
495
824
|
out: &mut Vec<(usize, usize, f64)>,
|
|
496
825
|
) {
|
|
497
826
|
let (di, wi) = c.row(i);
|
|
498
|
-
|
|
499
|
-
// probe's k lowest-rank dims). One sqrt per probe dim, reused across all its candidates.
|
|
500
|
-
s.xpn.clear();
|
|
501
|
-
s.xpn.push(0.0);
|
|
502
|
-
let mut sq = 0.0f64;
|
|
503
|
-
for &w in wi {
|
|
504
|
-
sq += w * w;
|
|
505
|
-
s.xpn.push(sq.sqrt());
|
|
506
|
-
}
|
|
507
|
-
// `‖probe‖` = the full prefix norm. Since `xpn` is monotonic, `xpn[kstar] ≤ xnorm` for every
|
|
508
|
-
// candidate, so `a + xnorm·pnorm ≥ a + xpn[kstar]·pnorm` (the exact bound). A candidate failing
|
|
509
|
-
// the cheap `a + xnorm·pnorm < need` test therefore also fails the exact bound — prune it without
|
|
510
|
-
// the per-candidate `partition_point`. The survivor set is bit-identical; only the binary search
|
|
511
|
-
// is skipped for the ~99% of touched candidates that can't reach `t` on dense corpora.
|
|
512
|
-
let xnorm = sq.sqrt();
|
|
513
|
-
let prebound = di.len() >= PREBOUND_MIN_DIMS;
|
|
827
|
+
let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
|
|
514
828
|
let need = t - PRUNE_SLACK;
|
|
515
|
-
let Scratch { acc, touched, xpn } = s;
|
|
829
|
+
let Scratch { acc, touched, xpn, .. } = s;
|
|
516
830
|
// (Software-prefetching the candidate row a few ahead was tried + reverted: no measurable change
|
|
517
831
|
// — with all cores gathering at once the join is memory-*bandwidth*-bound, not per-access
|
|
518
832
|
// latency-bound, so prefetch can't add throughput.)
|
|
@@ -522,19 +836,73 @@ fn verify_pruned(
|
|
|
522
836
|
let a = unsafe { std::mem::replace(acc.get_unchecked_mut(yu), -1.0) };
|
|
523
837
|
// SAFETY: `yu < n`. One scattered load fetches both prune fields.
|
|
524
838
|
let bd = unsafe { *cached.bound.get_unchecked(yu) };
|
|
525
|
-
|
|
526
|
-
|
|
527
|
-
|
|
528
|
-
|
|
839
|
+
if candidate_survives(di, xpn, xnorm, prebound, need, a, bd) {
|
|
840
|
+
let cos = cos_full((di, wi), c.row(yu));
|
|
841
|
+
if cos >= t {
|
|
842
|
+
out.push((yu, i, cos));
|
|
843
|
+
}
|
|
844
|
+
}
|
|
845
|
+
}
|
|
846
|
+
}
|
|
847
|
+
|
|
848
|
+
/// Sparse-mode phase 1 — the same partial-dot accumulation as [`accumulate`], but into the sorted
|
|
849
|
+
/// `(y, acc)` list instead of the dense `acc[n]` array. Wins when a probe touches few candidates:
|
|
850
|
+
/// the list stays tiny and L1-resident, and the n-wide dense array's footprint never competes with
|
|
851
|
+
/// the postings for L2. The additions happen in the same per-dim/posting order, so the accumulated
|
|
852
|
+
/// values are bit-identical to the dense pass. Returns `false` on [`LIST_ACC_MAX`] overflow — the
|
|
853
|
+
/// caller must redo the probe densely (identical additions, identical order).
|
|
854
|
+
#[cfg_attr(feature = "profiling", inline(never))]
|
|
855
|
+
fn accumulate_list(
|
|
856
|
+
index: &[Vec<(u32, f64)>],
|
|
857
|
+
(di, wi): (&[u32], &[f64]),
|
|
858
|
+
cutoff: u32,
|
|
859
|
+
s: &mut Scratch,
|
|
860
|
+
) -> bool {
|
|
861
|
+
let Scratch { list, .. } = s;
|
|
862
|
+
list.clear();
|
|
863
|
+
for (&d, &w) in di.iter().zip(wi) {
|
|
864
|
+
for &(y, wy) in &index[d as usize] {
|
|
865
|
+
if y >= cutoff {
|
|
866
|
+
break;
|
|
867
|
+
}
|
|
868
|
+
let term = w * wy;
|
|
869
|
+
match list.binary_search_by_key(&y, |&(yy, _)| yy) {
|
|
870
|
+
Ok(k) => list[k].1 += term,
|
|
871
|
+
Err(k) => {
|
|
872
|
+
if list.len() >= LIST_ACC_MAX {
|
|
873
|
+
list.clear();
|
|
874
|
+
return false;
|
|
875
|
+
}
|
|
876
|
+
list.insert(k, (y, term));
|
|
877
|
+
}
|
|
878
|
+
}
|
|
529
879
|
}
|
|
530
|
-
|
|
531
|
-
|
|
532
|
-
|
|
533
|
-
|
|
534
|
-
|
|
535
|
-
|
|
536
|
-
|
|
537
|
-
|
|
880
|
+
}
|
|
881
|
+
true
|
|
882
|
+
}
|
|
883
|
+
|
|
884
|
+
/// Sparse-mode phase 2 — the same prune + score as [`verify_pruned`], over the sorted accumulator
|
|
885
|
+
/// list. No reset needed: the next probe's [`accumulate_list`] clears the list. Candidate order is
|
|
886
|
+
/// by id (the dense pass visits first-touch order); both are "arbitrary order" for the caller, and
|
|
887
|
+
/// the per-candidate arithmetic is identical, so the emitted pair set + scores are bit-identical.
|
|
888
|
+
#[cfg_attr(feature = "profiling", inline(never))]
|
|
889
|
+
fn verify_list(
|
|
890
|
+
c: &Corpus,
|
|
891
|
+
i: usize,
|
|
892
|
+
t: f64,
|
|
893
|
+
s: &mut Scratch,
|
|
894
|
+
cached: &Cached,
|
|
895
|
+
out: &mut Vec<(usize, usize, f64)>,
|
|
896
|
+
) {
|
|
897
|
+
let (di, wi) = c.row(i);
|
|
898
|
+
let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
|
|
899
|
+
let need = t - PRUNE_SLACK;
|
|
900
|
+
let Scratch { list, xpn, .. } = s;
|
|
901
|
+
for &(y, a) in list.iter() {
|
|
902
|
+
let yu = y as usize;
|
|
903
|
+
// SAFETY: `yu < n` (candidate ids come from the postings, same as in `accumulate`).
|
|
904
|
+
let bd = unsafe { *cached.bound.get_unchecked(yu) };
|
|
905
|
+
if candidate_survives(di, xpn, xnorm, prebound, need, a, bd) {
|
|
538
906
|
let cos = cos_full((di, wi), c.row(yu));
|
|
539
907
|
if cos >= t {
|
|
540
908
|
out.push((yu, i, cos));
|
|
@@ -592,34 +960,21 @@ fn index_suffix(
|
|
|
592
960
|
const GPU_FILTER_MARGIN: f64 = 1e-4;
|
|
593
961
|
|
|
594
962
|
/// Like the verify half of [`verify_pruned`], but instead of scoring, **collects** each surviving
|
|
595
|
-
/// `(candidate, probe)` pair (`candidate < probe`) for batch scoring elsewhere. The bound
|
|
596
|
-
///
|
|
963
|
+
/// `(candidate, probe)` pair (`candidate < probe`) for batch scoring elsewhere. The bound test is
|
|
964
|
+
/// [`candidate_survives`] — the same function `verify_pruned` uses — so the survivor set is
|
|
965
|
+
/// identical by construction.
|
|
597
966
|
#[cfg(all(target_os = "macos", feature = "gpu"))]
|
|
598
967
|
fn collect_survivors(c: &Corpus, i: usize, t: f64, s: &mut Scratch, cached: &Cached, out: &mut Vec<(u32, u32)>) {
|
|
599
968
|
let (di, wi) = c.row(i);
|
|
600
|
-
s.xpn
|
|
601
|
-
s.xpn.push(0.0);
|
|
602
|
-
let mut sq = 0.0f64;
|
|
603
|
-
for &w in wi {
|
|
604
|
-
sq += w * w;
|
|
605
|
-
s.xpn.push(sq.sqrt());
|
|
606
|
-
}
|
|
607
|
-
let xnorm = sq.sqrt();
|
|
608
|
-
let prebound = di.len() >= PREBOUND_MIN_DIMS;
|
|
969
|
+
let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
|
|
609
970
|
let need = t - PRUNE_SLACK;
|
|
610
|
-
let Scratch { acc, touched, xpn } = s;
|
|
971
|
+
let Scratch { acc, touched, xpn, .. } = s;
|
|
611
972
|
for &y in touched.iter() {
|
|
612
973
|
let yu = y as usize;
|
|
613
974
|
// SAFETY: `yu < n` (same provenance as in `accumulate`).
|
|
614
975
|
let a = unsafe { std::mem::replace(acc.get_unchecked_mut(yu), -1.0) };
|
|
615
976
|
let bd = unsafe { *cached.bound.get_unchecked(yu) };
|
|
616
|
-
|
|
617
|
-
if prebound && a + xnorm * bd.pnorm < need {
|
|
618
|
-
continue;
|
|
619
|
-
}
|
|
620
|
-
let kstar = di.partition_point(|&d| d < bd.split);
|
|
621
|
-
let bound = a + unsafe { xpn.get_unchecked(kstar) } * bd.pnorm;
|
|
622
|
-
if bound >= need {
|
|
977
|
+
if candidate_survives(di, xpn, xnorm, prebound, need, a, bd) {
|
|
623
978
|
out.push((y, i as u32)); // (candidate, probe), candidate < probe
|
|
624
979
|
}
|
|
625
980
|
}
|
|
@@ -706,7 +1061,7 @@ fn survivor_pairs(c: &Corpus, t: f64) -> (Vec<u32>, Vec<u32>) {
|
|
|
706
1061
|
.into_par_iter()
|
|
707
1062
|
.with_min_len(256)
|
|
708
1063
|
.map_init(
|
|
709
|
-
|| Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new() },
|
|
1064
|
+
|| Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() },
|
|
710
1065
|
|scratch, i| {
|
|
711
1066
|
let (di, wi) = c.row(i);
|
|
712
1067
|
accumulate(&index, (di, wi), i as u32, scratch);
|
|
@@ -734,31 +1089,19 @@ pub fn cosine_join_counts(c: &Corpus, t: f64) -> (u64, u64, u64) {
|
|
|
734
1089
|
let (di, wi) = c.row(i);
|
|
735
1090
|
index_suffix(c, i, (di, wi), t, &mut index, &mut cached);
|
|
736
1091
|
}
|
|
737
|
-
let mut s = Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new() };
|
|
1092
|
+
let mut s = Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() };
|
|
738
1093
|
let (mut ncand, mut survivors, mut pairs) = (0u64, 0u64, 0u64);
|
|
739
1094
|
let need = t - PRUNE_SLACK;
|
|
740
1095
|
for i in 0..n {
|
|
741
1096
|
let (di, wi) = c.row(i);
|
|
742
1097
|
accumulate(&index, (di, wi), i as u32, &mut s);
|
|
743
1098
|
ncand += s.touched.len() as u64;
|
|
744
|
-
s.xpn
|
|
745
|
-
s.xpn.push(0.0);
|
|
746
|
-
let mut sq = 0.0f64;
|
|
747
|
-
for &w in wi {
|
|
748
|
-
sq += w * w;
|
|
749
|
-
s.xpn.push(sq.sqrt());
|
|
750
|
-
}
|
|
751
|
-
let xnorm = sq.sqrt();
|
|
752
|
-
let prebound = di.len() >= PREBOUND_MIN_DIMS;
|
|
1099
|
+
let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
|
|
753
1100
|
for &y in &s.touched {
|
|
754
1101
|
let yu = y as usize;
|
|
755
1102
|
let a = std::mem::replace(&mut s.acc[yu], -1.0);
|
|
756
1103
|
let bd = cached.bound[yu];
|
|
757
|
-
if
|
|
758
|
-
continue;
|
|
759
|
-
}
|
|
760
|
-
let kstar = di.partition_point(|&d| d < bd.split);
|
|
761
|
-
if a + s.xpn[kstar] * bd.pnorm >= need {
|
|
1104
|
+
if candidate_survives(di, &s.xpn, xnorm, prebound, need, a, bd) {
|
|
762
1105
|
survivors += 1;
|
|
763
1106
|
if cos_full((di, wi), c.row(yu)) >= t {
|
|
764
1107
|
pairs += 1;
|
|
@@ -787,7 +1130,10 @@ pub fn cosine_join_bruteforce(c: &Corpus, t: f64) -> Vec<(usize, usize, f64)> {
|
|
|
787
1130
|
|
|
788
1131
|
#[cfg(test)]
|
|
789
1132
|
mod tests {
|
|
790
|
-
use super::{
|
|
1133
|
+
use super::{
|
|
1134
|
+
cosine_join, cosine_join_bruteforce, cosine_join_double, cosine_join_for_mode, AccMode,
|
|
1135
|
+
Corpus,
|
|
1136
|
+
};
|
|
791
1137
|
|
|
792
1138
|
fn xorshift(seed: u64) -> impl FnMut() -> u64 {
|
|
793
1139
|
let mut s = seed;
|
|
@@ -827,6 +1173,145 @@ mod tests {
|
|
|
827
1173
|
}
|
|
828
1174
|
}
|
|
829
1175
|
|
|
1176
|
+
/// `Corpus::from_rows` must be a deterministic function of its input: building the same rows
|
|
1177
|
+
/// twice yields the same corpus, so the join output is bit-identical. (The previous
|
|
1178
|
+
/// HashMap-based per-row merge summed weights in per-call-random hash order, so this did not
|
|
1179
|
+
/// hold; it is now a pinned contract.)
|
|
1180
|
+
#[test]
|
|
1181
|
+
fn from_rows_is_deterministic() {
|
|
1182
|
+
let mut next = xorshift(0x6f6c_6c61_6279_6e6f);
|
|
1183
|
+
for _ in 0..20 {
|
|
1184
|
+
let n = (next() % 30 + 2) as usize;
|
|
1185
|
+
let dim_space = next() % 12 + 1;
|
|
1186
|
+
let rows: Vec<Vec<(u32, f64)>> = (0..n)
|
|
1187
|
+
.map(|_| {
|
|
1188
|
+
let nnz = (next() % 9) as usize;
|
|
1189
|
+
let mut r: Vec<(u32, f64)> = (0..nnz)
|
|
1190
|
+
.map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
|
|
1191
|
+
.collect();
|
|
1192
|
+
// Force duplicate dims so the merge path runs and the sum order matters.
|
|
1193
|
+
if nnz > 1 {
|
|
1194
|
+
r[1].0 = r[0].0;
|
|
1195
|
+
}
|
|
1196
|
+
r
|
|
1197
|
+
})
|
|
1198
|
+
.collect();
|
|
1199
|
+
let a = Corpus::from_rows(&rows);
|
|
1200
|
+
let b = Corpus::from_rows(&rows);
|
|
1201
|
+
for &t in &[0.3_f64, 0.7] {
|
|
1202
|
+
assert_eq!(
|
|
1203
|
+
sort_pairs(cosine_join(&a, t)),
|
|
1204
|
+
sort_pairs(cosine_join(&b, t)),
|
|
1205
|
+
"n={n} t={t}"
|
|
1206
|
+
);
|
|
1207
|
+
}
|
|
1208
|
+
}
|
|
1209
|
+
}
|
|
1210
|
+
|
|
1211
|
+
/// Stripping `df = 1` dims from the corpus must be invisible to the join: a hapax dim is shared
|
|
1212
|
+
/// by no pair, so it contributes exactly `0.0` to every pairwise dot (the branchless `cos_full`
|
|
1213
|
+
/// accumulates `+0.0` no-ops on non-negative sums), its posting is never walked, and the tighter
|
|
1214
|
+
/// Cauchy–Schwarz bounds only prune more. The stripped and un-stripped builds must therefore
|
|
1215
|
+
/// return bit-identical pair sets, and each must match its own brute-force oracle.
|
|
1216
|
+
#[test]
|
|
1217
|
+
fn hapax_stripping_is_invisible() {
|
|
1218
|
+
let mut next = xorshift(0x51a7_4e5f_0c92_83d6);
|
|
1219
|
+
for _ in 0..200 {
|
|
1220
|
+
let n = (next() % 40 + 2) as usize;
|
|
1221
|
+
let dim_space = next() % 12 + 1;
|
|
1222
|
+
let rows: Vec<Vec<(u32, f64)>> = (0..n)
|
|
1223
|
+
.map(|i| {
|
|
1224
|
+
let nnz = (next() % 8) as usize;
|
|
1225
|
+
let mut r: Vec<(u32, f64)> = (0..nnz)
|
|
1226
|
+
.map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
|
|
1227
|
+
.collect();
|
|
1228
|
+
// Every row carries one unique dim — a guaranteed hapax the strip must remove
|
|
1229
|
+
// invisibly (its mass participates in the L2 normalisation of the full row).
|
|
1230
|
+
r.push((1000 + i as u32, 1.0));
|
|
1231
|
+
r
|
|
1232
|
+
})
|
|
1233
|
+
.collect();
|
|
1234
|
+
let full = Corpus::from_rows_with(&rows, false);
|
|
1235
|
+
let stripped = Corpus::from_rows_with(&rows, true);
|
|
1236
|
+
for &t in &[0.1_f64, 0.25, 0.5, 0.75, 0.9, 1.0] {
|
|
1237
|
+
let want = sort_pairs(cosine_join_bruteforce(&stripped, t));
|
|
1238
|
+
assert_eq!(
|
|
1239
|
+
sort_pairs(cosine_join(&stripped, t)),
|
|
1240
|
+
want,
|
|
1241
|
+
"stripped vs oracle, n={n} t={t}"
|
|
1242
|
+
);
|
|
1243
|
+
assert_eq!(
|
|
1244
|
+
sort_pairs(cosine_join(&full, t)),
|
|
1245
|
+
want,
|
|
1246
|
+
"stripped vs full, n={n} t={t}"
|
|
1247
|
+
);
|
|
1248
|
+
assert_eq!(
|
|
1249
|
+
sort_pairs(cosine_join(&full, t)),
|
|
1250
|
+
sort_pairs(cosine_join_bruteforce(&full, t)),
|
|
1251
|
+
"full vs its own oracle, n={n} t={t}"
|
|
1252
|
+
);
|
|
1253
|
+
}
|
|
1254
|
+
}
|
|
1255
|
+
}
|
|
1256
|
+
|
|
1257
|
+
/// Both accumulator strategies must agree bit-for-bit with brute force. The shape-based mode
|
|
1258
|
+
/// pick in `cosine_join` means any given corpus exercises only one strategy — force each here
|
|
1259
|
+
/// (the dense pass is what the large-corpus shape selects, the sparse list the duplicate-heavy
|
|
1260
|
+
/// consumer shape).
|
|
1261
|
+
#[test]
|
|
1262
|
+
fn both_acc_modes_match_bruteforce() {
|
|
1263
|
+
let mut next = xorshift(0x5eed_0f1a_2b3c_4d5e);
|
|
1264
|
+
for _ in 0..200 {
|
|
1265
|
+
let n = (next() % 40 + 2) as usize;
|
|
1266
|
+
let dim_space = next() % 15 + 1;
|
|
1267
|
+
let rows: Vec<Vec<(u32, f64)>> = (0..n)
|
|
1268
|
+
.map(|_| {
|
|
1269
|
+
let nnz = (next() % 8) as usize;
|
|
1270
|
+
(0..nnz)
|
|
1271
|
+
.map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
|
|
1272
|
+
.collect()
|
|
1273
|
+
})
|
|
1274
|
+
.collect();
|
|
1275
|
+
let c = Corpus::from_rows(&rows);
|
|
1276
|
+
for &t in &[0.1_f64, 0.25, 0.5, 0.75, 0.9, 1.0] {
|
|
1277
|
+
let want = sort_pairs(cosine_join_bruteforce(&c, t));
|
|
1278
|
+
for mode in [AccMode::Dense, AccMode::Sparse] {
|
|
1279
|
+
let got = sort_pairs(cosine_join_for_mode(&c, t, mode));
|
|
1280
|
+
assert_eq!(got, want, "n={n} t={t} mode={mode:?}");
|
|
1281
|
+
}
|
|
1282
|
+
}
|
|
1283
|
+
}
|
|
1284
|
+
}
|
|
1285
|
+
|
|
1286
|
+
/// The fused two-corpus join must return, per half, the bit-identical pair set of a separate
|
|
1287
|
+
/// `cosine_join` on that corpus — fusion shares only schedule and scratch, never arithmetic.
|
|
1288
|
+
#[test]
|
|
1289
|
+
fn double_join_matches_single() {
|
|
1290
|
+
let mut next = xorshift(0xd0a2_7b7e_9c31_5f84);
|
|
1291
|
+
for _ in 0..200 {
|
|
1292
|
+
let n = (next() % 40 + 2) as usize;
|
|
1293
|
+
let dim_space = next() % 15 + 1;
|
|
1294
|
+
let mut gen = || -> Vec<Vec<(u32, f64)>> {
|
|
1295
|
+
(0..n)
|
|
1296
|
+
.map(|_| {
|
|
1297
|
+
let nnz = (next() % 8) as usize;
|
|
1298
|
+
(0..nnz)
|
|
1299
|
+
.map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
|
|
1300
|
+
.collect()
|
|
1301
|
+
})
|
|
1302
|
+
.collect()
|
|
1303
|
+
};
|
|
1304
|
+
let (rows_a, rows_b) = (gen(), gen());
|
|
1305
|
+
let a = Corpus::from_rows(&rows_a);
|
|
1306
|
+
let b = Corpus::from_rows(&rows_b);
|
|
1307
|
+
for &t in &[0.1_f64, 0.3, 0.5, 0.7, 0.9, 1.0] {
|
|
1308
|
+
let (ga, gb) = cosine_join_double(&a, &b, t);
|
|
1309
|
+
assert_eq!(sort_pairs(ga), sort_pairs(cosine_join(&a, t)), "half a n={n} t={t}");
|
|
1310
|
+
assert_eq!(sort_pairs(gb), sort_pairs(cosine_join(&b, t)), "half b n={n} t={t}");
|
|
1311
|
+
}
|
|
1312
|
+
}
|
|
1313
|
+
}
|
|
1314
|
+
|
|
830
1315
|
/// The CPU+GPU hybrid [`super::cosine_join_gpu`] must return bit-identical results to the pure-CPU
|
|
831
1316
|
/// [`cosine_join`] on fuzzed corpora — the GPU is only a conservative filter; every emitted score
|
|
832
1317
|
/// is the exact CPU `f64` value. Skips when no Metal device is present.
|
|
@@ -164,7 +164,14 @@ impl BatchCosineGpu {
|
|
|
164
164
|
}
|
|
165
165
|
|
|
166
166
|
/// Copy `data` into a fresh shared (UMA) buffer — one memcpy into unified memory (zero device copy).
|
|
167
|
+
/// An empty slice gets a 1-byte zeroed buffer instead: `new_buffer_with_data` would memcpy 1 byte
|
|
168
|
+
/// from `data.as_ptr()`, which is a dangling pointer for empty slices (a real SIGSEGV once hapax
|
|
169
|
+
/// stripping made fully-empty CSR arrays reachable). The kernel never reads the buffer when the
|
|
170
|
+
/// corpus has no entries (every row's `indptr` window is empty), so the 1-byte placeholder is inert.
|
|
167
171
|
fn upload<T: Copy>(device: &Device, data: &[T]) -> Buffer {
|
|
168
|
-
|
|
172
|
+
if data.is_empty() {
|
|
173
|
+
return device.new_buffer(1, MTLResourceOptions::StorageModeShared);
|
|
174
|
+
}
|
|
175
|
+
let bytes = std::mem::size_of_val(data) as NSUInteger;
|
|
169
176
|
device.new_buffer_with_data(data.as_ptr().cast::<c_void>(), bytes, MTLResourceOptions::StorageModeShared)
|
|
170
177
|
}
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|