Read FIF simple numeric tags through a memory map

FIF raw segments are read as byte-offset views into a PID-keyed memory
map of the file instead of open/seek/read per call; buffer entries are
selected with searchsorted on the sorted bounds. gzip, file-like objects,
and non-simple tag types keep the legacy path. The generic memory-map
cache in _read_segments_file also serves the other binary readers.

Author: bruAristimunha

mne/_fiff/_mmap_cache.py

@@ -0,0 +1,38 @@
+"""PID-keyed memmap caching for direct byte-offset reads.
+
+Used by readers that need random access into raw data files (currently the
+FIF raw reader). Keyed by PID so forked worker processes (e.g., PyTorch
+DataLoader workers) create their own mapping instead of sharing a parent's,
+and validated against file size/mtime so stale mappings are never reused.
+"""
+
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors
+
+import os
+
+import numpy as np
+
+_MAX_CACHE = 16
+_cache = {}
+
+
+def get_u8_memmap(path):
+    """Return a uint8 memmap of *path* (PID-keyed), or None on any failure."""
+    try:
+        st = os.stat(path)
+        key = (os.getpid(), str(path))
+        hit = _cache.get(key)
+        if hit is not None:
+            mm, mtime_ns, size = hit
+            if mtime_ns == st.st_mtime_ns and size == st.st_size:
+                return mm
+            _cache.pop(key, None)
+        mm = np.memmap(str(path), dtype=np.uint8, mode="r")
+    except Exception:
+        return None
+    _cache[key] = (mm, st.st_mtime_ns, st.st_size)
+    while len(_cache) > _MAX_CACHE:
+        _cache.pop(next(iter(_cache)))
+    return mm

mne/_fiff/utils.py

@@ -84,14 +84,15 @@ def _mult_cal_one(data_view, one, idx, cals, mult):
     else:
         assert cals is not None
         if isinstance(idx, slice):
-            # Hot path: gather + type-cast + calibration in a single pass
-            # (was three passes plus a full float64 temporary).
-            # Benchmark (128 ch x 1024 samples): ~85 -> ~30 us per call
-            # on BrainVision/FIF window reads.
+            # Hot path: gather + type-cast + calibration in a single pass,
+            # without materializing an intermediate float64 copy of `one`
+            # (`one[idx]` is a view for basic slices). Numerically identical
+            # to cast-then-scale because both are elementwise.
             np.multiply(one[idx], cals.reshape(-1, 1), out=data_view,
                         casting="unsafe")
         else:
             one = np.asarray(one, dtype=data_view.dtype)
+            # faster than doing one = one[idx]
             np.take(one, idx, axis=0, out=data_view)
             data_view *= cals
 
