core-labs/core
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

Patch PR #907 to make scalar Rust inputs fully zero-copy without stack array copies (#908)

Browse source 
(cherry picked from commit 59b2a8f526)
This commit is contained in:
Shay 2026-06-24 13:36:36 -07:00 committed by GitHub
parent 99a107dd39
commit e77b0299a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 22 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
core-rs/src/lib.rs
View file

@ -39,9 +39,9 @@ fn geometric_product(
a: numpy::PyReadonlyArray1<'_, f32>, a: numpy::PyReadonlyArray1<'_, f32>,
b: numpy::PyReadonlyArray1<'_, f32>, b: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<PyObject> { ) -> PyResult<PyObject> {
let a_slice = read_f32_cl41_mv(a)?; let a_slice = read_f32_cl41_mv(&a)?;
let b_slice = read_f32_cl41_mv(b)?; let b_slice = read_f32_cl41_mv(&b)?;
let result = geometric_product_raw(&a_slice, &b_slice) let result = geometric_product_raw(a_slice, b_slice)
.map_err(|e| PyValueError::new_err(e.to_string()))?; .map_err(|e| PyValueError::new_err(e.to_string()))?;
f32_array_to_numpy(py, &result) f32_array_to_numpy(py, &result)
} }
@ -68,9 +68,9 @@ fn versor_apply_with_closure(
v: numpy::PyReadonlyArray1<'_, f32>, v: numpy::PyReadonlyArray1<'_, f32>,
f: numpy::PyReadonlyArray1<'_, f32>, f: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<PyObject> { ) -> PyResult<PyObject> {
let v_slice = read_f32_cl41_mv(v)?; let v_slice = read_f32_cl41_mv(&v)?;
let f_slice = read_f32_cl41_mv(f)?; let f_slice = read_f32_cl41_mv(&f)?;
let result = versor_apply_closed(&v_slice, &f_slice) let result = versor_apply_closed(v_slice, f_slice)
.map_err(|e| PyValueError::new_err(e.to_string()))?; .map_err(|e| PyValueError::new_err(e.to_string()))?;
f32_array_to_numpy(py, &result) f32_array_to_numpy(py, &result)
} }
@ -84,9 +84,9 @@ fn versor_apply_with_closure_f64(
v: numpy::PyReadonlyArray1<'_, f64>, v: numpy::PyReadonlyArray1<'_, f64>,
f: numpy::PyReadonlyArray1<'_, f64>, f: numpy::PyReadonlyArray1<'_, f64>,
) -> PyResult<PyObject> { ) -> PyResult<PyObject> {
let v_slice = read_f64_cl41_mv(v)?; let v_slice = read_f64_cl41_mv(&v)?;
let f_slice = read_f64_cl41_mv(f)?; let f_slice = read_f64_cl41_mv(&f)?;
let result = versor_apply_closed_f64(&v_slice, &f_slice) let result = versor_apply_closed_f64(v_slice, f_slice)
.map_err(|e| PyValueError::new_err(e.to_string()))?; .map_err(|e| PyValueError::new_err(e.to_string()))?;
f64_array_to_numpy(py, &result) f64_array_to_numpy(py, &result)
} }
@ -94,8 +94,8 @@ fn versor_apply_with_closure_f64(
/// ||F*reverse(F) - 1||_F. Returns scalar f32. /// ||F*reverse(F) - 1||_F. Returns scalar f32.
#[pyfunction] #[pyfunction]
fn versor_condition(f: numpy::PyReadonlyArray1<'_, f32>) -> PyResult<f32> { fn versor_condition(f: numpy::PyReadonlyArray1<'_, f32>) -> PyResult<f32> {
let f_slice = read_f32_cl41_mv(f)?; let f_slice = read_f32_cl41_mv(&f)?;
versor_condition_raw(&f_slice).map_err(|e| PyValueError::new_err(e.to_string())) versor_condition_raw(f_slice).map_err(|e| PyValueError::new_err(e.to_string()))
} }
/// Project F onto versor manifold: F / sqrt(|F*rev(F)|). /// Project F onto versor manifold: F / sqrt(|F*rev(F)|).
@ -116,9 +116,9 @@ fn cga_inner(
x: numpy::PyReadonlyArray1<'_, f32>, x: numpy::PyReadonlyArray1<'_, f32>,
y: numpy::PyReadonlyArray1<'_, f32>, y: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<f32> { ) -> PyResult<f32> {
let x_slice = read_f32_cl41_mv(x)?; let x_slice = read_f32_cl41_mv(&x)?;
let y_slice = read_f32_cl41_mv(y)?; let y_slice = read_f32_cl41_mv(&y)?;
cga_inner_raw(&x_slice, &y_slice).map_err(|e| PyValueError::new_err(e.to_string())) 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 (zero-copy). /// Parallel top-k vault recall by CGA inner product (zero-copy).
@ -260,7 +260,7 @@ fn diffusion_step<'py>(
Ok((numpy::IntoPyArray::into_pyarray_bound(arr, py), delta)) Ok((numpy::IntoPyArray::into_pyarray_bound(arr, py), delta))
} }
fn read_f32_cl41_mv(arr: numpy::PyReadonlyArray1<'_, f32>) -> PyResult<[f32; 32]> { fn read_f32_cl41_mv<'a>(arr: &'a numpy::PyReadonlyArray1<'a, f32>) -> PyResult<&'a [f32; 32]> {
let len = arr.len()?; let len = arr.len()?;
if len != 32 { if len != 32 {
return Err(PyValueError::new_err(format!( return Err(PyValueError::new_err(format!(
@ -274,12 +274,12 @@ fn read_f32_cl41_mv(arr: numpy::PyReadonlyArray1<'_, f32>) -> PyResult<[f32; 32]
e e
)) ))
})?; })?;
let mut out = [0f32; 32]; slice.try_into().map_err(|_| {
out.copy_from_slice(slice); PyValueError::new_err("expected contiguous float32 array of length 32")
Ok(out) })
} }
fn read_f64_cl41_mv(arr: numpy::PyReadonlyArray1<'_, f64>) -> PyResult<[f64; 32]> { fn read_f64_cl41_mv<'a>(arr: &'a numpy::PyReadonlyArray1<'a, f64>) -> PyResult<&'a [f64; 32]> {
let len = arr.len()?; let len = arr.len()?;
if len != 32 { if len != 32 {
return Err(PyValueError::new_err(format!( return Err(PyValueError::new_err(format!(
@ -293,9 +293,9 @@ fn read_f64_cl41_mv(arr: numpy::PyReadonlyArray1<'_, f64>) -> PyResult<[f64; 32]
e e
)) ))
})?; })?;
let mut out = [0f64; 32]; slice.try_into().map_err(|_| {
out.copy_from_slice(slice); PyValueError::new_err("expected contiguous float64 array of length 32")
Ok(out) })
} }
fn extract_f32_slice(obj: &pyo3::types::PyAny) -> PyResult<[f32; 32]> { fn extract_f32_slice(obj: &pyo3::types::PyAny) -> PyResult<[f32; 32]> {