Your browser sucks, JavaScript disabled! Gunnar's console needs JavaScript and WebAssembly. Git over SSH on port 2222 does not.

nordisk/nornir

src/vector/mod.rs

//! Vector (semantic) search index — a hand-written **exact-flat** ANN over
//! `f32` vectors, keyed by stable `u64` ids.
//!
//! Design priorities (per `plan.md`): **maximum precision** and speed, 100%
//! Rust, no C/FFI, self-contained airgapped binary.
//!
//! - **Exact, not approximate.** Every query scores against every stored
//!   vector, so recall is 100% — no quantization loss, no graph-traversal
//!   miss. The embedding model already spends storage on precision
//!   (`jina-v2-base-code`, 768-dim); the index does not throw that away.
//! - **Cosine similarity.** Vectors are L2-normalized on insert and the query
//!   is normalized per search, so the score is a plain dot product (cosine).
//!   Higher score = closer.
//! - **SIMD, runtime-detected.** The per-vector dot product dispatches once
//!   per search to the best kernel the *running* CPU supports: AVX-512F →
//!   AVX2+FMA → scalar. The binary builds and runs everywhere; it just goes
//!   faster where the silicon allows (e.g. AVX-512 on Zen 4).
//! - **int8 quantization + VNNI (G2).** Because stored vectors are
//!   L2-normalized (components in `[-1,1]`), they quantize losslessly-enough to
//!   `i8` (scale 127). The int8 cosine matches the f32 cosine to ~1e-2, runs on
//!   a single AVX-512 **VNNI** `vpdpbusd` (64 int8 MACs/instr vs 16 f32 FMA
//!   lanes), and quarters the bytes/row the (memory-bound) scoring loop streams.
//!   See [`score_i8_batch`] / [`VectorIndex::search_i8`]; both fall back to a
//!   scalar int8 kernel where VNNI is absent. [`bench_kernels`] times every
//!   path and is exercised by `nornir vector bench`.
//! - **Multicore.** For large corpora the scoring loop is split across cores
//!   via scoped threads (no `Arc`, no new dependency), each computing a local
//!   top-k that is merged into the global top-k.
//!
//! Ids map back to warehouse rows (chunk id → `{repo, git_sha, model, file,
//! span, excerpt}`), so the index stays a pure derived artifact — the same
//! shape as the Tantivy full-text index, and snapshot/restore-able the same
//! way. The embedding model that produces the `f32` vectors (Candle, feature
//! `embed-tract` / `embed-ort`) is a separate layer; this module is
//! model-agnostic and
//! only cares about dimensionality.
//!
//! Cargo feature: `vector`.

pub mod chunk;
// The ONE writer of "where does this item start and end" — nested, so a method
// of an `impl` is its own span rather than a slice of an impl-block average.
// Consumed by `chunk` (leaf spans), by the sub-clone detector (enclosing-item
// labels) and by the warehouse's missing `symbol_facts.end_line` (mimir-2).
pub mod spans;
pub mod store;
// One vector per function node of the code room: the chunk is the node's own
// text (cut by the rustc driver's span), the key node id → content_hash → vector.
// What the two vector lenses (`similar_code`, `semantic_clones`) read.
pub mod fnvec;
/// The embed MEMORY CEILING (`--max-memory` / `NORNIR_EMBED_MAX_MEMORY_MB`) and
/// the fan-out plan the backends run under it — gap 3 of
/// INVENTORY-semantic-twins-2026-09-10 §5. Backend-independent, `vector`-gated only.
pub mod memory;

// The interchangeable embedding-model registry (plan item G3). Pure-std, no
// deps — also `include!`d by `build.rs` so the model it fetches matches the one
// the embedder loads. Available whenever `vector` is on (the CLI reports the
// selected model), not just under an embed backend.
pub mod embed_registry;

// Shared embedder support — compiled when either backend is enabled.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
pub mod embed_support;

// Why we are on the CPU when the box wanted a GPU. Compiled with either backend
// — the tract-only build is precisely the shape that needs to confess loudest.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
pub mod fallback;

// Embedder backends (the `store::Embedder` trait is the interface). Both run
// the same jina code ONNX model; pick by Cargo feature.
#[cfg(feature = "embed-ort")]
pub mod cuda; // runtime CUDA-lib + onnxruntime discovery so the ort NVIDIA EP "just works"
#[cfg(feature = "embed-tract")]
pub mod embed; // tract-onnx, CPU, pure Rust
#[cfg(feature = "embed-ort")]
pub mod embed_ort; // ort / ONNX Runtime, GPU (CUDA/ROCm) or CPU
#[cfg(feature = "embed-ort-rocm")]
pub mod rocm; // runtime ROCm-lib discovery + probe for the ort AMD EP (G1)

// GPU-batched kNN over the SAME q8 corpus `VectorIndex::quantized()` hands the
// VNNI kernel — a WebGPU compute pipeline, not an FFI binding. Independent of
// the embed backends: this accelerates SEARCH, they accelerate INFERENCE.
#[cfg(feature = "vector-gpu")]
pub mod gpu;

// Hybrid ranking: the dense score the GPU produced, combined with the SCIP call
// graph's exact relationality. A PURE fold over (cosine, hop count) — no device,
// no warehouse — so alpha/beta are calibrated by measurement on a headless
// runner, and the lens boundary is never crossed.
pub mod hybrid;

// The exact cosine self-join of the function vectors, kept per process and
// brought forward by what changed (the vector lenses' candidate scan).
pub mod pair_index;
/// The tiled exact self-join (the pair index's kernel replicas).
pub mod tiled;

