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).

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01814GKVzZVL3FuF3DeiZXhv

src/simurg/_core.py

@@ -0,0 +1,24 @@
+"""Optional compiled kernels (``rust/simurg-core``, plan §18 P9).
+
+``core`` is the ``simurg_core`` extension module when it is installed, else ``None``; every caller keeps a pure
+Python/NumPy path with bit-identical results. ``SIMURG_PURE_PYTHON=1`` forces the pure paths (tests, debugging).
+"""
+
+from __future__ import annotations
+
+import os
+from types import ModuleType
+
+core: ModuleType | None
+if os.environ.get("SIMURG_PURE_PYTHON", "") not in ("", "0"):
+    core = None
+else:
+    try:
+        import simurg_core as core  # type: ignore[no-redef]
+    except ImportError:  # pragma: no cover - depends on the installation
+        core = None
+
+
+def backend() -> str:
+    """``"rust <version>"`` or ``"python"`` (recorded in run manifests)."""
+    return f"rust {core.__version__}" if core is not None else "python"

src/simurg/formats/bpm.py

@@ -32,6 +32,7 @@ from typing import Any
 import numpy as np
 import numpy.typing as npt
 
+from simurg._core import core
 from simurg.core.errors import FormatError
 from simurg.formats.binutil import BinaryReader, load_bytes
 
