src/vector/gpu.rs
//! GPU-batched kNN over the **q8 vector space** — the same `i8` corpus
//! [`VectorIndex::quantized`] already hands the VNNI kernel, scored on a WebGPU
//! compute pipeline instead of on cores.
//!
//! # Why WebGPU and not CUDA
//!
//! MEASURED on oden 2026-08-25, not assumed:
//!
//! - `libcuda.so.1` + `libnvidia-ml.so.1` are present (driver 595.84), but there
//! is **no CUDA toolkit**: no `nvcc`, no `libnvrtc`, no `libcublas`. A CUDA
//! kernel path would need either a build-time toolchain this machine does not
//! have, or hand-written PTX loaded through the driver API — FFI, in a crate
//! whose stated design priority is "100% Rust, no C/FFI, self-contained
//! airgapped binary" ([`crate::vector`] module docs).
//! - `wgpu 29.0.4` is **already in `Cargo.lock`** (pulled by `eframe` for the viz
//! tabs), as are `pollster 0.4` and `bytemuck 1.25`. The GPU path therefore
//! adds **zero packages** to the resolution graph.
//! - The 4090 answers on Vulkan 1.4 with `VK_KHR_shader_integer_dot_product`, so
//! `dot4I8Packed` lowers to a native `OpSDot` — the DP4a instruction, the exact
//! GPU analogue of the `vpdpbusd` the CPU kernel wants and this box's Zen 2
//! Threadripper does not have.
//!
//! So the modern-platform rule (workspace LAW 4) and the no-FFI rule point at the
//! same answer, and it happens to be the free one.
//!
//! # Why the result is IDENTICAL, not "close"
//!
//! The int8 cosine is `Σ qᵢ·rᵢ / (Q·Q)` with `Q = 127`: an **integer** dot,
//! divided once at the end. Integer addition is associative and here cannot
//! overflow (`|qᵢ·rᵢ| ≤ 127² = 16 129`, so `|Σ| ≤ 1.06e9 < 2³¹` for any
//! `dim ≤ 65 536`), so **every** accumulation order — 32 CPU cores, 16 384 CUDA
//! lanes, one scalar loop — produces the *same* `i32`. One `f32` division of the
//! same integer then produces the same `f32`.
//!
//! That is why [`GpuKnn::knn`] is checked against
//! [`znippy_zoomies::vann::score_i8_batch`] for **bit-identity**, not for
//! agreement within a tolerance: a tolerance-based check on this pipeline would
//! be a hollow guard, because the correct answer is exact.
//! `gpu_knn_matches_cpu_int8_bit_for_bit` is red-proven by
//! [`BREAK_ENV`] — see `a_corrupted_gpu_corpus_fails_the_oracle`.
//!
//! # No silent fallback
//!
//! A Vulkan enumeration on this box returns **two** adapters: the 4090 and
//! `llvmpipe`, Mesa's software rasteriser. `llvmpipe` is a CPU wearing a GPU's
//! API — it would satisfy every assertion in this module while making every
//! timing a lie. [`open`] refuses it by `DeviceType` *and* by name and says so,
//! following the precedent in `.nornir/build-variants.toml`: *"A host without a
//! `RenderState` gets a loud line and the CPU painter, never a silent
//! downgrade."*
//!
//! # Measured throughput — **on a SYNTHETIC corpus**
//!
//! Read the provenance line before the numbers. The rows below were **not**
//! embeddings of anything. They are `tests::fixture(n, 768, 64, 0.02)`: LCG
//! output (seed `0xC0FFEE`) shaped into 64 near-orthogonal centroids with a
//! tight noise ball around each, then L2-normalised. **No deps-embedding corpus
//! exists on this disk** — the warehouse that would hold one was erased, and
//! nothing in this repo has ever embedded 2 000 000 rows. Every row count in
//! the table is a fixture in the right *shape and scale*, and nothing more.
//!
//! That is defensible for exactly one reason, and it is worth stating rather
//! than leaving implicit: this is an **exact flat scan**. Every row is read and
//! scored on every query regardless of its contents, so the wall time is
//! `O(n · dim)` in the *shape* of the corpus and is not a function of the
//! values in it. A synthetic corpus of the right dimensions therefore measures
//! the same thing a real one would — but it does NOT establish that a corpus of
//! that size exists, and the two claims were previously run together here.
//!
//! Anything printed from code carries its origin as data, not as prose: see
//! [`CorpusOrigin`] and [`Throughput::line`], which cannot render a figure
//! without its provenance.
//!
//! dim 768, k = 8, 1024 queries, RTX 4090 vs a 32-core Threadripper PRO 3975WX,
//! oden 2026-08-25:
//!
//! ```text
//! rows GPU CPU f32 ratio loadavg GPU rate corpus
//! 79 000 0.038 s 0.769 s 20x ~5 2.12 Gpair/s SYNTHETIC
//! 1 000 000 0.636 s 5.98 s 9x 26 1.61 Gpair/s SYNTHETIC
//! 2 000 000 1.744 s 37.17 s 21x 121 1.17 Gpair/s SYNTHETIC
//! ```
//!
//! No committed benchmark reproduces this table — it was an ad-hoc run, and
//! saying so is part of the figure.
//!
//! The ratio column is NOT a clean number and should not be quoted as one. The
//! CPU baseline was taken on a shared box whose load average ranged from 5 to
//! 121 across these runs, so it swings by 4x for reasons that have nothing to do
//! with this code. The GPU column is stable to about +/-1% run to run.
//!
//! The defensible statement is the one built from the stable side: the GPU holds
//! 1.2-2.1 Gpair/s while the CPU's BEST observed rate is 0.11 Gpair/s, so the
//! honest floor is **about 10x**, and today's 21x is a contended machine
//! flattering it.
//!
//! # How big is the deps corpus actually?
//!
//! The old text here said 1.7 ms at 2 000 000 rows is *"the millisecond class
//! the deps corpus needs"*, which quietly asserted that the deps corpus reaches
//! that size. Nobody had counted. MEASURED 2026-08-26 by running the real
//! [`crate::vector::chunk::chunk_file`] at `ChunkOptions::default()`
//! (`max_lines = 120`, `overlap = 16`) over every `.rs` file of every crate
//! source actually unpacked under `~/.cargo/registry/src` — see
//! [`DEPS_CORPUS_MEASURED`].
//!
//! **5 702 262 rows.** That is **2.85x** the largest synthetic point, so the
//! table does not flatter the workload — it understates it. Extrapolating the
//! flat scan linearly gives ~5.0 s per 1024-query batch at real deps scale,
//! ~4.9 ms per query: still millisecond class, and now resting on a counted
//! corpus instead of an assumed one. See [`DEPS_CORPUS_VS_FIXTURE`].
//!
//! This is a real count of real files. It is still **not** an embedding corpus:
//! nothing has embedded them, and until something does, every timing in this
//! module stays labelled SYNTHETIC.
//!
//! Cargo feature: `vector-gpu`.
use anyhow::{Context, Result, bail};
use znippy_zoomies::vann::VectorIndex;
/// The int8 quantization scale, identical to `vann`'s private `Q`. Kept as one
/// constant used by both the packer and the score reconstruction so the two
/// cannot drift (LAW 5).
const Q: f32 = 127.0;
/// Queries resident in one dispatch. The shader keeps the whole tile in
/// workgroup memory, so each corpus word is loaded once and consumed 16 times;
/// 16 × 192 words (dim 768) = 12 KiB, inside the 16 KiB workgroup-storage floor
/// every WebGPU implementation guarantees.
pub const QUERY_TILE: usize = 16;
/// Largest `k` the shader's fixed-size per-thread heap supports.
pub const K_MAX: usize = 8;
/// Row blocks per query in the top-k stage. Each block produces `k` candidates,
/// so the host merges `NBLOCKS · k` per query — small enough that the readback
/// is a rounding error next to the scan.
const NBLOCKS: usize = 64;
/// Set this to deliberately corrupt the packed corpus, to prove the oracle can
/// go RED. Reading it in the shipping path (rather than only in tests) is what
/// makes the red proof a property of the *real* code path and not of a parallel
/// test-only copy.
pub const BREAK_ENV: &str = "NORNIR_GPU_KNN_BREAK";
/// Why a GPU kNN did or did not happen. Always rendered, never swallowed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GpuVerdict {
/// Scored on a real discrete GPU.
Gpu {
/// Adapter name as the driver reports it.
adapter: String,
/// Graphics backend (`Vulkan`).
backend: String,
},
/// No usable GPU. `reason` is a complete sentence, and `rejected` names every
/// adapter that was seen and turned down — a `llvmpipe`-only box must read as
/// "refused a software adapter", never as "no adapter".
Cpu {
/// Full explanation, suitable for printing verbatim.
reason: String,
/// Adapters enumerated and rejected, `backend/type name`.
rejected: Vec<String>,
},
}
impl GpuVerdict {
/// `true` only when a real discrete GPU did the arithmetic.
pub fn is_gpu(&self) -> bool {
matches!(self, GpuVerdict::Gpu { .. })
}
/// The one line a caller must print. A CPU verdict is prefixed `⚠` because a
/// downgrade the operator cannot see is the failure mode this whole module
/// is written against.
pub fn line(&self) -> String {
match self {
GpuVerdict::Gpu { adapter, backend } => {
format!("GPU kNN on {adapter} ({backend}).")
}
GpuVerdict::Cpu { reason, rejected } => {
let seen = if rejected.is_empty() {
"no adapters enumerated".to_string()
} else {
format!("adapters seen and refused: {}", rejected.join(", "))
};
format!("⚠ GPU kNN UNAVAILABLE — {reason} ({seen}). Scoring on CPU instead.")
}
}
}
}
// ---------------------------------------------------------------------------
// Corpus provenance — a number may not be rendered without where it came from
// ---------------------------------------------------------------------------
/// **Where the rows behind a throughput figure came from.**
///
/// Every corpus figure this module has ever produced is synthetic, and that was
/// discoverable only by reading the test module. `GpuVerdict` already
/// establishes the rule for this file — the reason a run happened the way it did
/// is *returned as data* so the caller cannot forget to render it — and a corpus
/// figure needs the same treatment for the same reason. A row count that reads
/// as measured-on-real-data when it is not is the same class of defect as a
/// silent CPU downgrade, and it is harder to notice.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CorpusOrigin {
/// Rows generated in the right shape and scale, embedding nothing.
/// `generator` names the function, so the figure is reproducible and the
/// word SYNTHETIC is not a claim anyone has to take on trust.
Synthetic {
/// e.g. `"gpu::tests::fixture(n, 768, 64, 0.02)"`.
generator: &'static str,
/// The PRNG seed, so a re-run produces the same rows.
seed: u64,
},
/// Rows that are embeddings of real content, named by where it came from.
/// **Nothing in this repo constructs this variant yet** — that is the
/// honest state of the world, and it is expressed in the type rather than in
/// a comment somebody can stop reading.
Real {
/// What was embedded, e.g. a warehouse workspace + snapshot id.
source: String,
/// Which model produced the vectors.
model: String,
},
}
impl CorpusOrigin {
/// The prefix every rendered figure carries.
pub fn label(&self) -> String {
match self {
CorpusOrigin::Synthetic { generator, seed } => {
format!("SYNTHETIC corpus ({generator}, seed {seed:#x}) — embeddings of nothing")
}
CorpusOrigin::Real { source, model } => {
format!("REAL corpus ({source}, model {model})")
}
}
}
/// `true` only for rows that are embeddings of real content.
pub fn is_real(&self) -> bool {
matches!(self, CorpusOrigin::Real { .. })
}
}
/// One throughput figure, inseparable from its corpus provenance and its
/// verdict. Constructing this is the only supported way to print a rate from
/// this module.
#[derive(Debug, Clone, PartialEq)]
pub struct Throughput {
pub rows: usize,
pub dim: usize,
pub queries: usize,
pub secs: f64,
pub corpus: CorpusOrigin,
pub verdict: GpuVerdict,
}
impl Throughput {
/// Pair-comparisons per second — `rows × queries / secs`.
pub fn gpair_per_s(&self) -> f64 {
if self.secs <= 0.0 {
return 0.0;
}
(self.rows as f64 * self.queries as f64) / self.secs / 1e9
}
/// The one line a caller prints. The corpus label comes FIRST and is not
/// optional — `a_synthetic_figure_cannot_be_printed_without_saying_so` is the
/// guard, and it is written to go red if the label is ever moved, shortened
/// away, or made conditional.
pub fn line(&self) -> String {
format!(
"[{}] {} rows × {} queries, dim {} in {:.3} s = {:.2} Gpair/s ({:.2} ms/query). {}",
self.corpus.label(),
self.rows,
self.queries,
self.dim,
self.secs,
self.gpair_per_s(),
self.secs * 1000.0 / self.queries.max(1) as f64,
self.verdict.line()
)
}
}
/// **The real deps corpus, counted — not embedded.**
///
/// MEASURED 2026-08-26 by running [`crate::vector::chunk::chunk_file`] at
/// `ChunkOptions::default()` over every `.rs` file of every crate source
/// unpacked under `~/.cargo/registry/src` on oden. This is how many embedding
/// rows the dependency corpus WOULD have. It is a real count of real files, and
/// it is still not an embedding corpus: nothing has embedded them.
pub const DEPS_CORPUS_MEASURED: &str = "\
5 702 262 chunks (rows) — MEASURED 2026-08-26 on oden by running the real \
vector::chunk::chunk_file at ChunkOptions::default() (max_lines=120, overlap=16) \
over all 88 407 .rs files (1 405 266 470 bytes) of the 2 912 crate sources \
unpacked under ~/.cargo/registry/src; 0 unreadable, 217 s. These rows are \
COUNTED, NOT embedded — no model was run and no vectors exist. At dim 768 they \
would be 16.3 GiB as f32 and 4.1 GiB as the i8 corpus this module scores.";
/// What the measured deps size does to the 2 000 000-row figure above.
///
/// 5 702 262 is **2.85×** the largest synthetic point, so the synthetic table
/// does not overstate the workload — it *understates* it. Extrapolating the
/// flat scan linearly (which is what a flat scan does): ~5.0 s for a
/// 1024-query batch, ~4.9 ms/query. Still the millisecond class, and now that
/// claim rests on a counted corpus rather than on an assumed one. The i8 corpus
/// fits in a 4090's 24 GB; the f32 one, at 16.3 GiB, does not leave room to
/// work alongside it — which is a reason the q8 space is the one that gets
/// scored, and it was not previously written down anywhere.
pub const DEPS_CORPUS_VS_FIXTURE: &str = "\
the real deps corpus is 5 702 262 rows = 2.85x the 2 000 000-row synthetic \
point, so the fixture UNDERSTATES the workload; linear extrapolation of the \
flat scan gives ~5.0 s / 1024-query batch, ~4.9 ms per query";
/// An opened compute device plus the two compiled pipelines.
pub struct GpuKnn {
device: wgpu::Device,
queue: wgpu::Queue,
score: wgpu::ComputePipeline,
topk: wgpu::ComputePipeline,
layout: wgpu::BindGroupLayout,
verdict: GpuVerdict,
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct Params {
n: u32,
dim_words: u32,
q_tile: u32,
nblocks: u32,
k: u32,
q_base: u32,
cand_base: u32,
_pad: u32,
}
/// Enumerate adapters and pick the one real GPU, or explain — in a sentence a
/// human can act on — why there is none. Never returns a software adapter.
pub fn open() -> std::result::Result<GpuKnn, GpuVerdict> {
let inst = wgpu::Instance::new(wgpu::InstanceDescriptor {
backends: wgpu::Backends::VULKAN,
flags: wgpu::InstanceFlags::default(),
memory_budget_thresholds: Default::default(),
backend_options: Default::default(),
display: None,
});
let mut rejected = Vec::new();
let mut chosen = None;
for a in pollster::block_on(inst.enumerate_adapters(wgpu::Backends::VULKAN)) {
let info = a.get_info();
let tag = format!("{:?}/{:?} {}", info.backend, info.device_type, info.name);
if info.device_type == wgpu::DeviceType::DiscreteGpu && !is_software(&info.name) {
chosen = Some(a);
break;
}
rejected.push(tag);
}
let Some(adapter) = chosen else {
return Err(GpuVerdict::Cpu {
reason: "no DISCRETE, non-software Vulkan adapter is present".into(),
rejected,
});
};
let info = adapter.get_info();
let lim = adapter.limits();
let (device, queue) = match pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
label: Some("nornir-gpu-knn"),
required_features: wgpu::Features::empty(),
// The corpus is one storage buffer; the 128 MiB WebGPU default would cap
// it at ~170k rows at dim 768. Ask for what the adapter actually offers.
required_limits: wgpu::Limits {
max_storage_buffer_binding_size: lim.max_storage_buffer_binding_size,
max_buffer_size: lim.max_buffer_size,
..wgpu::Limits::default()
},
memory_hints: wgpu::MemoryHints::Performance,
experimental_features: wgpu::ExperimentalFeatures::disabled(),
trace: wgpu::Trace::Off,
})) {
Ok(v) => v,
Err(e) => {
return Err(GpuVerdict::Cpu {
reason: format!("{} refused a compute device: {e}", info.name),
rejected,
});
}
};
let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("knn"),
source: wgpu::ShaderSource::Wgsl(include_str!("knn.wgsl").into()),
});
let storage = |b: u32, read_only: bool| wgpu::BindGroupLayoutEntry {
binding: b,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
};
let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("knn"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
storage(1, true),
storage(2, true),
storage(3, false),
storage(4, false),
storage(5, false),
],
});
let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("knn"),
bind_group_layouts: &[Some(&layout)],
immediate_size: 0,
});
let mk = |ep: &'static str| {
device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(ep),
layout: Some(&pl),
module: &module,
entry_point: Some(ep),
compilation_options: Default::default(),
cache: None,
})
};
Ok(GpuKnn {
score: mk("score"),
topk: mk("topk"),
layout,
verdict: GpuVerdict::Gpu {
adapter: info.name.clone(),
backend: format!("{:?}", info.backend),
},
device,
queue,
})
}
/// Software rasterisers that answer the Vulkan API without being a GPU.
fn is_software(name: &str) -> bool {
let n = name.to_ascii_lowercase();
n.contains("llvmpipe") || n.contains("lavapipe") || n.contains("swiftshader")
}
/// Pack a row-major `i8` matrix into `u32` words, four lanes per word, in the
/// little-endian lane order `dot4I8Packed` reads.
pub fn pack(rows: &[i8], dim: usize) -> Vec<u32> {
assert_eq!(dim % 4, 0, "dim {dim} must be a multiple of 4");
rows.chunks_exact(4)
.map(|c| {
(c[0] as u8 as u32)
| ((c[1] as u8 as u32) << 8)
| ((c[2] as u8 as u32) << 16)
| ((c[3] as u8 as u32) << 24)
})
.collect()
}
/// The corpus in **word-transposed** order, `out[w · n + row]`.
///
/// This is not a micro-optimisation, it is 3.6× of the runtime. MEASURED at
/// n = 79 000, dim = 768, 1024 queries on a 4090: row-major 0.256 s, transposed
/// 0.071 s. Row-major makes a 32-lane warp touch 32 rows that are 768 B apart —
/// one cache line served per lane. Transposed, the same warp's loads at a fixed
/// word are consecutive `u32`, so one burst serves the warp.
pub fn pack_transposed(rows: &[i8], n: usize, dim: usize) -> Vec<u32> {
let dw = dim / 4;
let flat = pack(rows, dim);
let mut out = vec![0u32; n * dw];
for r in 0..n {
for w in 0..dw {
out[w * n + r] = flat[r * dw + w];
}
}
out
}
impl GpuKnn {
/// What this instance is running on. Print it; never assume it.
pub fn verdict(&self) -> &GpuVerdict {
&self.verdict
}
/// Batched exact-int8 kNN: `out[i]` is the top-`k` of `queries[i]` over every
/// row in `index`, as `(id, cosine)`.
///
/// The returned scores are the **int8** cosines, bit-identical to
/// [`znippy_zoomies::vann::score_i8_batch`] + [`znippy_zoomies::vann::top_k`]
/// on the same corpus — including the `(score DESC, id ASC)` tie-break, which
/// the shader reproduces so that the block-local lists merge into the same
/// global list rather than a plausible neighbour of it.
///
/// # Errors
/// If `queries` is not a whole number of `dim`-vectors, if `k` exceeds
/// [`K_MAX`], or if the corpus exceeds the adapter's storage-buffer limit.
pub fn knn(&self, index: &VectorIndex, queries: &[f32], k: usize) -> Result<Vec<Vec<(u64, f32)>>> {
self.knn_inner(index, queries, k, std::env::var_os(BREAK_ENV).is_some())
}
/// [`Self::knn`] with the sabotage decided by the CALLER rather than by the
/// process environment.
///
/// The red proof needs one run with a corrupt corpus beside one run without,
/// in the same test binary. Driving that through an env var would have the
/// two tests racing on process-global state — the corrupting test could flip
/// the switch under the oracle test running on another thread, and the suite
/// would be green or red by scheduling accident. A parameter cannot race.
/// [`BREAK_ENV`] stays, because the shipping `knn` reading it is what makes
/// the red proof reachable in the REAL path and not only in a test-only twin.
#[doc(hidden)]
pub fn knn_maybe_corrupted(
&self,
index: &VectorIndex,
queries: &[f32],
k: usize,
corrupt: bool,
) -> Result<Vec<Vec<(u64, f32)>>> {
self.knn_inner(index, queries, k, corrupt)
}
fn knn_inner(
&self,
index: &VectorIndex,
queries: &[f32],
k: usize,
corrupt: bool,
) -> Result<Vec<Vec<(u64, f32)>>> {
let dim = index.dim();
let n = index.len();
if k == 0 || k > K_MAX {
bail!("k must be in 1..={K_MAX}, got {k}");
}
if dim % 4 != 0 {
bail!("GPU kNN needs dim % 4 == 0 (dot4I8Packed reads four lanes per word); dim = {dim}");
}
if queries.len() % dim != 0 {
bail!("queries len {} is not a multiple of dim {dim}", queries.len());
}
if n == 0 {
bail!("the index is EMPTY — that is a missing producer, not zero neighbours");
}
let nq = queries.len() / dim;
let dw = dim / 4;
// Reuse the index's own memoized i8 matrix. Quantizing a second time here
// would be a second writer of the q8 space and the first thing to drift
// (LAW 5); the whole oracle rests on both sides reading the SAME bytes.
let (rows_i8, _row_sums) = index.quantized();
let mut packed = pack_transposed(&rows_i8, n, dim);
if corrupt {
// Deliberate sabotage, for the red proof. Zeroing ONE of `dw` words per
// row is a sub-percent corruption at dim 768 — small enough that a
// tolerance-based check would wave it through, which is exactly the point.
eprintln!("⚠ the GPU corpus is being corrupted ON PURPOSE (red proof).");
for r in 0..n {
packed[(dw - 1) * n + r] = 0;
}
}
// Quantize the queries exactly as `score_i8_batch` does: L2-normalize,
// then round-to-nearest at scale Q. Same functions, not a reimplementation.
let mut q_i8: Vec<i8> = Vec::with_capacity(nq.next_multiple_of(QUERY_TILE) * dim);
for qi in 0..nq {
let qn = znippy_zoomies::vann::normalized(&queries[qi * dim..(qi + 1) * dim]);
znippy_zoomies::vann::quantize_i8(&qn, &mut q_i8);
}
let ntiles = nq.div_ceil(QUERY_TILE);
q_i8.resize(ntiles * QUERY_TILE * dim, 0);
let nblocks = NBLOCKS.min(n.div_ceil(256)).max(1);
let cand_per_tile = QUERY_TILE * nblocks * k;
let cand_n = ntiles * cand_per_tile;
let out = self.dispatch(&packed, &q_i8, n, dw, ntiles, nblocks, k, cand_n)?;
// Merge each query's block-local lists with the SAME `top_k` the CPU path
// uses, so the tie-break is one implementation rather than two that agree.
let (sc, rw) = out;
let mut res = Vec::with_capacity(nq);
for qi in 0..nq {
let tile = qi / QUERY_TILE;
let ql = qi % QUERY_TILE;
let mut merged: Vec<(u64, f32)> = Vec::with_capacity(nblocks * k);
for b in 0..nblocks {
for j in 0..k {
let idx = tile * cand_per_tile + (ql * nblocks + b) * k + j;
if rw[idx] == u32::MAX {
continue;
}
merged.push((rw[idx] as u64, sc[idx] as f32 / (Q * Q)));
}
}
znippy_zoomies::vann::top_k(&mut merged, k);
res.push(merged);
}
Ok(res)
}
/// Record every tile's dispatch pair into ONE encoder and close the batch
/// with ONE readback. Mapping per tile costs a full round trip each; at
/// n = 79 000 that was 60× the arithmetic.
#[allow(clippy::too_many_arguments)]
fn dispatch(
&self,
packed: &[u32],
q_i8: &[i8],
n: usize,
dw: usize,
ntiles: usize,
nblocks: usize,
k: usize,
cand_n: usize,
) -> Result<(Vec<i32>, Vec<u32>)> {
use wgpu::util::DeviceExt;
let corpus_bytes = (packed.len() * 4) as u64;
let limit = self.device.limits().max_storage_buffer_binding_size as u64;
if corpus_bytes > limit {
bail!(
"corpus is {corpus_bytes} B but this adapter binds at most {limit} B per storage \
buffer — split the corpus or score on CPU; it must NOT be silently truncated"
);
}
let qpk = pack(q_i8, dw * 4);
let rows_buf = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("corpus"),
contents: bytemuck::cast_slice(packed),
usage: wgpu::BufferUsages::STORAGE,
});
let q_buf = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("queries"),
contents: bytemuck::cast_slice(&qpk),
usage: wgpu::BufferUsages::STORAGE,
});
let scores = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("scores"),
size: (QUERY_TILE * n * 4) as u64,
usage: wgpu::BufferUsages::STORAGE,
mapped_at_creation: false,
});
let mk = |l: &'static str| {
self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(l),
size: (cand_n * 4) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
})
};
let cs = mk("cand_score");
let cr = mk("cand_row");
let read = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("readback"),
size: (cand_n * 8) as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut bgs = Vec::with_capacity(ntiles);
for tile in 0..ntiles {
let p = Params {
n: n as u32,
dim_words: dw as u32,
q_tile: QUERY_TILE as u32,
nblocks: nblocks as u32,
k: k as u32,
q_base: (tile * QUERY_TILE * dw) as u32,
cand_base: (tile * QUERY_TILE * nblocks * k) as u32,
_pad: 0,
};
let pbuf = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("params"),
contents: bytemuck::bytes_of(&p),
usage: wgpu::BufferUsages::UNIFORM,
});
bgs.push(self.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: None,
layout: &self.layout,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: pbuf.as_entire_binding() },
wgpu::BindGroupEntry { binding: 1, resource: rows_buf.as_entire_binding() },
wgpu::BindGroupEntry { binding: 2, resource: q_buf.as_entire_binding() },
wgpu::BindGroupEntry { binding: 3, resource: scores.as_entire_binding() },
wgpu::BindGroupEntry { binding: 4, resource: cs.as_entire_binding() },
wgpu::BindGroupEntry { binding: 5, resource: cr.as_entire_binding() },
],
}));
}
let mut enc = self.device.create_command_encoder(&Default::default());
{
let mut cp = enc.begin_compute_pass(&Default::default());
let wg = n.div_ceil(256).clamp(1, 8192) as u32;
for bg in &bgs {
cp.set_bind_group(0, bg, &[]);
cp.set_pipeline(&self.score);
cp.dispatch_workgroups(wg, 1, 1);
cp.set_pipeline(&self.topk);
cp.dispatch_workgroups((QUERY_TILE * nblocks) as u32, 1, 1);
}
}
enc.copy_buffer_to_buffer(&cs, 0, &read, 0, (cand_n * 4) as u64);
enc.copy_buffer_to_buffer(&cr, 0, &read, (cand_n * 4) as u64, (cand_n * 4) as u64);
self.queue.submit([enc.finish()]);
let slice = read.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
self.device
.poll(wgpu::PollType::Wait { submission_index: None, timeout: None })
.context("polling the GPU queue")?;
rx.recv()
.context("the buffer-map callback never fired")?
.context("mapping the readback buffer")?;
// wgpu 30: the mapped range is a `Result` (the map can have been lost)
let data = slice.get_mapped_range().map_err(|e| anyhow::anyhow!("reading the mapped readback buffer: {e}"))?;
let sc: Vec<i32> = bytemuck::cast_slice(&data[..cand_n * 4]).to_vec();
let rw: Vec<u32> = bytemuck::cast_slice(&data[cand_n * 4..]).to_vec();
drop(data);
read.unmap();
Ok((sc, rw))
}
}
/// The CPU reference: the same top-`k`, in the same q8 space, through
/// `vann`'s own kernels. This is the oracle [`GpuKnn::knn`] is measured against
/// — and the honest CPU path when [`open`] refuses the box.
pub fn cpu_int8_knn(index: &VectorIndex, queries: &[f32], k: usize) -> Vec<Vec<(u64, f32)>> {
let dim = index.dim();
let (rows, sums) = index.quantized();
(0..queries.len() / dim)
.map(|qi| {
let s = znippy_zoomies::vann::score_i8_batch(
&queries[qi * dim..(qi + 1) * dim],
&rows,
dim,
&sums,
);
let mut v: Vec<(u64, f32)> =
s.iter().enumerate().map(|(i, &x)| (i as u64, x)).collect();
znippy_zoomies::vann::top_k(&mut v, k);
v
})
.collect()
}
/// Score on the GPU when there is one, on the CPU when there is not — and say
/// which, always. The verdict is returned rather than logged so the caller
/// cannot forget to render it.
pub fn knn_loud(
index: &VectorIndex,
queries: &[f32],
k: usize,
) -> (Vec<Vec<(u64, f32)>>, GpuVerdict) {
match open() {
Ok(g) => match g.knn(index, queries, k) {
Ok(v) => {
let verdict = g.verdict().clone();
(v, verdict)
}
Err(e) => (
cpu_int8_knn(index, queries, k),
GpuVerdict::Cpu {
reason: format!("the GPU dispatch FAILED: {e:#}"),
rejected: Vec::new(),
},
),
},
Err(v) => (cpu_int8_knn(index, queries, k), v),
}
}
#[cfg(test)]
mod tests {
use super::*;
/// A deterministic clustered fixture: `nclust` mutually near-orthogonal
/// centroids (random unit vectors in 768 dimensions are almost orthogonal),
/// each surrounded by a tight noise ball. Row `i` belongs to blob `i % nclust`
/// **by construction**, which is what makes the blob-identity assertion an
/// assertion and not a restatement of the output.
fn fixture(n: usize, dim: usize, nclust: usize, noise: f32) -> Vec<f32> {
let mut s = 0xC0FFEEu64;
let mut next = move || {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
((s >> 33) as f32 / (1u64 << 31) as f32) - 1.0
};
let cents: Vec<Vec<f32>> = (0..nclust)
.map(|_| {
znippy_zoomies::vann::normalized(
&(0..dim).map(|_| next()).collect::<Vec<f32>>(),
)
})
.collect();
let mut out = Vec::with_capacity(n * dim);
for i in 0..n {
let c = ¢s[i % nclust];
let v: Vec<f32> = (0..dim).map(|j| c[j] + noise * next()).collect();
out.extend_from_slice(&znippy_zoomies::vann::normalized(&v));
}
out
}
fn index_of(flat: &[f32], n: usize, dim: usize) -> VectorIndex {
let mut idx = VectorIndex::new(dim).unwrap();
idx.add(flat, &(0..n as u64).collect::<Vec<u64>>()).unwrap();
idx
}
/// A refused box must still be a LOUD box. This runs everywhere, GPU or not.
#[test]
fn a_cpu_verdict_names_what_it_refused_and_never_reads_as_success() {
let v = GpuVerdict::Cpu {
reason: "no DISCRETE, non-software Vulkan adapter is present".into(),
rejected: vec!["Vulkan/Cpu llvmpipe (LLVM 21.1.8, 256 bits)".into()],
};
assert!(!v.is_gpu());
let line = v.line();
assert!(line.starts_with('⚠'), "a downgrade must be visually loud: {line}");
assert!(line.contains("llvmpipe"), "must NAME the refused adapter: {line}");
assert!(line.contains("CPU"), "must say where the work went: {line}");
// The failure this guards: a Cpu verdict whose line is indistinguishable
// from a Gpu one.
let g = GpuVerdict::Gpu { adapter: "NVIDIA GeForce RTX 4090".into(), backend: "Vulkan".into() };
assert_ne!(g.line(), line);
assert!(!g.line().starts_with('⚠'));
}
#[test]
fn software_adapters_are_recognised_by_name() {
for n in ["llvmpipe (LLVM 21.1.8, 256 bits)", "SwiftShader Device", "lavapipe"] {
assert!(is_software(n), "{n} must be refused");
}
assert!(!is_software("NVIDIA GeForce RTX 4090"));
assert!(!is_software("AMD Radeon RX 7900 XTX"));
}
/// The transpose is the layout the shader indexes; if it is wrong every score
/// is wrong. Checked against the row-major packing it permutes, elementwise.
#[test]
fn the_transposed_packing_is_a_permutation_of_the_row_major_one() {
let dim = 8;
let n = 5;
let rows: Vec<i8> = (0..(n * dim)).map(|i| (i as i32 % 127 - 63) as i8).collect();
let flat = pack(&rows, dim);
let t = pack_transposed(&rows, n, dim);
let dw = dim / 4;
assert_eq!(flat.len(), t.len());
for r in 0..n {
for w in 0..dw {
assert_eq!(t[w * n + r], flat[r * dw + w], "row {r} word {w}");
}
}
// RED-guard: a transpose that just copied would pass an "is a permutation"
// check on square input, so assert the layouts genuinely DIFFER here.
assert_ne!(flat, t, "n != dim, so the two layouts must not coincide");
}
/// THE ORACLE. Same corpus, same q8 space, two engines — and because the
/// dot is an exact i32 the answers must be equal BYTE FOR BYTE, not close.
///
/// Skipped only when the box has no GPU, and then LOUDLY: a silently skipped
/// oracle is a green nobody has seen red.
#[test]
fn gpu_knn_matches_cpu_int8_bit_for_bit() {
let (n, dim, k) = (6_000usize, 128usize, 8usize);
let flat = fixture(n, dim, 32, 0.05);
let idx = index_of(&flat, n, dim);
let nq = 64;
let q = &flat[..nq * dim];
let g = match open() {
Ok(g) => g,
Err(v) => {
eprintln!("SKIPPED gpu_knn_matches_cpu_int8_bit_for_bit — {}", v.line());
return;
}
};
let gpu = g.knn(&idx, q, k).expect("gpu knn");
let cpu = cpu_int8_knn(&idx, q, k);
assert_eq!(gpu.len(), cpu.len());
for qi in 0..nq {
assert_eq!(
gpu[qi], cpu[qi],
"query {qi}: the GPU and the CPU disagree in a space where the dot is an \
EXACT i32 — one of the two kernels is wrong"
);
}
}
/// THE RED PROOF for the oracle above. With the corpus deliberately corrupted
/// by one word in 32 (a 3% perturbation of the dot), the oracle must FAIL.
/// If this test cannot make it fail, the oracle is blind and its green means
/// nothing.
#[test]
fn a_corrupted_gpu_corpus_fails_the_oracle() {
let (n, dim, k) = (6_000usize, 128usize, 8usize);
let flat = fixture(n, dim, 32, 0.05);
let idx = index_of(&flat, n, dim);
let nq = 64;
let q = &flat[..nq * dim];
let g = match open() {
Ok(g) => g,
Err(v) => {
eprintln!("SKIPPED a_corrupted_gpu_corpus_fails_the_oracle — {}", v.line());
return;
}
};
let cpu = cpu_int8_knn(&idx, q, k);
let broken = g
.knn_maybe_corrupted(&idx, q, k, true)
.expect("gpu knn (deliberately broken)");
// The SAME call with the sabotage off must agree, or "differing" below
// would prove nothing about the corruption -- it could be proving the
// kernel is simply wrong.
let clean = g.knn_maybe_corrupted(&idx, q, k, false).expect("gpu knn (clean)");
assert_eq!(clean, cpu, "the UNcorrupted run must match, or this test is not isolating the corruption");
let differing = (0..nq).filter(|&qi| broken[qi] != cpu[qi]).count();
assert!(
differing > 0,
"a corrupted corpus produced the SAME top-{k} for all {nq} queries — the oracle \
cannot see a wrong kernel and is therefore worthless"
);
}
/// The mimir_clones-shaped invariant, in this module's terms: on WELL-SEPARATED
/// blobs the GPU must assign every neighbour to the same blob the exact f32
/// CPU search does. That is what a clustering fold consumes; the tie order
/// inside one blob is not.
#[test]
fn well_separated_blobs_get_identical_membership_from_gpu_and_exact_cpu() {
let (n, dim, k, nclust) = (4_096usize, 128usize, 8usize, 16usize);
let flat = fixture(n, dim, nclust, 0.03);
let idx = index_of(&flat, n, dim);
let nq = 64;
let q = &flat[..nq * dim];
let g = match open() {
Ok(g) => g,
Err(v) => {
eprintln!("SKIPPED well_separated_blobs… — {}", v.line());
return;
}
};
let gpu = g.knn(&idx, q, k).expect("gpu knn");
let exact = idx.search_batch(q, k);
let blob = |r: u64| (r as usize) % nclust;
for qi in 0..nq {
let want = blob(qi as u64);
let gpu_off = gpu[qi].iter().filter(|x| blob(x.0) != want).count();
let cpu_off = exact[qi].iter().filter(|x| blob(x.0) != want).count();
assert_eq!(
gpu_off, cpu_off,
"query {qi}: GPU put {gpu_off} neighbours outside blob {want}, exact-f32 put \
{cpu_off} — the two methods do NOT cluster identically on separated data"
);
assert_eq!(gpu_off, 0, "query {qi}: separated blobs must not leak neighbours");
}
}
/// An empty index is a MISSING PRODUCER, never "no neighbours". Same law the
/// lenses' MISSING FEEDS note encodes.
#[test]
fn an_empty_index_is_an_error_not_an_empty_result() {
let idx = VectorIndex::new(64).unwrap();
let Ok(g) = open() else { return };
let e = g.knn(&idx, &vec![0.1f32; 64], 4).unwrap_err();
assert!(format!("{e:#}").contains("EMPTY"), "{e:#}");
}
// ---- corpus provenance: a figure may not be printed without its origin ----
fn synthetic_figure() -> Throughput {
Throughput {
rows: 2_000_000,
dim: 768,
queries: 1024,
secs: 1.744,
corpus: CorpusOrigin::Synthetic { generator: "gpu::tests::fixture", seed: 0xC0FFEE },
verdict: GpuVerdict::Gpu {
adapter: "NVIDIA GeForce RTX 4090".into(),
backend: "Vulkan".into(),
},
}
}
/// **The honesty guard.** Every rendered figure states SYNTHETIC, names the
/// generator and the seed, and does it BEFORE the number — the same
/// discipline `GpuVerdict::line` applies to a CPU downgrade. A reader who
/// stops after the first bracket has already been told.
#[test]
fn a_synthetic_figure_cannot_be_printed_without_saying_so() {
let t = synthetic_figure();
let l = t.line();
assert!(l.starts_with("[SYNTHETIC corpus"), "provenance comes FIRST: {l}");
assert!(l.contains("gpu::tests::fixture"), "and names the generator: {l}");
assert!(l.contains("0xc0ffee"), "and the seed, so it is reproducible: {l}");
assert!(l.contains("embeddings of nothing"), "in words, not a bare tag: {l}");
assert!(!t.corpus.is_real());
// The number is still there and still right.
assert!(l.contains("2000000 rows"), "{l}");
assert!((t.gpair_per_s() - 1.174).abs() < 0.01, "{}", t.gpair_per_s());
}
/// The red half: the guard must be able to tell the two apart, or it is not
/// checking anything. Nothing in the shipping tree constructs `Real` — this
/// is the only place it exists, and that is the point.
#[test]
fn a_real_corpus_would_read_differently_which_is_how_we_know_the_label_is_load_bearing() {
let t = Throughput {
corpus: CorpusOrigin::Real {
source: "warehouse nordisk@snapshot-42".into(),
model: "jina-embeddings-v2-base-code".into(),
},
..synthetic_figure()
};
let l = t.line();
assert!(l.starts_with("[REAL corpus"), "{l}");
assert!(!l.contains("SYNTHETIC"), "{l}");
assert!(t.corpus.is_real());
}
/// The recorded deps-corpus size must be a MEASUREMENT with its method
/// attached, never a round number somebody remembered.
#[test]
fn the_deps_corpus_figure_states_its_method_and_that_it_is_not_embedded() {
let d = DEPS_CORPUS_MEASURED;
assert!(d.contains("chunk_file"), "the tool that counted it: {d}");
assert!(d.contains("registry/src"), "what it counted: {d}");
assert!(d.contains("NOT embedded"), "and what it is not: {d}");
}
}
// ----- the `knn_neighbors` producer ------------------------------------------
/// One row of the `knn_neighbors` warehouse table
/// ([`crate::warehouse::iceberg_schema::knn_neighbors`]).
///
/// This is the shape that crosses the boundary from "GPU compute" to "a table a
/// pure lens fold can read". Nothing downstream of here knows a GPU exists.
#[derive(Debug, Clone, PartialEq)]
pub struct NeighborRow {
/// Repo the source symbol belongs to.
pub repo: String,
/// Embedding space these neighbours live in (`model_profile`).
pub model_profile: String,
/// `k` the list was built with — the fold parameter, carried so a consumer
/// cannot silently read a `k = 8` answer out of a `k = 32` list.
pub k: i32,
/// Distance metric; `"cosine_i8"` for this engine.
pub metric: String,
/// FQN of the query symbol.
pub src_symbol: String,
/// FQN of the neighbour.
pub dst_symbol: String,
/// 0-based position in `src_symbol`'s list.
pub rank: i32,
/// The int8 cosine — bit-identical on GPU and CPU.
pub cosine: f32,
/// `true` when `src` is also inside `dst`'s own top-`k`.
pub mutual: bool,
/// `gpu:<adapter>` or `cpu:<simd>`; the durable record of where the
/// arithmetic happened.
pub producer: String,
}
/// Turn a batched kNN result into `knn_neighbors` rows, computing `mutual` ONCE.
///
/// `symbols[i]` is the FQN of index row `i`; `queries_of[q]` is the row index
/// query `q` was built from. `neighbors[q]` is that query's `(row_id, cosine)`
/// list.
///
/// # Why `mutual` is computed here and not by the consumer
///
/// `mimir_clones`' mutual-kNN retain step needs "is `a` in `b`'s top-k AND `b`
/// in `a`'s top-k". That predicate needs the WHOLE list; a consumer reading one
/// symbol's rows cannot evaluate it, and a consumer reading all of them would be
/// re-deriving, per lens, something the producer already knew. One writer, not
/// one per reader (LAW 5).
///
/// A symbol is never its own neighbour: the self-hit (cosine 1.0 by
/// construction, since every query here IS a corpus row) is dropped, because a
/// self-edge is not a similarity finding and would silently inflate every
/// downstream degree count by exactly one.
pub fn neighbor_rows(
repo: &str,
model_profile: &str,
producer: &str,
k: usize,
symbols: &[String],
queries_of: &[usize],
neighbors: &[Vec<(u64, f32)>],
) -> Result<Vec<NeighborRow>> {
if queries_of.len() != neighbors.len() {
bail!(
"queries_of has {} entries but neighbors has {} — these must be the same query set",
queries_of.len(),
neighbors.len()
);
}
// row -> the set of rows it lists, so `mutual` is a lookup rather than a scan.
let mut listed: std::collections::HashMap<usize, std::collections::HashSet<u64>> =
std::collections::HashMap::with_capacity(queries_of.len());
for (qi, &src) in queries_of.iter().enumerate() {
listed.insert(
src,
neighbors[qi]
.iter()
.map(|&(id, _)| id)
.filter(|&id| id as usize != src)
.collect(),
);
}
let mut out = Vec::new();
for (qi, &src) in queries_of.iter().enumerate() {
let src_name = symbols
.get(src)
.with_context(|| format!("row {src} has no symbol name"))?;
let mut rank = 0i32;
for &(dst, cos) in &neighbors[qi] {
let d = dst as usize;
if d == src {
continue; // never a self-edge
}
let dst_name = symbols
.get(d)
.with_context(|| format!("row {d} has no symbol name"))?;
let mutual = listed
.get(&d)
.is_some_and(|s| s.contains(&(src as u64)));
out.push(NeighborRow {
repo: repo.to_string(),
model_profile: model_profile.to_string(),
k: k as i32,
metric: "cosine_i8".to_string(),
src_symbol: src_name.clone(),
dst_symbol: dst_name.clone(),
rank,
cosine: cos,
mutual,
producer: producer.to_string(),
});
rank += 1;
}
}
Ok(out)
}
#[cfg(test)]
mod producer_tests {
use super::*;
fn names(n: usize) -> Vec<String> {
(0..n).map(|i| format!("crate::f{i}")).collect()
}
/// `mutual` must be TRUE only when the listing is reciprocal. The failure this
/// guards is the easy one: returning `true` for every edge, which every
/// "does it produce rows" test passes.
#[test]
fn mutual_is_reciprocal_and_not_a_constant() {
let syms = names(4);
// 0 lists 1; 1 lists 0 -> reciprocal.
// 2 lists 3; 3 lists 0 -> NOT reciprocal either way.
let neighbors = vec![
vec![(0u64, 1.0f32), (1, 0.9)],
vec![(1u64, 1.0f32), (0, 0.9)],
vec![(2u64, 1.0f32), (3, 0.8)],
vec![(3u64, 1.0f32), (0, 0.7)],
];
let rows = neighbor_rows("r", "jina-768", "gpu:test", 2, &syms, &[0, 1, 2, 3], &neighbors)
.unwrap();
let find = |s: &str, d: &str| {
rows.iter()
.find(|r| r.src_symbol == s && r.dst_symbol == d)
.unwrap_or_else(|| panic!("no row {s} -> {d}"))
};
assert!(find("crate::f0", "crate::f1").mutual);
assert!(find("crate::f1", "crate::f0").mutual);
assert!(!find("crate::f2", "crate::f3").mutual, "3 does not list 2");
assert!(!find("crate::f3", "crate::f0").mutual, "0 does not list 3");
// The anti-constant assertion: both values must actually occur.
assert!(rows.iter().any(|r| r.mutual));
assert!(rows.iter().any(|r| !r.mutual));
}
/// The self-hit is dropped, and dropping it must not leave a gap in `rank`.
#[test]
fn the_self_edge_is_dropped_and_ranks_stay_dense_from_zero() {
let syms = names(3);
let neighbors = vec![vec![(0u64, 1.0f32), (1, 0.9), (2, 0.8)]];
let rows =
neighbor_rows("r", "m", "cpu:avx2", 3, &syms, &[0], &neighbors).unwrap();
assert_eq!(rows.len(), 2, "the self-hit must be gone: {rows:?}");
assert!(rows.iter().all(|r| r.src_symbol != r.dst_symbol));
assert_eq!(rows[0].rank, 0);
assert_eq!(rows[1].rank, 1, "ranks must be dense, not 0 then 2");
}
/// A mismatched query set is a wiring bug, and must be an ERROR — never a
/// short table that reads like "these symbols have no neighbours".
#[test]
fn a_mismatched_query_set_is_an_error_not_a_short_table() {
let syms = names(2);
let e = neighbor_rows("r", "m", "p", 1, &syms, &[0, 1], &[vec![(1u64, 0.5f32)]])
.unwrap_err();
assert!(format!("{e:#}").contains("same query set"), "{e:#}");
}
/// The producer string is what an operator reads to know a GPU ran. It must
/// reach every row, not just the first.
#[test]
fn every_row_carries_the_producer_and_the_k_it_was_built_with() {
let syms = names(3);
let neighbors = vec![vec![(1u64, 0.9f32), (2, 0.8)], vec![(0u64, 0.9f32), (2, 0.7)]];
let rows = neighbor_rows(
"nornir",
"jina-v2-base-code",
"gpu:NVIDIA GeForce RTX 4090",
2,
&syms,
&[0, 1],
&neighbors,
)
.unwrap();
assert!(!rows.is_empty());
assert!(rows.iter().all(|r| r.producer.starts_with("gpu:")));
assert!(rows.iter().all(|r| r.k == 2));
assert!(rows.iter().all(|r| r.metric == "cosine_i8"));
assert!(rows.iter().all(|r| r.model_profile == "jina-v2-base-code"));
}
}