@@ -220,6 +221,8 @@ def _read_segments_file(
     if n_channels is None:
         n_channels = raw._raw_extras[fi]["orig_nchan"]
 
+    import os as _os
+
     n_bytes = np.dtype(dtype).itemsize
     # data_offset and data_left count data samples (channels x time points),
     # not bytes.
@@ -229,6 +232,49 @@ def _read_segments_file(
     # Read up to 100 MB of data at a time, block_size is in data samples
     block_size = ((int(100e6) // n_bytes) // n_channels) * n_channels
     block_size = min(data_left, block_size)
+
+    # Reuse a memory map across calls (keyed by PID so forked processes --
+    # e.g., PyTorch DataLoader workers -- create their own mapping instead of
+    # sharing one). This removes the per-call open/seek/syscall overhead.
+    ex = raw._raw_extras[fi] if fi < len(raw._raw_extras) else {}
+    mm = ex.get("_mm") if isinstance(ex, dict) else None
+    if mm is not None and ex.get("_mm_pid") != _os.getpid():
+        mm = None
+    if mm is not None and (
+        mm.dtype != np.dtype(dtype) or mm.size * n_bytes < data_offset + data_left * n_bytes
+    ):
+        mm = None
+    if mm is None and isinstance(ex, dict):
+        try:
+            mm = np.memmap(raw.filenames[fi], dtype=dtype, mode="r")
+            ex["_mm"] = mm
+            ex["_mm_pid"] = _os.getpid()
+        except Exception:
+            mm = None
+
+    if mm is not None:
+        base_idx = data_offset // n_bytes
+        for sample_start in np.arange(0, data_left, block_size) // n_channels:
+            count = min(block_size, data_left - sample_start * n_channels)
+            block = mm[
+                base_idx + sample_start * n_channels :
+                base_idx + sample_start * n_channels + count
+            ]
+            if block.size != count:
+                raise RuntimeError(
+                    f"Incorrect number of samples ({block.size} != {count}), "
+                    "please report this error to MNE-Python developers"
+                )
+            block = block.reshape(n_channels, -1, order="F")
+            n_samples = block.shape[1]
+            sample_stop = sample_start + n_samples
+            if trigger_ch is not None:
+                stim_ch = trigger_ch[start:stop][sample_start:sample_stop]
+                block = np.vstack((block, stim_ch))
+            data_view = data[:, sample_start:sample_stop]
+            _mult_cal_one(data_view, block, idx, cals, mult)
+        return
+
     with open(raw.filenames[fi], "rb", buffering=0) as fid:
         fid.seek(data_offset)
         # extract data in chunks

mne/io/fiff/raw.py

@@ -9,10 +9,11 @@
 
 import numpy as np
 
+from ..._fiff._mmap_cache import get_u8_memmap
 from ..._fiff.constants import FIFF
 from ..._fiff.meas_info import read_meas_info
 from ..._fiff.open import _fiff_get_fid, _get_next_fname, fiff_open
-from ..._fiff.tag import _call_dict, read_tag
+from ..._fiff.tag import _call_dict, _simple_dict, read_tag
 from ..._fiff.tree import dir_tree_find
 from ..._fiff.utils import _mult_cal_one
 from ...annotations import Annotations, _read_annotations_fif
@@ -403,13 +404,84 @@ def _dtype(self):
     def _read_segment_file(self, data, idx, fi, start, stop, cals, mult):
         """Read a segment of data from a file."""
         n_bad = 0
-        with _fiff_get_fid(self._raw_extras[fi]["filename"]) as fid:
-            bounds = self._raw_extras[fi]["bounds"]
-            ents = self._raw_extras[fi]["ent"]
-            nchan = self._raw_extras[fi]["orig_nchan"]
-            use = (stop > bounds[:-1]) & (start < bounds[1:])
+        bounds = self._raw_extras[fi]["bounds"]
+        ents = self._raw_extras[fi]["ent"]
+        nchan = self._raw_extras[fi]["orig_nchan"]
+        fname = self._raw_extras[fi]["filename"]
+        # Entries overlapping [start, stop) via binary search on sorted bounds
+        # (O(log n) instead of a mask over every entry; matters for long
+        # recordings with thousands of buffer entries).
+        eis = range(
+            max(np.searchsorted(bounds, start, side="right") - 1, 0),
+            min(np.searchsorted(bounds, stop, side="left"), len(bounds) - 1),
+        )
+
+        # Fast path: read tag payloads directly through a PID-keyed memory map,
+        # skipping per-call open/seek/read syscalls. Only taken for uncompressed
+        # real files whose touched tags are simple numeric types with the
+        # expected sizes; everything else falls back to the legacy loop below.
+        mm = None
+        if (
+            isinstance(fname, Path)
+            and len(fname.suffixes) > 0
+            and fname.suffixes[-1] != ".gz"
+        ):
+            mm = get_u8_memmap(fname)
+        if mm is not None:
+            for ei in eis:
+                ent = ents[ei]
+                if ent is None or ent.type not in _simple_dict:
+                    mm = None
+                    break
+                nsamp_ei = bounds[ei + 1] - bounds[ei]
+                itemsize = np.dtype(_simple_dict[ent.type]).itemsize
+                if getattr(ent, "size", None) != nsamp_ei * nchan * itemsize:
+                    mm = None
+                    break
+        if mm is not None:
+            offset = 0
+            for ei in eis:
+                first = bounds[ei]
+                last = bounds[ei + 1]
+                nsamp = last - first
+                ent = ents[ei]
+                first_pick = max(start - first, 0)
+                last_pick = min(nsamp, stop - first)
+                picksamp = last_pick - first_pick
+                this_start = offset
+                offset += picksamp
+                this_stop = offset
+                if ent is None:
+                    continue  # gaps were zero-initialized by the caller
+                dtype_s = _simple_dict[ent.type]
+                itemsize = np.dtype(dtype_s).itemsize
+                nbytes = picksamp * nchan * itemsize
+                base = ent.pos + 16 + first_pick * nchan * itemsize
+                one = np.frombuffer(
+                    mm[base : base + nbytes], dtype=dtype_s, count=picksamp * nchan
+                )
+                if one.size != picksamp * nchan:
+                    n_bad += picksamp
+                    continue
+                one = one.reshape(picksamp, nchan)
+                _mult_cal_one(
+                    data[:, this_start:this_stop],
+                    one.T,
+                    idx,
+                    cals,
+                    mult,
+                )
+            if n_bad:
+                warn(
+                    f"FIF raw buffer could not be read, acquisition error "
+                    f"likely: {n_bad} samples set to zero"
+                )
+            assert offset == stop - start
+            return
+
+        with _fiff_get_fid(fname) as fid:
             offset = 0
-            for ei in np.where(use)[0]:
+            for ei in eis:
                 first = bounds[ei]
                 last = bounds[ei + 1]
                 nsamp = last - first