@@ -223,7 +224,9 @@ def read_bpm(path: str | Path) -> Manifest:
 
 
 def _parse_locus_entries(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -> dict[str, Any]:
-    """Hot loop: parse ``n`` LocusEntry records (optimised pure Python; Rust in phase P9)."""
+    """Hot loop: parse ``n`` LocusEntry records (compiled in ``simurg_core`` when installed; same results)."""
+    if core is not None:
+        return _parse_locus_entries_core(r, n, name_to_idx)
     view = memoryview(r._buf)
     size = r.size
     pos = r.pos
@@ -416,6 +419,53 @@ def _parse_locus_entries(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -
     }
 
 
+_STR_COLS = (
+    "ilmn_ids",
+    "snps",
+    "chroms",
+    "ilmn_strands",
+    "source_strands",
+    "ref_strands",
+    "ploidies",
+    "species",
+    "genome_builds",
+    "sources",
+    "source_versions",
+    "customer_strands",
+    "source_seqs",
+    "top_genomic_seqs",
+    "allele_a_probe_seqs",
+    "allele_b_probe_seqs",
+)
+
+
+def _parse_locus_entries_core(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -> dict[str, Any]:
+    """``_parse_locus_entries`` through ``simurg_core.bpm_locus_entries`` (identical columns and error messages)."""
+    assert core is not None
+    pos = r.pos
+    try:
+        lists, map_infos, a, b, assay, exp_cl, int_only, versions, fracs, used = core.bpm_locus_entries(
+            bytes(r._view[pos:]), n, list(name_to_idx)
+        )
+    except ValueError as e:
+        raise r.fail(str(e)) from None
+    out: dict[str, Any] = dict(zip(_STR_COLS, lists, strict=True))
+    for i in np.flatnonzero(assay < 0):
+        assay[i] = infer_assay_type(out["allele_a_probe_seqs"][i], out["allele_b_probe_seqs"][i], out["source_seqs"][i])
+    r.pos = pos + used
+    return {
+        **out,
+        "map_infos": map_infos,
+        "address_a": a,
+        "address_b": b,
+        "assay_types": assay.astype(np.uint8),
+        "exp_clusters": exp_cl,
+        "intensity_only": int_only,
+        "entry_versions": versions,
+        "base_fractions": fracs,
+    }
+
+
 _IUPAC = {
     "A": 1,
     "C": 2,

src/simurg/genotype/baf_lrr.py

@@ -11,6 +11,7 @@ import math
 import numpy as np
 import numpy.typing as npt
 
+from simurg._core import core
 from simurg.genotype.clusters import ClusterArrays
 from simurg.genotype.gencall import _libm
 
@@ -22,6 +23,17 @@ NEG_NAN = np.frombuffer(np.uint32(0xFFC00000).tobytes(), dtype=np.float32)[0]  #
 def baf_lrr(
     theta: F32Arr, r: F32Arr, raw_zero: npt.NDArray[np.bool_], c: ClusterArrays, intensity_only: npt.NDArray[np.bool_]
 ) -> tuple[F32Arr, F32Arr]:
+    if core is not None:
+        return core.baf_lrr(  # type: ignore[no-any-return]
+            np.ascontiguousarray(theta, np.float32),
+            np.ascontiguousarray(r, np.float32),
+            np.ascontiguousarray(raw_zero, np.bool_),
+            np.ascontiguousarray(c.present, np.bool_),
+            np.ascontiguousarray(c.theta_mean, np.float32).ravel(),
+            np.ascontiguousarray(c.r_mean, np.float32).ravel(),
+            np.ascontiguousarray(c.mean_r, np.float32),
+            np.ascontiguousarray(intensity_only, np.bool_),
+        )
     t = np.asarray(theta, np.float32)
     rr = np.asarray(r, np.float32)
     at, bt, ct = c.theta_mean[:, 0], c.theta_mean[:, 1], c.theta_mean[:, 2]

src/simurg/genotype/gencall.py

@@ -15,6 +15,7 @@ from dataclasses import dataclass
 import numpy as np
 import numpy.typing as npt
 
+from simurg._core import core
 from simurg.genotype.clusters import ClusterArrays
 
 F32 = np.float32
@@ -42,8 +43,17 @@ def _c_call(fn, v: float) -> float:  # type: ignore[no-untyped-def]
         return math.inf
 
 
+_CORE_FN = {math.exp: "exp", math.log10: "log10", math.log2: "log2"}
+
+
 def _libm(fn, x: F64Arr) -> F64Arr:  # type: ignore[no-untyped-def]
-    """Element-wise libm call (bit-identical to C)."""
+    """Element-wise libm call (bit-identical to C); compiled when ``simurg_core`` is installed."""
+    if core is not None:
+        xa = np.ascontiguousarray(x, dtype=np.float64).ravel()
+        if (e := getattr(fn, "exponent", None)) is not None:
+            return core.pow(xa, e).reshape(np.shape(x))  # type: ignore[no-any-return]
+        if (name := _CORE_FN.get(fn)) is not None:
+            return getattr(core, name)(xa).reshape(np.shape(x))  # type: ignore[no-any-return]
     return np.fromiter((_c_call(fn, v) for v in x.tolist()), dtype=np.float64, count=len(x))
 
 
@@ -51,6 +61,7 @@ def _pow(e: float):  # type: ignore[no-untyped-def]
     def f(v: float) -> float:
         return math.pow(v, e)
 
+    f.exponent = e  # type: ignore[attr-defined]
     return f
 
 
@@ -137,6 +148,23 @@ def call_genotypes(
     cutoff: float = DEFAULT_CUTOFF,
 ) -> Calls:
     """Bayesian calls and GenCall scores for one sample (plan §7.5)."""
+    if core is not None:
+        f32 = [np.ascontiguousarray(a, np.float32).ravel() for a in (c.theta_mean, c.theta_dev, c.r_mean, c.r_dev)]
+        geno, score = core.call_genotypes(
+            np.ascontiguousarray(theta, np.float32),
+            np.ascontiguousarray(r, np.float32),
+            np.ascontiguousarray(raw_zero, np.bool_),
+            np.ascontiguousarray(c.present, np.bool_),
+            np.ascontiguousarray(c.n, np.int32).ravel(),
+            *f32,
+            np.ascontiguousarray(c.intensity_threshold, np.float32),
+            np.ascontiguousarray(c.gentrain, np.float32),
+            np.ascontiguousarray(intensity_only, np.bool_),
+            F32(cutoff),
+        )
+        return Calls(
+            genotypes=geno, scores=score, num_calls=int((geno > 0).sum()), num_intensity_only=int(intensity_only.sum())
+        )
     n = len(theta)
     t = np.asarray(theta, np.float32)
     rr = np.asarray(r, np.float32)

src/simurg/pipeline/genotype.py

@@ -160,7 +160,7 @@ def _process(job: Job) -> dict[str, Any]:
     try:
         grn, red = read_idat(job.grn), read_idat(job.red)
         check_pair(grn, red)
-        _, notes = check_coverage(grn, ctx.manifest)
+        _, notes = check_coverage(grn, ctx.manifest, ctx.cache)
         row["notes"].extend(notes)
         g = genotype_sample(grn, red, ctx, params, _W["started"])
         write_gtc(g, Path(job.gtc_path))

src/simurg/signal/normalization/fit.py

@@ -14,6 +14,7 @@ from dataclasses import dataclass
 import numpy as np
 import numpy.typing as npt
 
+from simurg._core import core
 from simurg.signal.normalization.numerics import (
     F32,
     TINY,
@@ -86,6 +87,8 @@ def closest_points(xref: FloatArr, yref: FloatArr, x: FloatArr, y: FloatArr) ->
 
 def _closest_slow(xref: FloatArr, yref: FloatArr, x: FloatArr, y: FloatArr) -> IdxArr:
     """Exact search: squared distances in float32; the first minimum wins (C strict ``<``)."""
+    if core is not None and len(xref) and all(a.dtype == np.float32 for a in (xref, yref, x, y)):
+        return core.closest_exact(*(np.ascontiguousarray(a) for a in (xref, yref, x, y)))  # type: ignore[no-any-return]
     dx = x[:, None] - xref[None, :]
     dy = y[:, None] - yref[None, :]
     return np.argmin(dx * dx + dy * dy, axis=1).astype(np.int64)
@@ -118,6 +121,12 @@ def _closest_binned(raw_x: FloatArr, raw_y: FloatArr, axis_x: FloatArr, axis_y:
     low = rem < F32(0.25)
     choice = np.where(mid, k, np.where(low, np.where(k == 0, k, 12 + k - 1), np.where(k == 11, k, 12 + k)))
 
+    if core is not None and all(a.dtype == np.float32 for a in (raw_x, raw_y, axis_x, axis_y)):
+        i64 = np.int64
+        return core.closest_binned(  # type: ignore[no-any-return]
+            *(np.ascontiguousarray(a, np.float32) for a in (raw_x, raw_y, axis_x, axis_y)),
+            *(np.ascontiguousarray(a, i64) for a in (choice, eligible, b1, b2)),
+        )
     out = np.full(len(axis_a), -1, dtype=np.int64)
     for c in np.unique(choice):
         cand = eligible[b1 == c] if c < 12 else eligible[b2 == c - 12]

src/simurg/signal/raw_xy.py

@@ -90,9 +90,11 @@ def address_coverage(idat: Idat, manifest: Manifest) -> float:
     return float(found.mean())
 
 
-def check_coverage(idat: Idat, manifest: Manifest) -> tuple[float, list[tuple[str, str]]]:
+def check_coverage(
+    idat: Idat, manifest: Manifest, cache: LocusMapCache | None = None
+) -> tuple[float, list[tuple[str, str]]]:
     """IN-027: < 0.99 → E_IDAT_MANIFEST_MISMATCH (raised); < 1 → W_IDAT_LOW_ADDRESS_COVERAGE (returned)."""
-    cov = address_coverage(idat, manifest)
+    cov = cache.coverage(idat) if cache is not None and cache.manifest is manifest else address_coverage(idat, manifest)
     notes: list[tuple[str, str]] = []
     if cov < COVERAGE_ERROR:
         raise SimurgError(
@@ -149,10 +151,24 @@ class LocusMapCache:
         self.manifest = manifest
         self.caddr = control_addresses(manifest)
         self._maps: dict[bytes, LocusMap] = {}
+        self._coverage: dict[bytes, float] = {}
         self.max_entries = max_entries
 
+    @staticmethod
+    def _key(idat: Idat) -> bytes:
+        return hashlib.blake2b(np.ascontiguousarray(idat.address_ids).tobytes(), digest_size=16).digest()
+
+    def coverage(self, idat: Idat) -> float:
+        """``address_coverage`` memoised per address layout."""
+        key = self._key(idat)
+        if (cov := self._coverage.get(key)) is None:
+            if len(self._coverage) >= self.max_entries:
+                self._coverage.pop(next(iter(self._coverage)))
+            cov = self._coverage[key] = address_coverage(idat, self.manifest)
+        return cov
+
     def get(self, idat: Idat) -> LocusMap:
-        key = hashlib.blake2b(np.ascontiguousarray(idat.address_ids).tobytes(), digest_size=16).digest()
+        key = self._key(idat)
         lm = self._maps.get(key)
         if lm is None:
             if len(self._maps) >= self.max_entries: