Speed up raw IO for many-small-read workloads (deep learning) - #14209
Closed
bruAristimunha wants to merge 6 commits into
Closed
Speed up raw IO for many-small-read workloads (deep learning)#14209bruAristimunha wants to merge 6 commits into
bruAristimunha wants to merge 6 commits into
Conversation
- Raw.get_data: skip full times-axis allocation when tmin/tmax unset; resolve picks=None to arange directly instead of string machinery - _picks_to_idx: fast return for integer arrays already in range - _mult_cal_one: fuse gather+cast+calibration into one pass - FIF reader: read simple numeric tags through a PID-keyed memory map; select buffer entries via searchsorted on bounds - EDF/BDF reader: vectorized uniform-sfreq fast path writing straight into the output buffer (incl. uniform stim channels via legacy bitmask), optional numba kernels for window decode and int24 samples with numpy fallbacks (all outputs bit-identical to the previous implementation) - BaseRaw._get_windows: internal batched window reader reusing one buffer Benchmarks (128ch x 1800s @512Hz, random 2s windows): EDF/BDF/BrainVision/ FIF per-call latency reduced ~4.3-7x vs main; preloaded access now at the numpy floor.
engine='edfio' parses EDF via the optional edfio package into a preloaded Raw (uniform sampling rates; all channels EEG; no meas_date). Decoding stacks digital samples once and applies calibration in two fused passes; output matches the native engine within 1 ulp.
_gdf_edf_get_fid now returns PID-keyed LRU-cached binary handles wrapped in a no-close proxy so reader context managers detach instead of closing, and reuse seeks to position 0 to preserve fresh-open semantics. Removes the per-call open/close pair on every windowed read.
Add decode_window_into numba kernel writing calibrated samples into a strided destination slice, removing the per-chunk temporary and copy when no projector/compensation is active. Byteswap big-endian chunks to native before the kernel (real-world EDF); kernel is gated to the direct path so projection handling keeps its exact legacy route.
Preloading into a memmap file now skips decoding entirely when a file with the expected size already exists, and BaseRaw.__del__ no longer deletes memmap-backed files so caches persist across sessions. Reopening a 944 MB recording drops from ~450 ms of decoding to ~15 ms, and windowed reads via the public API run at the preloaded-data floor (~84 us).
preload="memmap" now requests an automatically managed memory-map cache keyed by source path/mtime/size: the first read decodes into it, later reads (including from fresh processes) mmap it directly instead of decoding again. Directory defaults to ~/.cache/mne/memmap and can be redirected via the MNE_MEMMAP_DIR config/environment variable.
Member
|
For EDF/BDF I wonder if we can optimize our own reading such that |
Contributor
Author
|
Yes, we can @larsoner, I will split this PR into four steps, and leave the other dependencies out. |
Contributor
Author
|
Split into a dependency-stacked series of focused drafts (merge in order; each stacks on the previous):
Each branch contains only its own files/hunks and passes its relevant suites; this umbrella stays as the overview until all parts merge. |
Member
|
I think we can close this and work with the individual PRs, but feel free to reopen if I'm mistaken! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Reference issue (if any)
None. Umbrella draft; being split into #14212 → #14213 → #14214 → #14215 → #14216 (merge in that order).
What does this implement/fix?
Speeds up raw IO for workloads that make many small reads (DL training loops drawing random 2 s windows). Best of 4 interleaved runs vs
main, 128 ch × 1800 s @ 512 Hz fixtures:EDF 1472 → ~300 µs · BDF 2484 → ~406 µs · BrainVision 973 → ~200 µs · FIF 1305 → ~250 µs per window read. Preloading through a managed memmap cache drops reopening a 944 MB recording from ~450 ms to ~15 ms and puts window reads at the numpy floor (~84 µs).
Bit-identical output except where noted in the part descriptions.
Additional information
AI disclosure: I directed the work and reviewed/tested everything; Claude Code (Opus 4.1) wrote most of the code under my direction. Benchmark harness on the branch (
benchmarks/io_dl/), kept out of review scope.