diff --git a/algebra/backend.py b/algebra/backend.py index 28b6d11b..da810fcf 100644 --- a/algebra/backend.py +++ b/algebra/backend.py @@ -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. diff --git a/core-rs/Cargo.toml b/core-rs/Cargo.toml index 8230b931..c31a92e7 100644 --- a/core-rs/Cargo.toml +++ b/core-rs/Cargo.toml @@ -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"] } diff --git a/core-rs/src/lib.rs b/core-rs/src/lib.rs index bc1f724a..8b071708 100644 --- a/core-rs/src/lib.rs +++ b/core-rs/src/lib.rs @@ -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 { 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> { - 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())) } diff --git a/core-rs/src/vault.rs b/core-rs/src/vault.rs index 6d63acfa..f095ee4f 100644 --- a/core-rs/src/vault.rs +++ b/core-rs/src/vault.rs @@ -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, 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) }