Optimize SamplingDelayAlignment for MLX and right-size its fractional-delay FIR

docs/source/guides/processing.md

@@ -92,6 +92,50 @@ Set `filter_len` to `0` to disable alignment entirely — the transformer become
 a pass-through that returns its input unchanged, handy for A/B comparisons or for
 leaving the unit wired in but inert.
 
+### Choosing `filter_len`
+
+Longer filters flatten the response nearer Nyquist but add latency: the bulk
+delay is `(filter_len-1)//2` samples. The table below is the worst case over all
+32 within-bank fractional delays at 30 kHz, and is pinned by
+`tests/test_sampling_delay_alignment.py`:
+
+| Passband | `filter_len` | Max phase error | Max magnitude error | Bulk delay |
+|---|---|---|---|---|
+| 0–500 Hz (LFP) | 7 | 0.0009° | 0.00001 dB | 3 samples (100 µs) |
+| 0–3 kHz | 9 | 0.0038° | 0.0001 dB | 4 samples (133 µs) |
+| 0–7.5 kHz (broadband/spike) | 13 *(default)* | 0.015° | 0.0015 dB | 6 samples (200 µs) |
+| 0–7.5 kHz | 33 | 0.0028° | 0.0004 dB | 16 samples (533 µs) |
+
+The requirement the default is chosen against is 0.05° of phase error and 0.01 dB
+of magnitude error over the intended passband — roughly three orders of magnitude
+below the ~81° of skew being corrected. 13 taps is the shortest odd length that
+holds that over the full broadband band (11 taps misses it at 0.09°), so it is
+the default. Drop to 7 in an LFP-only pipeline to halve the latency again; 33
+was the previous default and still works, but buys accuracy that is already far
+below the noise floor at three times the latency.
+
+### Performance
+
+The transformer runs once per message, so per-message cost at live chunk sizes —
+tens of samples, not thousands — is what matters. The portable Array-API
+implementation (the per-tap multiply-add loop and a log-depth rail scan) runs on
+any backend; two optional fast paths sit on top of it:
+
+- **MLX.** The FIR is evaluated as a single depthwise convolution (one group per
+  channel, kernel laid out once at state reset) rather than one dispatched
+  multiply-add per tap — about 2–3× faster per message and roughly flat from 3
+  to 1000 samples. The rail forward-fill's running max uses `mx.cummax`.
+- **numpy, with the optional `accel` extra (`pip install ezmsg-blackrock[accel]`,
+  which pulls in `numba`).** Both the FIR and the rail forward-fill run as fused
+  single-pass JIT-compiled kernels — roughly 3× faster than the tap loop at live
+  chunk sizes, and, via a threaded variant that engages automatically for large
+  buffers, an order of magnitude faster for offline batch processing of whole
+  recordings. Without the extra, numpy uses the portable path (the rail scan
+  taking the `maximum.accumulate` one-pass form), so nothing here is required.
+
+`examples/bench_sampling_delay_alignment.py` reports all of this across channel
+counts, message sizes, rail handling on/off, and fast-path versus portable.
+
 A few things to keep in mind:
 
 - **Latency.** The causal FIR adds a common bulk delay of `(filter_len-1)//2`

examples/bench_sampling_delay_alignment.py

