Calibrate raw buffers with an optional numba kernel - #14253
Calibrate raw buffers with an optional numba kernel#14253bruAristimunha wants to merge 4 commits into
Conversation
FIFF stores data sample-major, so _read_segment_file hands _mult_cal_one the
transpose of a big-endian tag: consecutive samples of one channel sit a whole
row apart. np.multiply has no fast loop for input that is byte-swapped *and*
strided, so it swaps element by element while it also casts to float64 and
transposes.
Swapping up front in one contiguous pass is cheaper. On a 133 MB int16 file
raw.get_data() goes 109.3 -> 91.9 ms (1.19x); the 13 MB shipped test_raw.fif
goes 9.27 -> 7.72 ms (1.20x).
Instrumenting the read loop shows why this is where the time is: _mult_cal_one
is 86-93% of it, while the tag reads are 7-13% and seek/reshape are noise.
The guard is narrow because the win is. Measured across every dtype and layout
the readers actually produce:
- big-endian int, strided (FIFF, EGI simple-binary): 1.28-1.33x
- big-endian int, contiguous (CTF): ~1.00x, so skipped
- big-endian float (Artemis123, float32 FIFF): 0.85-1.15x, unreliable,
so skipped -- an early version without the float exclusion regressed
Artemis123 to 0.90x
- 8-byte input: 0.52-0.76x, skipped
- native byte order: untouched
Output is bit-identical: full get_data() compared with np.array_equal across
FIFF (int16, float32), CTF, EGI, Artemis123, BrainVision and KIT.
_mult_cal_one is 86-93% of the FIFF read loop (tag I/O is 7-13%, seek and
reshape are noise). It asks a single np.multiply to byte-swap, cast to float64,
transpose and scale at once. NumPy has no fused loop for that, so it falls back
to nditer buffering; a jitted loop does the same work in one pass.
raw.get_data(), interleaved cross-process medians:
main byteswap only + kernel
FIFF 133 MB int16 109.6 ms 91.4 (1.20x) 59.5 (1.84x)
FIFF 263 MB float32 99.5 ms 101.2 (0.98x) 65.2 (1.53x)
FIFF test_raw.fif 9.1 ms 7.7 (1.18x) 4.9 (1.87x)
CTF 1.34 ms 1.35 (0.99x) 1.03 (1.31x)
The kernel also rescues float32 FIFF, which the byte-order commit alone cannot
help: swapping floats in NumPy costs about what it saves, but numba refuses
byte-swapped input outright, so the swap becomes worthwhile once the kernel
follows it.
Follows the existing numba pattern: the jitted function lives in its own
lazily-imported module, is gated on has_numba, and falls back to the NumPy path
(kept intact, including its integer byte-order fix) when numba is missing or
MNE_USE_NUMBA=false.
TRADE-OFF for review: this puts numba on the raw-reading path, so the first
get_data() in a process pays ~0.09 s, which is almost entirely 'import numba'
-- the very cost mne/_numba.py's docstring says it keeps out of mne.fixes.
Break-even is roughly 200 MB of FIFF per session; below that it is a small
regression. MNE_USE_NUMBA=false opts out. Gating the kernel on total read size
would avoid the cost for small reads but needs a size hint threaded through
_read_segment_file, which touches every reader's signature.
Verified on both paths: 910 tests pass with and without numba, and get_data()
is bit-identical to main for 37/37 readable fixtures across 18 formats.
Measuring the kernel's activation cost properly changes the picture in this
PR, so the gate it asked for is now implemented -- and the headline numbers in
the PR description need correcting.
Resolving the kernel costs ~159 ms in a fresh process and ~90 ms when numba is
already imported. That does NOT amortize across processes: with numba's on-disk
cache hit (verified: 1 cache hit, 0 misses) the artifact still has to be loaded,
and two consecutive processes both paid ~90 ms. The earlier '~0.09 s' figure was
measured with numba pre-imported and understated the cold case.
Against a measured saving of ~0.25 ms per MB of integer payload and ~0.145 ms
per MB of float payload, break-even is ~640 MB (int) / ~1.1 GB (float) cold, and
~360 MB / ~620 MB warm. _read_segment_file is the only place that knows the size
of a whole request, so the decision is made there, once per read, and the
resolved kernel (or None) is passed to _mult_cal_one. Requests below the
threshold never import numba at all -- verified.
What a user actually sees, fresh process, first get_data():
payload mne-tools#14252 only + gated kernel
13 MB 11.2 ms 10.5 ms 1.06x (kernel off)
130 MB 110.1 ms 109.2 ms 1.01x (kernel off)
433 MB 305.4 ms 302.4 ms 1.01x (kernel off)
920 MB 625.7 ms 556.2 ms 1.12x (kernel on)
The kernel only pays off when its cost is spread over several reads in one
process -- a pipeline or a notebook rather than a one-shot script:
920 MB payload, one process, cumulative
after 1 read 671 -> 604 ms 1.11x
after 3 reads 2022 -> 1480 ms 1.37x
after 5 reads 3335 -> 2345 ms 1.42x
So the 1.84x in the PR description is a warm-process number and overstates the
one-shot case; it is corrected there.
Also sets fastmath=False on the kernel. This decodes stored values rather than
approximating a computation, so reassociation buys nothing and only risks
changing them; output stays bit-identical with the kernel forced on.
910 tests pass with the gate, 236 with MNE_USE_NUMBA=false.
|
Closing this: having measured the activation cost properly, I do not think it The blocker is numba's per-process dispatcher load, and it cannot be Against a measured saving of ~0.25 ms/MB of integer payload and ~0.145 ms/MB of With an honest size gate, this is what a user actually gets on a first
It only approaches its potential when the cost is spread across several reads in That is a thin return for putting numba on MNE's most common read path plus two I also want to correct the record on this PR's original numbers: the 1.84× #14252 stands on its own and is the part worth having: 1.19–1.20× on integer The one thing that would revive this approach is avoiding the per-process |
1 participant
Stacked on #14252 — that PR carries the NumPy fallback this one keeps
intact. Please read #14252 first.
What does this implement/fix?
_mult_cal_oneis 86–93% of the FIFF read loop (boundary instrumentation in#14252). It asks a single
np.multiplyto byte-swap, cast to float64,transpose and scale at once. NumPy has no fused loop for that, so it falls back
to
nditerbuffering; a jitted loop does the same work in one pass and is3.2× faster in isolation. Hand-tiling it for cache locality made it slower
(1.8–2.4×), so the naive row-major loop is what ships.
The activation cost, measured
Resolving the kernel costs ~159 ms in a fresh process and ~90 ms when
numba is already imported. Crucially it does not amortize across
processes: with numba's on-disk cache hit — verified, 1 cache hit and 0 misses
— the artifact still has to be loaded, and two consecutive processes each paid
~90 ms.
Against a measured saving of ~0.25 ms/MB of integer payload and ~0.145 ms/MB of
float payload, break-even is ~640 MB (int) / ~1.1 GB (float) cold.
The gate
_read_segment_fileis the only place that knows the size of a whole request,so the decision is made there, once per read, and the resolved kernel (or
None) is passed to_mult_cal_one. This needs no change to any otherreader's signature, and it takes
_get_scale_kernel()out of the per-bufferloop. Requests below the threshold never import numba at all (verified).
What a user actually sees — fresh process, first
get_data()The kernel only pays off when its cost is spread over several reads in one
process — a pipeline or a notebook, not a one-shot script:
Also in this update
fastmath=Falseon the kernel. This decodes stored values rather thanapproximating a computation, so reassociation buys nothing and only risks
changing them. Output stays bit-identical with the kernel forced on.
Correctness
has_numba;import mnestill does not import numba.MNE_USE_NUMBA=false.get_data()bit-identical tomainon 37/37 readable fixtures across 18formats, on both paths, and with the kernel forced on.
My recommendation
Merge #14252; I would not merge this one as it stands. With the gate honest,
the kernel is off for every file under ~640 MB, gives 1.12× on a 920 MB
one-shot read, and only reaches ~1.4× in a long-lived process. That is a thin
return for putting numba on MNE's most common read path plus two
hardware-specific thresholds that would need revisiting per platform.
I am happy to close this and keep only #14252, or to keep it if you think the
pipeline/notebook case justifies it. Two things that would change my view: a
way to avoid numba's ~90 ms per-process dispatcher load, or a compiled kernel
shipped with MNE rather than JIT-ed.
Additional information
AI disclosure: I directed the work and reviewed and tested every change; Claude
Code (Claude Opus 5) wrote and benchmarked the kernel variants, measured the
activation cost and cache behaviour, and made the edits under my direction.