// Static token→vector table (model2vec / potion) — no graph, no inference runtime.
// Deliberately NOT behind its own Cargo feature: it adds no dependency (safetensors is
// hand-parsed; `tokenizers` and `serde_json` are already here), so gating it would only
// create a build shape in which `Backend::StaticTable` names an unreachable impl. It
// rides `any(embed-tract, embed-ort)` because that is exactly where `tokenizers` and
// `embed_support` live.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
pub mod embed_static;

// Runtime backend selector (#9). Active only when BOTH the CPU (tract) and GPU
// (ort) embedders are compiled — the default `cargo install nornir` build — so
// the single binary picks CUDA/ROCm/CPU at runtime. With only one backend
// compiled, `load_embedder` selects it directly (below).
#[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
pub mod select;

// An embedder on another box (`$NORNIR_EMBED_REMOTE`): a GPU worker's forward
// passes for a server without a GPU. The client rides the wire stubs; the
// worker serves them (`nornir vector serve-embed`).
#[cfg(all(any(feature = "embed-tract", feature = "embed-ort"), feature = "wire-client"))]
pub mod remote;

/// Load the default embedder as a trait object, choosing the best available
/// backend **at runtime** (#9): when both backends are compiled (the default
/// build) the [`select`] module probes the box (NVIDIA → AMD → CPU) and loads
/// the right one; otherwise the single compiled backend is loaded directly.
/// Both backends produce vectors with the same `model_profile`, so they
/// interoperate in the warehouse.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
#[allow(clippy::needless_return)] // returns disambiguate the cfg branches
pub fn load_embedder() -> anyhow::Result<Box<dyn store::Embedder>> {
    // A worker named: its GPU, never a local fallback — an unattended bulk embed
    // that silently moved to this box's CPU would run for hours.
    #[cfg(feature = "wire-client")]
    if remote::configured().is_some() {
        return Ok(Box::new(remote::RemoteEmbedder::from_env()?));
    }
    // Default build: both compiled → runtime selection.
    #[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
    {
        return select::load();
    }
    // Single-backend builds. The static-table check comes FIRST in both: a static model
    // has no ONNX graph on disk, so handing it to either graph loader is a guaranteed
    // "load model.onnx: No such file" whose real cause (the selected MODEL, not the
    // backend) is nowhere in the message.
    #[cfg(all(feature = "embed-ort", not(feature = "embed-tract")))]
    {
        if embed_registry::selected_kind_is_static() {
            return Ok(Box::new(embed_static::StaticEmbedder::load()?));
        }
        return Ok(Box::new(embed_ort::OrtEmbedder::load()?));
    }
    // Only tract compiled.
    #[cfg(all(feature = "embed-tract", not(feature = "embed-ort")))]
    {
        if embed_registry::selected_kind_is_static() {
            return Ok(Box::new(embed_static::StaticEmbedder::load()?));
        }
        return Ok(Box::new(embed::JinaEmbedder::load()?));
    }
}

/// Human-readable name of the backend [`load_embedder`] selects. In the default
/// (both-backend) build this **probes the running machine** and reports the
/// runtime choice (CUDA / ROCm / CPU); with a single backend it names it.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
pub fn embedder_backend() -> &'static str {
    #[cfg(feature = "wire-client")]
    if let Some(url) = remote::configured() {
        static LABEL: std::sync::OnceLock<String> = std::sync::OnceLock::new();
        return LABEL.get_or_init(|| format!("remote embed worker {url} (${})", remote::REMOTE_ENV)).as_str();
    }
    #[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
    {
        return select::chosen_backend().label();
    }
    #[cfg(all(feature = "embed-ort", not(feature = "embed-tract")))]
    {
        return "ort (ONNX Runtime, CUDA→CPU)";
    }
    #[cfg(all(feature = "embed-tract", not(feature = "embed-ort")))]
    {
        "tract-onnx (CPU, pure Rust)"
    }
}

/// `id` of the active embedding model (from `$NORNIR_EMBED_MODEL` or the
/// registry default). Available whenever `vector` is on, for diagnostics — the
/// model is selectable independently of the backend. See
/// [`embed_registry::selected`].
pub fn selected_model_id() -> &'static str {
    embed_registry::selected()
        .map(|m| m.id)
        .unwrap_or("<invalid>")
}

/// Human-readable `"<model-name> (<dim>-dim)"` of the active model, for CLI /
/// diagnostics. Reports the registry default (jina-v2-base-code, 768-dim)
/// unless `$NORNIR_EMBED_MODEL` selects another.
pub fn selected_model_desc() -> String {
    match embed_registry::selected() {
        Ok(m) => format!("{} ({}-dim)", m.model_name, m.dim),
        Err(e) => e,
    }
}

// ----- vector-ANN engine ------------------------------------------------------
//
// The pure compute engine — the exact-flat `VectorIndex`, the SIMD/int8 cosine
// kernels, and the `bench_kernels` sweep — now lives in `znippy_zoomies::vann`
// (extracted verbatim; nornir already depended on znippy-zoomies for gatling).
// It is re-exported here so every existing `nornir::vector::…` path still
// resolves unchanged (`VectorIndex`, `quantize_i8`, `score_i8_batch`, …). The
// embedder glue above (chunk/store/embed*/registry) stays in nornir.
pub use znippy_zoomies::vann::{
    BenchReport, HnswIndex, HnswParams, I8_SCAN_THRESHOLD, KernelTiming, REALISTIC_SPREAD,
    VectorIndex, active_simd, bench_kernels, clustered_corpus, clustered_row, normalized,
    quantize_i8, score_i8_batch, top_k, vnni_available,
};