@@ -0,0 +1,163 @@
+"""Benchmark SamplingDelayAlignment on short, live-acquisition-sized chunks.
+
+The transformer runs once per incoming message, so what matters is per-message
+wall time at the chunk sizes an NSP actually delivers (a few samples to a few
+hundred), not throughput on a long buffer. This script measures that, with each
+message evaluated to completion so MLX's lazy graph construction can't hide in
+the timing.
+
+It reports, per configuration:
+
+* **fir** -- alignment only (``rail_threshold=None``).
+* **rail+fir** -- with the rail forward-fill enabled, so the fill's share of the
+  message cost is visible next to the FIR it feeds.
+* **portable** -- the same transformer with its backend fast path suppressed:
+  the per-tap multiply-add loop and the log-depth rail scan. This is the
+  before/after -- the MLX depthwise conv on the ``mlx`` backend, the numba
+  kernels on ``numpy`` (shown only when numba is installed).
+
+To see the parallel FIR kernel, pass large ``--sizes`` (e.g. ``30000,300000``);
+below a few thousand samples the numba path runs its serial kernel.
+
+Usage:
+    python examples/bench_sampling_delay_alignment.py
+    python examples/bench_sampling_delay_alignment.py --channels 256 --reps 500
+    python examples/bench_sampling_delay_alignment.py --backends numpy --sizes 300,30000,300000
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import numpy as np
+from ezmsg.util.messages.axisarray import AxisArray, LinearAxis
+
+import ezmsg.blackrock.sampling_delay_alignment as sda
+from ezmsg.blackrock.sampling_delay_alignment import (
+    SamplingDelayAlignmentSettings,
+    SamplingDelayAlignmentTransformer,
+)
+
+FS = 30000.0
+RAIL = 8000.0
+
+try:
+    import mlx.core as mx
+except ImportError:  # pragma: no cover - MLX is an optional, Apple-silicon dep
+    mx = None
+
+
+def _message(data: np.ndarray, backend: str, offset: float) -> AxisArray:
+    payload = mx.array(data) if backend == "mlx" else data
+    return AxisArray(
+        data=payload,
+        dims=["time", "ch"],
+        axes={"time": LinearAxis(offset=offset, gain=1.0 / FS)},
+        key="bench",
+    )
+
+
+def _time_stream(
+    backend: str,
+    n_ch: int,
+    sizes: list[int],
+    filter_len: int,
+    rail: bool,
+    reps: int,
+    force_portable: bool = False,
+) -> float:
+    """Mean milliseconds per message, cycling through ``sizes``.
+
+    ``force_portable`` suppresses the backend fast path (numba kernels on numpy,
+    depthwise conv on MLX) so the tap loop / log-scan baseline can be timed. Each
+    message is synchronized before the clock is read, so MLX timings measure
+    execution rather than graph construction.
+    """
+    rng = np.random.default_rng(0)
+    # Hiding the numba kernels must happen before the first message (state reset
+    # is where nb_w is chosen); restore afterward so other configs still use it.
+    nb_saved = sda._nb
+    if force_portable:
+        sda._nb = None
+    try:
+        proc = SamplingDelayAlignmentTransformer(
+            settings=SamplingDelayAlignmentSettings(
+                filter_len=filter_len,
+                rail_threshold=RAIL if rail else None,
+            )
+        )
+        chunks = [
+            _message(rng.standard_normal((size, n_ch)).astype(np.float32), backend, i / FS)
+            for i, size in enumerate(sizes)
+        ]
+
+        def run_once(i: int) -> None:
+            out = proc(chunks[i % len(chunks)])
+            if backend == "mlx":
+                mx.eval(out.data)
+
+        run_once(0)
+        if force_portable and backend == "mlx":
+            # The sample shape is fixed across this stream, so the state (and its
+            # cached kernel) is not rebuilt -- dropping it here selects the
+            # tap-sum for the rest of the run.
+            proc.state.conv_w = None
+        for i in range(min(reps, 20)):  # warm up
+            run_once(i)
+
+        start = time.perf_counter()
+        for i in range(reps):
+            run_once(i)
+        return (time.perf_counter() - start) / reps * 1e3
+    finally:
+        sda._nb = nb_saved
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--channels", default="96,256", help="Comma-separated channel counts.")
+    parser.add_argument("--sizes", default="3,30,300,1000", help="Comma-separated samples/message.")
+    parser.add_argument("--filter-len", type=int, default=13, help="FIR length.")
+    parser.add_argument("--reps", type=int, default=300, help="Messages timed per configuration.")
+    parser.add_argument(
+        "--backends",
+        default="numpy,mlx",
+        help="Comma-separated backends to time (numpy, mlx).",
+    )
+    args = parser.parse_args()
+
+    channels = [int(c) for c in args.channels.split(",")]
+    sizes = [int(s) for s in args.sizes.split(",")]
+    backends = [b.strip() for b in args.backends.split(",")]
+    if "mlx" in backends and mx is None:
+        print("mlx not installed; skipping the mlx backend.")
+        backends = [b for b in backends if b != "mlx"]
+
+    numba_on = sda._nb is not None
+    print(f"filter_len={args.filter_len}  fs={FS:.0f} Hz  reps={args.reps}  (ms/message)")
+    print(f"numba: {'installed' if numba_on else 'not installed'}\n")
+    for backend in backends:
+        # A portable-baseline column wherever a fast path exists to compare to:
+        # always on MLX, on numpy only when numba is actually installed.
+        show_portable = backend == "mlx" or numba_on
+        for n_ch in channels:
+            print(f"--- {backend}, {n_ch} channels ---")
+            header = f"{'samples/msg':>12}  {'fir':>9}  {'rail+fir':>9}"
+            if show_portable:
+                header += f"  {'portable':>9}  {'speedup':>8}"
+            print(header)
+            # Each fixed size, then the mixed stream a live pipeline really sees.
+            for label, stream in [(str(s), [s]) for s in sizes] + [("mixed", sizes)]:
+                fir = _time_stream(backend, n_ch, stream, args.filter_len, False, args.reps)
+                both = _time_stream(backend, n_ch, stream, args.filter_len, True, args.reps)
+                row = f"{label:>12}  {fir:9.3f}  {both:9.3f}"
+                if show_portable:
+                    port = _time_stream(backend, n_ch, stream, args.filter_len, False, args.reps, force_portable=True)
+                    row += f"  {port:9.3f}  {port / fir:7.2f}x"
+                print(row)
+            print()
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -16,6 +16,14 @@ dependencies = [
     "scipy",
 ]
 
+[project.optional-dependencies]
+accel = [
+    # JIT-compiled kernels for SamplingDelayAlignment's numpy path (FIR and
+    # rail forward-fill). Strictly optional: without numba the portable
+    # Array-API implementation is used.
+    "numba>=0.59",
+]
+
 [dependency-groups]
 dev = [
   "typer>=0.12.5",
@@ -30,6 +38,8 @@ test = [
   "pytest",
   "pytest-timeout",
   "ezmsg-neo>=0.5.0",
+  # So CI exercises the numba-jitted SamplingDelayAlignment kernels.
+  "numba>=0.59",
 ]
 docs = [
   "sphinx>=7.0",

src/ezmsg/blackrock/_numba_kernels.py

@@ -0,0 +1,117 @@
+"""Optional numba-jitted kernels for SamplingDelayAlignment's numpy path.
+
+Imported best-effort by :mod:`ezmsg.blackrock.sampling_delay_alignment`; when
+numba is not installed the transformer falls back to its portable Array-API
+implementation, so nothing here is required. Install the ``accel`` extra to get
+it (``pip install ezmsg-blackrock[accel]``).
+
+Both operations the numpy path spends time in become a single fused pass here:
+
+* **FIR** (:func:`fir`) -- ``y[i, c] = sum_k w[k, c] * xext[i + k, c]`` with the
+  per-column taps already reversed into ``w``. The tap loop is ``filter_len``
+  vectorized multiply-adds, each a full read+write of the whole buffer; this
+  reads each input once, several times faster at live chunk sizes and, with the
+  threaded variant, ~20x faster on the large buffers offline batch processing
+  feeds through. The column (channel) loop is innermost so it stays contiguous
+  and auto-vectorizes; accumulation is in the array dtype, so the result matches
+  the tap sum within the float32 tolerance the tests pin.
+* **Rail forward-fill** (:func:`fill_rails`) -- the single left-to-right pass the
+  operation actually is, rather than the O(n log n) Hillis-Steele scan the
+  Array-API standard forces (it has no cumulative max). This is the numpy path's
+  dominant per-message cost once the FIR is fast.
+
+Both take and return 2D ``(time, column)`` arrays; the caller flattens the
+sample shape to columns (matching the MLX conv layout) and reshapes back.
+"""
+
+from __future__ import annotations
+
+import numpy.typing as npt
+from numba import njit, prange
+
+# Below this many samples the parallel FIR's thread fork-join (~0.1 ms) costs
+# more than it saves, and -- more importantly -- spawning a thread pool per
+# message inside a live ezmsg graph invites latency jitter and oversubscription.
+# Live acquisition chunks are far smaller than this, so they always take the
+# serial kernel; only the large buffers of offline batch processing parallelize.
+PARALLEL_MIN_SAMPLES = 4096
+
+
+@njit(cache=True, fastmath=True)
+def _fir_serial(xext: npt.NDArray, w: npt.NDArray, out: npt.NDArray) -> None:
+    n, n_cols = out.shape
+    n_taps = w.shape[0]
+    for i in range(n):
+        for c in range(n_cols):  # tap 0 seeds the row (avoids a separate zeroing pass)
+            out[i, c] = w[0, c] * xext[i, c]
+        for k in range(1, n_taps):
+            row = i + k
+            for c in range(n_cols):  # innermost + contiguous -> auto-vectorized
+                out[i, c] += w[k, c] * xext[row, c]
+
+
+@njit(cache=True, fastmath=True, parallel=True)
+def _fir_parallel(xext: npt.NDArray, w: npt.NDArray, out: npt.NDArray) -> None:
+    n, n_cols = out.shape
+    n_taps = w.shape[0]
+    for i in prange(n):
+        for c in range(n_cols):
+            out[i, c] = w[0, c] * xext[i, c]
+        for k in range(1, n_taps):
+            row = i + k
+            for c in range(n_cols):
+                out[i, c] += w[k, c] * xext[row, c]
+
+
+def fir(xext: npt.NDArray, w: npt.NDArray, out: npt.NDArray) -> None:
+    """Fill ``out[i, c] = sum_k w[k, c] * xext[i + k, c]`` (reversed taps in w).
+
+    ``xext`` has ``out.shape[0] + w.shape[0] - 1`` rows (the carried history
+    prepended). Picks the threaded kernel only for large buffers; see
+    :data:`PARALLEL_MIN_SAMPLES`.
+    """
+    if out.shape[0] >= PARALLEL_MIN_SAMPLES:
+        _fir_parallel(xext, w, out)
+    else:
+        _fir_serial(xext, w, out)
+
+
+@njit(cache=True)
+def _fill_rails_serial(x: npt.NDArray, thresh: float, out: npt.NDArray) -> None:
+    for c in range(x.shape[1]):
+        _fill_rails_column(x, thresh, out, c)
+
+
+@njit(cache=True, parallel=True)
+def _fill_rails_parallel(x: npt.NDArray, thresh: float, out: npt.NDArray) -> None:
+    # Columns are independent, so the forward-fill's per-column time recurrence
+    # parallelizes cleanly across channels -- which the vectorized-across-columns
+    # scan it replaces cannot do (its recurrence runs along time).
+    for c in prange(x.shape[1]):
+        _fill_rails_column(x, thresh, out, c)
+
+
+@njit(cache=True, inline="always")
+def _fill_rails_column(x: npt.NDArray, thresh: float, out: npt.NDArray, c: int) -> None:
+    last_valid = x[0, c]  # leading-rail fallback: the first sample, as-is
+    seen = False
+    for i in range(x.shape[0]):
+        v = x[i, c]
+        if v >= thresh or v <= -thresh:
+            out[i, c] = last_valid if seen else x[0, c]
+        else:
+            out[i, c] = v
+            last_valid = v
+            seen = True
+
+
+def fill_rails(x: npt.NDArray, thresh: float, out: npt.NDArray) -> None:
+    """Forward-fill (hold last valid) over railed samples, per column, in place
+    of the portable scan. Matches its semantics exactly: a run of ``|x| >=
+    thresh`` holds the last valid value, and a leading rail (nothing valid seen
+    yet) falls back to the column's first sample. Parallelizes over columns only
+    for large buffers; see :data:`PARALLEL_MIN_SAMPLES`."""
+    if x.shape[0] >= PARALLEL_MIN_SAMPLES:
+        _fill_rails_parallel(x, thresh, out)
+    else:
+        _fill_rails_serial(x, thresh, out)