P9 (WIP): compiled kernels (rust/simurg-core, PyO3) with bit-identical pure-Python fallback
Rust ports of the hot paths found by profiling: fused per-locus GenCall calling and scoring, BAF/LRR, libm maps with CPython math semantics, both nearest-neighbour searches of the normalization fit, and the BPM LocusEntry parser. Address coverage is memoised per IDAT layout. SIMURG_PURE_PYTHON=1 forces the pure paths. On the Pi: genotype_sample 3.4 s -> 1.1 s, 36-sample GSA run 131 s -> 46 s, BPM read 4x faster; 36 GSA + 6 Omni2.5 GTCs and QC tables byte-identical to the pure path; tests/unit/test_core_kernels.py checks every kernel on edge cases. bench/profile_genotype.py, bench/scaling.py, cloud task p9 (thread scaling).
This commit is contained in:
kemal committed
1 parent
f11b39b5a3
commit
6524a6d1b0
19 files changed
+1755
-6
No files matched your search
@@ -0,0 +1,226 @@
|
||||
//! Compiled kernels for SimurgArray (plan §18 P9, §19).
|
||||
//!
|
||||
//! Every kernel is a drop-in replacement for a pure-Python/NumPy reference in the `simurg` package and returns
|
||||
//! bit-identical results (checked by `tests/unit/test_core_kernels.py`):
|
||||
//!
|
||||
//! * `exp`, `log10`, `log2`, `pow` — element-wise glibc libm calls with the error semantics of CPython's `math`
|
||||
//! module as used by `simurg.genotype.gencall._c_call` (NaN input returned as is, a domain error gives the
|
||||
//! canonical +NaN, `log10(0)` = −inf, overflow = +inf). The libm symbols are the ones CPython's `math` calls, so
|
||||
//! results equal the C reference implementations (gtc2vcf, MIT) that the Python port follows.
|
||||
//! * `closest_binned` — the per-bin candidate scan of the vendor's binned nearest-neighbour search
|
||||
//! (`simurg.signal.normalization.fit._closest_binned`): float32 squared distances, first minimum wins.
|
||||
//!
|
||||
//! * `call_genotypes`, `baf_lrr` — fused per-locus GenCall calling/scoring and BAF/LRR (`gencall.rs`).
|
||||
//! * `bpm_locus_entries` — the BPM LocusEntry parser (`simurg.formats.bpm._parse_locus_entries`).
|
||||
//!
|
||||
//! No floating-point contraction (FMA) happens: Rust never fuses `a*b+c` implicitly.
|
||||
|
||||
mod bpm;
|
||||
mod gencall;
|
||||
|
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[link(name = "m")]
|
||||
unsafe extern "C" {
|
||||
#[link_name = "exp"]
|
||||
fn c_exp(x: f64) -> f64;
|
||||
#[link_name = "log10"]
|
||||
fn c_log10(x: f64) -> f64;
|
||||
#[link_name = "log2"]
|
||||
fn c_log2(x: f64) -> f64;
|
||||
#[link_name = "pow"]
|
||||
fn c_pow(x: f64, y: f64) -> f64;
|
||||
}
|
||||
|
||||
/// CPython's `float('nan')`.
|
||||
const PY_NAN: f64 = f64::from_bits(0x7ff8_0000_0000_0000);
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn math_exp(v: f64) -> f64 {
|
||||
if v.is_nan() {
|
||||
return v;
|
||||
}
|
||||
let r = unsafe { c_exp(v) };
|
||||
if r.is_nan() { PY_NAN } else if r.is_infinite() && v.is_finite() { f64::INFINITY } else { r }
|
||||
}
|
||||
|
||||
/// CPython `m_log10`/`m_log2` + `_c_call`: x > 0 → libm; 0 → −inf (log10) or NaN (log2); x < 0, −inf → NaN.
|
||||
#[inline]
|
||||
pub(crate) fn math_log(v: f64, ten: bool) -> f64 {
|
||||
if v.is_nan() {
|
||||
return v;
|
||||
}
|
||||
if v.is_infinite() {
|
||||
return if v > 0.0 { v } else { PY_NAN };
|
||||
}
|
||||
if v > 0.0 {
|
||||
return unsafe { if ten { c_log10(v) } else { c_log2(v) } };
|
||||
}
|
||||
if v == 0.0 && ten { f64::NEG_INFINITY } else { PY_NAN }
|
||||
}
|
||||
|
||||
/// CPython `math.pow(v, e)` for a finite exponent e > 0 that is not an integer, + `_c_call`.
|
||||
#[inline]
|
||||
pub(crate) fn math_pow(v: f64, e: f64) -> f64 {
|
||||
if v.is_nan() {
|
||||
return v;
|
||||
}
|
||||
if v.is_infinite() {
|
||||
return f64::INFINITY;
|
||||
}
|
||||
if v == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
let r = unsafe { c_pow(v, e) };
|
||||
if r.is_nan() { PY_NAN } else if r.is_infinite() { f64::INFINITY } else { r }
|
||||
}
|
||||
|
||||
fn map_f64<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>, f: impl Fn(f64) -> f64) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
let s = x.as_slice()?;
|
||||
let out: Vec<f64> = s.iter().map(|&v| f(v)).collect();
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn exp<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
map_f64(py, x, math_exp)
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn log10<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
map_f64(py, x, |v| math_log(v, true))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn log2<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
map_f64(py, x, |v| math_log(v, false))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn pow<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>, e: f64) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
if !(e.is_finite() && e > 0.0 && e.fract() != 0.0) {
|
||||
return Err(PyValueError::new_err("pow kernel: exponent must be finite, positive and not an integer"));
|
||||
}
|
||||
map_f64(py, x, |v| math_pow(v, e))
|
||||
}
|
||||
|
||||
/// For each sweep point `p` searching bin `choice[p]` (0..11 regular, 12..22 shifted; any other value: none), the
|
||||
/// index of the closest candidate among `eligible[i]` with `b1[i] == choice` (regular) or `b2[i] == choice − 12`
|
||||
/// (shifted), in ascending order; squared distance in float32; −1 if no candidate or the best is not < 1e20.
|
||||
#[pyfunction]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn closest_binned<'py>(
|
||||
py: Python<'py>,
|
||||
raw_x: PyReadonlyArray1<'py, f32>,
|
||||
raw_y: PyReadonlyArray1<'py, f32>,
|
||||
axis_x: PyReadonlyArray1<'py, f32>,
|
||||
axis_y: PyReadonlyArray1<'py, f32>,
|
||||
choice: PyReadonlyArray1<'py, i64>,
|
||||
eligible: PyReadonlyArray1<'py, i64>,
|
||||
b1: PyReadonlyArray1<'py, i64>,
|
||||
b2: PyReadonlyArray1<'py, i64>,
|
||||
) -> PyResult<Bound<'py, PyArray1<i64>>> {
|
||||
let (rx, ry) = (raw_x.as_slice()?, raw_y.as_slice()?);
|
||||
let (ax, ay, ch) = (axis_x.as_slice()?, axis_y.as_slice()?, choice.as_slice()?);
|
||||
let (el, b1, b2) = (eligible.as_slice()?, b1.as_slice()?, b2.as_slice()?);
|
||||
if ax.len() != ay.len() || ax.len() != ch.len() || el.len() != b1.len() || el.len() != b2.len() || rx.len() != ry.len() {
|
||||
return Err(PyValueError::new_err("closest_binned: length mismatch"));
|
||||
}
|
||||
let mut groups: Vec<Vec<usize>> = vec![Vec::new(); 23];
|
||||
for (i, &e) in el.iter().enumerate() {
|
||||
let e = usize::try_from(e).map_err(|_| PyValueError::new_err("closest_binned: negative index"))?;
|
||||
if e >= rx.len() {
|
||||
return Err(PyValueError::new_err("closest_binned: index out of range"));
|
||||
}
|
||||
if (0..12).contains(&b1[i]) {
|
||||
groups[b1[i] as usize].push(e);
|
||||
}
|
||||
if (0..11).contains(&b2[i]) {
|
||||
groups[12 + b2[i] as usize].push(e);
|
||||
}
|
||||
}
|
||||
let out: Vec<i64> = (0..ax.len())
|
||||
.map(|p| {
|
||||
let c = ch[p];
|
||||
if !(0..23).contains(&c) {
|
||||
return -1;
|
||||
}
|
||||
let cand = &groups[c as usize];
|
||||
let mut best_j = usize::MAX;
|
||||
let mut best = f32::INFINITY;
|
||||
for (j, &e) in cand.iter().enumerate() {
|
||||
let dx = rx[e] - ax[p];
|
||||
let dy = ry[e] - ay[p];
|
||||
let d = dx * dx + dy * dy;
|
||||
// numpy argmin: the first NaN wins, otherwise the first strict minimum
|
||||
if d.is_nan() {
|
||||
best_j = j;
|
||||
best = d;
|
||||
break;
|
||||
}
|
||||
if best_j == usize::MAX || d < best {
|
||||
best_j = j;
|
||||
best = d;
|
||||
}
|
||||
}
|
||||
if best_j == usize::MAX || !((best as f64) < 1e20) { -1 } else { cand[best_j] as i64 }
|
||||
})
|
||||
.collect();
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
/// Exact nearest neighbour (`simurg.signal.normalization.fit._closest_slow`): for each sweep point (x, y) the
|
||||
/// index of the reference point with the smallest float32 squared distance `(x-xr)²+(y-yr)²`; first minimum wins,
|
||||
/// a NaN distance wins as in `numpy.argmin`. `xref` must not be empty.
|
||||
#[pyfunction]
|
||||
fn closest_exact<'py>(
|
||||
py: Python<'py>,
|
||||
xref: PyReadonlyArray1<'py, f32>,
|
||||
yref: PyReadonlyArray1<'py, f32>,
|
||||
x: PyReadonlyArray1<'py, f32>,
|
||||
y: PyReadonlyArray1<'py, f32>,
|
||||
) -> PyResult<Bound<'py, PyArray1<i64>>> {
|
||||
let (xr, yr, xs, ys) = (xref.as_slice()?, yref.as_slice()?, x.as_slice()?, y.as_slice()?);
|
||||
if xr.is_empty() || xr.len() != yr.len() || xs.len() != ys.len() {
|
||||
return Err(PyValueError::new_err("closest_exact: empty reference or length mismatch"));
|
||||
}
|
||||
let out: Vec<i64> = xs
|
||||
.iter()
|
||||
.zip(ys)
|
||||
.map(|(&px, &py_)| {
|
||||
let mut best_j = 0usize;
|
||||
let mut best = f32::INFINITY;
|
||||
for (j, (&rx, &ry)) in xr.iter().zip(yr).enumerate() {
|
||||
let dx = px - rx;
|
||||
let dy = py_ - ry;
|
||||
let d = dx * dx + dy * dy;
|
||||
if d.is_nan() {
|
||||
return j as i64;
|
||||
}
|
||||
if j == 0 || d < best {
|
||||
best_j = j;
|
||||
best = d;
|
||||
}
|
||||
}
|
||||
best_j as i64
|
||||
})
|
||||
.collect();
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn simurg_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
|
||||
m.add_function(wrap_pyfunction!(exp, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(log10, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(log2, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(pow, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(closest_binned, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(closest_exact, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(bpm::bpm_locus_entries, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(gencall::call_genotypes, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(gencall::baf_lrr, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
Reference in new issue
Block a user