perf(rust): zero-copy vault_recall — Rust beats Python at scale
ADR-0020 follow-on (task #35). Two-pronged fix: 1. Kernel: ported ADR-0019 Stage 1 diagonal-metric kernel to core-rs/src/vault.rs. Per-versor scoring is now 32 multiplies + 32 adds via the precomputed Cl(4,1) metric, not the 1024-op full geometric_product the prior path computed. Bit-identity preserved by serial fold order matching Python. 2. Zero-copy marshalling: replaced Vec<&PyAny> + extract-per- versor with PyReadonlyArray2<f32> via the numpy Rust crate. The Rust binding now reads a slice view directly into the numpy buffer — no Python→Rust copy, no Vec<[f32;32]> re-chunk. Python caller passes the (N, 32) ndarray as-is (ascontiguousarray ensures C-contiguous f32). Result: N python rust speedup 1k 0.20ms 0.26ms 0.77x (Python wins on fixed overhead) 10k 1.62ms 1.45ms 1.12x 100k 19.22ms 12.93ms 1.49x 1M 251.50ms 131.36ms 1.91x Parity bit-identical (raw f32 bytes) at every scale across the parameterised test in tests/test_vault_recall_rust_parity.py. Both ADR-0020 first-surface gates now pass: parity AND performance at the scales where Rust is meant to win. Python remains the default per CLAUDE.md sequencing rule 5; CORE_BACKEND=rust is now a legitimate opt-in acceleration. Smoke 27/27, algebra 70/70, runtime 19/19 all green.
This commit is contained in:
Shay
parent
44e98790c3
commit
e36998d25d
4 changed files with 158 additions and 31 deletions
|
|
@ -95,15 +95,21 @@ def vault_recall(versors: list, query: np.ndarray, top_k: int = 5) -> list:
|
|||
order; the only thing the vectorisation replaces is the
|
||||
per-element Python dispatch loop. ADR-0019 Stage 1.
|
||||
"""
|
||||
if _RUST:
|
||||
try:
|
||||
return _rs.vault_recall(versors, query, top_k)
|
||||
except Exception:
|
||||
pass
|
||||
if not versors:
|
||||
return []
|
||||
q = np.asarray(query, dtype=np.float32)
|
||||
M = np.asarray(versors, dtype=np.float32)
|
||||
if _RUST and M.ndim == 2 and M.shape[1] == 32:
|
||||
try:
|
||||
# Pass the (N, 32) numpy buffer directly — the Rust
|
||||
# binding reads it zero-copy via PyReadonlyArray2 (task
|
||||
# #35). ascontiguousarray ensures C-contiguous f32
|
||||
# layout, which the zero-copy slice requires.
|
||||
Mc = np.ascontiguousarray(M, dtype=np.float32)
|
||||
qc = np.ascontiguousarray(q, dtype=np.float32)
|
||||
return _rs.vault_recall(Mc, qc, top_k)
|
||||
except Exception:
|
||||
pass
|
||||
if M.ndim != 2:
|
||||
# Heterogeneous shapes — fall back to the scalar path rather
|
||||
# than coerce silently.
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ crate-type = ["cdylib", "rlib"]
|
|||
|
||||
[dependencies]
|
||||
pyo3 = { version = "0.21" }
|
||||
numpy = "0.21"
|
||||
rayon = "1.10"
|
||||
nalgebra = "0.33"
|
||||
ndarray = { version = "0.16", features = ["rayon"] }
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ pub mod versor;
|
|||
use cga::cga_inner_raw;
|
||||
use cl41::geometric_product_raw;
|
||||
use diffusion::{graph_diffusion_step, unitize_f32};
|
||||
#[allow(unused_imports)]
|
||||
use vault::vault_recall_raw;
|
||||
use versor::{normalize_to_versor_raw, versor_apply_closed, versor_apply_raw, versor_condition_raw};
|
||||
|
||||
|
|
@ -94,19 +95,50 @@ fn cga_inner(x: &pyo3::types::PyAny, y: &pyo3::types::PyAny) -> PyResult<f32> {
|
|||
cga_inner_raw(&x_slice, &y_slice).map_err(|e| PyValueError::new_err(e.to_string()))
|
||||
}
|
||||
|
||||
/// Parallel top-k vault recall by CGA inner product.
|
||||
/// Parallel top-k vault recall by CGA inner product (zero-copy).
|
||||
///
|
||||
/// Per ADR-0020 follow-on (task #35): accepts a 2D numpy
|
||||
/// (N, 32) float32 array via `PyReadonlyArray2`, which exposes a
|
||||
/// view *directly into the numpy buffer*. No Python→Rust copy,
|
||||
/// no re-chunking — Rayon scores straight off the source slice.
|
||||
/// This is the load-bearing reason for the Rust path: NumPy
|
||||
/// already SIMD-vectorises the same kernel; the only win Rust
|
||||
/// can offer is *avoiding the marshalling tax* and adding
|
||||
/// thread-parallel scoring on top.
|
||||
#[pyfunction]
|
||||
fn vault_recall(
|
||||
versors: Vec<&pyo3::types::PyAny>,
|
||||
query: &pyo3::types::PyAny,
|
||||
versors: numpy::PyReadonlyArray2<'_, f32>,
|
||||
query: numpy::PyReadonlyArray1<'_, f32>,
|
||||
top_k: usize,
|
||||
) -> PyResult<Vec<(usize, f32)>> {
|
||||
let query_slice = extract_f32_slice(query)?;
|
||||
let mut slices: Vec<[f32; 32]> = Vec::with_capacity(versors.len());
|
||||
for v in &versors {
|
||||
slices.push(extract_f32_slice(v)?);
|
||||
let view = versors.as_array();
|
||||
let shape = view.shape();
|
||||
if shape.len() != 2 || shape[1] != 32 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"versors must be shape (N, 32), got {:?}",
|
||||
shape
|
||||
)));
|
||||
}
|
||||
vault_recall_raw(&slices, &query_slice, top_k)
|
||||
let n = shape[0];
|
||||
let q_slice = query.as_slice().map_err(|e| {
|
||||
PyValueError::new_err(format!("query must be contiguous f32 (32,): {}", e))
|
||||
})?;
|
||||
if q_slice.len() != 32 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"query must have length 32, got {}",
|
||||
q_slice.len()
|
||||
)));
|
||||
}
|
||||
let v_slice = versors.as_slice().map_err(|e| {
|
||||
PyValueError::new_err(format!(
|
||||
"versors must be C-contiguous f32 (N, 32): {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
let mut q_arr = [0f32; 32];
|
||||
q_arr.copy_from_slice(q_slice);
|
||||
|
||||
crate::vault::vault_recall_flat(v_slice, n, &q_arr, top_k)
|
||||
.map_err(|e| PyValueError::new_err(e.to_string()))
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,15 +1,23 @@
|
|||
//! VaultStore hot path: parallel top-k CGA inner product scan.
|
||||
//!
|
||||
//! Uses Rayon for data-parallel scoring across all stored versors.
|
||||
//! Each worker computes cga_inner(query, v) independently — no shared state,
|
||||
//! no locks. Results are merged with a partial sort for top-k.
|
||||
//! Each worker computes the diagonal-metric CGA inner product
|
||||
//! independently — no shared state, no locks. Top-k is finalised
|
||||
//! with a stable sort that mirrors Python's ascending-index
|
||||
//! tie-break (ADR-0019 Stage 1 + ADR-0020 first-surface parity).
|
||||
//!
|
||||
//! This is the primary reason the vault scan is in Rust:
|
||||
//! Python cannot release the GIL across a list comprehension.
|
||||
//! Rayon gives us true multithreaded scoring with zero-copy slice access.
|
||||
//! The CGA inner product in Cl(4,1) is structurally diagonal with
|
||||
//! ±1 metric values, so per-versor scoring collapses to
|
||||
//!
|
||||
//! sum_i metric[i] * v[i] * q[i]
|
||||
//!
|
||||
//! which is 32 multiplies + 32 adds, not the 1024-op full
|
||||
//! geometric_product the reference scalar path computes. Bit-
|
||||
//! identity with Python's vectorised path is preserved because
|
||||
//! the serial fold order is identical (i = 0..32, left-to-right
|
||||
//! accumulation) and float32 multiply/add are deterministic.
|
||||
|
||||
use rayon::prelude::*;
|
||||
use crate::cga::cga_inner_raw;
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
|
|
@ -18,14 +26,87 @@ pub enum VaultError {
|
|||
Cga(String),
|
||||
}
|
||||
|
||||
/// Diagonal Cl(4,1) CGA inner-product metric. Derived once at
|
||||
/// build time from cga_inner(e_i, e_i) over the 32-component
|
||||
/// basis. See `tests/test_vault_recall_vectorised.py` (Python
|
||||
/// side) for the empirical derivation that pins this vector.
|
||||
const CGA_INNER_METRIC: [f32; 32] = [
|
||||
1.0, 1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0,
|
||||
-1.0, 1.0, -1.0, -1.0, 1.0, -1.0, 1.0, 1.0,
|
||||
-1.0, -1.0, 1.0, -1.0, 1.0, 1.0, -1.0, 1.0,
|
||||
1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0, -1.0,
|
||||
];
|
||||
|
||||
/// Per-versor diagonal-metric CGA inner product. Same arithmetic
|
||||
/// order as Python's `(metric[i] * v[i]) * q[i]` serial fold —
|
||||
/// bit-identical scores by construction.
|
||||
#[inline(always)]
|
||||
fn diagonal_score(v: &[f32; 32], q: &[f32; 32]) -> f32 {
|
||||
let mut s: f32 = 0.0;
|
||||
for i in 0..32 {
|
||||
let t = CGA_INNER_METRIC[i] * v[i];
|
||||
s += t * q[i];
|
||||
}
|
||||
s
|
||||
}
|
||||
|
||||
/// Zero-copy parallel top-k recall by CGA inner product, over a
|
||||
/// flat (N*32,) f32 slice viewed directly from a numpy buffer.
|
||||
///
|
||||
/// `versors_flat` must hold N consecutive [f32; 32] blocks in C
|
||||
/// (row-major) order. No copies are made; Rayon scores straight
|
||||
/// off the source slice with stride 32.
|
||||
pub fn vault_recall_flat(
|
||||
versors_flat: &[f32],
|
||||
n: usize,
|
||||
query: &[f32; 32],
|
||||
top_k: usize,
|
||||
) -> Result<Vec<(usize, f32)>, VaultError> {
|
||||
if n == 0 {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
debug_assert_eq!(versors_flat.len(), n * 32);
|
||||
|
||||
let mut scores: Vec<(usize, f32)> = (0..n)
|
||||
.into_par_iter()
|
||||
.map(|i| {
|
||||
let v = &versors_flat[i * 32..(i + 1) * 32];
|
||||
let mut s: f32 = 0.0;
|
||||
for j in 0..32 {
|
||||
let t = CGA_INNER_METRIC[j] * v[j];
|
||||
s += t * query[j];
|
||||
}
|
||||
(i, s)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let k = top_k.min(scores.len());
|
||||
if k < scores.len() {
|
||||
scores.select_nth_unstable_by(k.saturating_sub(1), |a, b| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then(a.0.cmp(&b.0))
|
||||
});
|
||||
scores.truncate(k);
|
||||
}
|
||||
scores.sort_by(|a, b| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then(a.0.cmp(&b.0))
|
||||
});
|
||||
|
||||
Ok(scores)
|
||||
}
|
||||
|
||||
/// Parallel top-k recall by CGA inner product.
|
||||
///
|
||||
/// versors: slice of [f32; 32] stored versors
|
||||
/// query: [f32; 32] query versor
|
||||
/// top_k: number of results to return
|
||||
///
|
||||
/// Returns Vec<(index, score)> sorted descending by score.
|
||||
/// Thread-safe: Rayon spawns workers per chunk, no locks required.
|
||||
/// Returns Vec<(index, score)> sorted descending by score, with
|
||||
/// ascending-index tie-break. Thread-safe: Rayon spawns workers
|
||||
/// per chunk, no locks required.
|
||||
pub fn vault_recall_raw(
|
||||
versors: &[[f32; 32]],
|
||||
query: &[f32; 32],
|
||||
|
|
@ -35,23 +116,30 @@ pub fn vault_recall_raw(
|
|||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
// Score all versors in parallel
|
||||
// Score all versors in parallel via the diagonal kernel.
|
||||
let mut scores: Vec<(usize, f32)> = versors
|
||||
.par_iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| {
|
||||
let score = cga_inner_raw(v, query).unwrap_or(f32::NEG_INFINITY);
|
||||
(i, score)
|
||||
})
|
||||
.map(|(i, v)| (i, diagonal_score(v, query)))
|
||||
.collect();
|
||||
|
||||
// Partial sort: bring top_k to the front
|
||||
// Stable top-k order: descending score, ascending index on
|
||||
// ties (mirrors Python list.sort with key=lambda x: -x[1] on
|
||||
// an enumerated list).
|
||||
let k = top_k.min(scores.len());
|
||||
scores.select_nth_unstable_by(k.saturating_sub(1), |a, b| {
|
||||
b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
|
||||
if k < scores.len() {
|
||||
scores.select_nth_unstable_by(k.saturating_sub(1), |a, b| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then(a.0.cmp(&b.0))
|
||||
});
|
||||
scores.truncate(k);
|
||||
}
|
||||
scores.sort_by(|a, b| {
|
||||
b.1.partial_cmp(&a.1)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then(a.0.cmp(&b.0))
|
||||
});
|
||||
scores.truncate(k);
|
||||
scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
Ok(scores)
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue