Calibrate raw buffers with an optional numba kernel

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

Author: bruAristimunha

doc/changes/dev/14251.newfeature.rst

@@ -0,0 +1 @@
+Speed up raw data reading by calibrating each buffer with an optional numba kernel, by `Bruno Aristimunha`_.

mne/_fiff/_utils_numba.py

@@ -0,0 +1,20 @@
+"""Jitted helpers for :mod:`mne._fiff.utils`.
+
+Kept in its own module so that importing mne does not import numba.
+"""
+
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+from .._numba import jit
+
+
+@jit()
+def _scale_into(one, cals, data_view):
+    """Cast, transpose and calibrate one raw buffer in a single pass."""
+    n_channels, n_times = one.shape
+    for ci in range(n_channels):
+        cal = cals[ci]
+        for ti in range(n_times):
+            data_view[ci, ti] = one[ci, ti] * cal

mne/_fiff/utils.py

@@ -71,6 +71,24 @@ def _find_channels(ch_names, ch_type="EOG"):
     return eog_idx
 
 
+_SCALE_KERNEL = None
+
+
+def _get_scale_kernel():
+    """Return the jitted calibration kernel, or None if numba is unavailable."""
+    global _SCALE_KERNEL
+    if _SCALE_KERNEL is None:
+        from .._numba import has_numba
+
+        if not has_numba:
+            _SCALE_KERNEL = False
+        else:
+            from ._utils_numba import _scale_into
+
+            _SCALE_KERNEL = _scale_into
+    return _SCALE_KERNEL or None
+
+
 def _mult_cal_one(data_view, one, idx, cals, mult):
     """Take a chunk of raw data, multiply by mult or cals, and store."""
     assert data_view.shape[1] == one.shape[1], (
@@ -89,18 +107,27 @@ def _mult_cal_one(data_view, one, idx, cals, mult):
             # Benchmark (128 ch x 1024 samples): ~85 -> ~30 us per call
             # on BrainVision/FIF window reads.
             one = one[idx]
-            swapped_ints = not one.dtype.isnative and one.dtype.kind in "iu"
-            if swapped_ints and not one.flags.c_contiguous:
-                # FIFF stores sample-major, so `one` is the transpose of a
-                # big-endian integer tag and consecutive samples of a channel
-                # sit a 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 and transposes; swapping up front in one
-                # pass is cheaper. Floats are excluded because there the swap
-                # costs about what it saves, and contiguous input already has a
-                # fast loop.
-                one = one.astype(one.dtype.newbyteorder("="))
-            np.multiply(one, cals.reshape(-1, 1), out=data_view, casting="unsafe")
+            kernel = _get_scale_kernel()
+            if kernel is not None:
+                # numba refuses byte-swapped input outright, so normalize
+                # first; the kernel more than pays for the extra pass
+                if not one.dtype.isnative:
+                    one = one.astype(one.dtype.newbyteorder("="))
+                kernel(one, cals.reshape(-1), data_view)
+            else:
+                # NumPy fallback. FIFF hands us the transpose of a big-endian
+                # tag, and np.multiply has no fast loop for input that is
+                # byte-swapped *and* strided, so it swaps element by element
+                # while it also casts and transposes. Swapping up front in one
+                # pass is cheaper, but only for integers -- for floats it costs
+                # about what it saves.
+                if (
+                    not one.dtype.isnative
+                    and one.dtype.kind in "iu"
+                    and not one.flags.c_contiguous
+                ):
+                    one = one.astype(one.dtype.newbyteorder("="))
+                np.multiply(one, cals.reshape(-1, 1), out=data_view, casting="unsafe")
         else:
             one = np.asarray(one, dtype=data_view.dtype)
             np.take(one, idx, axis=0, out=data_view)