Filter subsets of channels at a time to reduce peak memory

[h-mayorquin]

Follows the discussion in #4813 (Filter in float32 if data is not float64) and #4837 (Implement filtering by banks of channels to reduce RAM).
sosfiltfilt and filtfilt make several float64 copies of their input, so filtering a 1 GiB int16 trace peaks at 10 times its size. This PR filters one channel at a time in FilterRecordingSegment.get_traces and writes each result into a preallocated output, which takes the peak down to 1.02 times the trace size (the output array is almost all of it). The output is identical to main for the sos and ba forms, forward-backward and causal.

I did not add a parameter for how many channels to filter at a time, as I don't want to get into how to expose it in this PR. Larger blocks are a bit faster (best is 8 to 16 channels, 9.7 s against 12.7 s) and use more memory, as the figure shows, but the gain is small next to the memory gain here. I would merge this and then discuss separately a common interface to run preprocessing operations per channel, as that needs more research first. Basically, I don't want to get us bogged down discussing an interface where there is an obvious gain to be picked.

sosfiltfilt, float64 coefficients, default bandpass (5th order Butterworth, 300 to 6000 Hz), 1 GiB int16 trace (384 channels, 46.6 s at 30 kHz). The peak includes the output array, the time is the best of 3 runs.

The ba form (filtfilt) behaves the same way on the same trace, one channel at a time against all channels at once:

filter	peak memory (% of trace)	time
bandpass ba, main	1000%	26.3 s
bandpass ba, this PR	102%	15.1 s
notch (iirnotch), main	1000%	20.4 s
notch (iirnotch), this PR	102%	13.4 s

[JoeZiminski]

cool! that's a great improvement. I'm surprised looping over channels is faster than numpy's internal axis handling. I guess like you show there is some speedup for lower channel counts but for large counts something is degrading in numpys implementation. Still surprised a python loop is faster than whatever they are doing!

very cool, I also agree I don't think there's much to be gained by exposing num channels in the block through the API for that small speed up.

[h-mayorquin]

The thing is that their functions loop inside. Usually, python operations are slower but the loop itself is not that bad. My hypothesis is that we are comparing a c loop vs a python loop and there is not much difference and that the speedup comes from avoiding the big temporary allocations.

[alejoe91]

@h-mayorquin Just a tiny comment / suggestion

[h-mayorquin]

One final thing I just thought about before waking up. I did not measure using really small chunks (time chunks) as you guys tend to use. It is me who uses really large chunks.

[alejoe91]

what chunk size did you use for the plots? Can you try to run it with 1s chunks? I guess things should be similar no?

[h-mayorquin]

The chunk size was 1 GiB trace which is what I have used on the past.

I did some benchmarks closer to our defaults and I would like to delay merging this and/or keep it out of the release.

The default is chunk_duration="1s", which is 23 MB for 384 channels at 30 kHz. This is sosfiltfilt with the default bandpass on 384 int16 channels, best of 3, as time compared with main (all channels in one call), for 1 channel per call (this PR) and for blocks of 4, 16 and 64 channels:

So what do I think:

• Note that the RAM savings in percentage are the same on every chunk size but they start to matter once the chunks become large.
• Right now the filter uses 1000% of the chunk size in RAM, so a 1 GiB chunk takes about 10 GiB. I think it is fair to say that large chunks are not viable when the pipeline includes the filter. My feeling is that this is mostly historic as the main users of the pipeline mostly work with small chunks but I might be wrong about this @alejoe91 @samuelgarcia .
• Enabling large chunks is important for writing or for operations vectorized across times where one would expect the benefits of large chunks to accrue. One important case is loading large amounts of data for sharding.
• That said, enabling those RAM savings will make the common workflow slower: on the default chunk, filtering that takes 1 s now will take about 1.7 s. That is why I want to defer. I think we should benchmark the speed of complete pipelines including writing so we can make an informed decision about this.

[alejoe91]

I think it's fair to postpone to 0.106.0, since this will be mainly impactful on Zarr writes + sharding.

Given your initial benchmark and the 1s one, it seems that 16 channels banks could be a good default? I don't mind exposing the new parameter and set a good default honestly...

@h-mayorquin does that sound fair?

Still under discussion

[galenlynch]

@alejoe91 16-channel banks look like a good default to me too. My numbers agree with @h-mayorquin's plot: banks of 16 are slightly faster than main even on one thread, and keep peak RAM at about 140% of the chunk instead of 900%. One more option, which I would prefer: sosfiltfilt releases the GIL, so a thread pool can filter the banks in parallel within a single get_traces call, about 5x faster than main at about 400% peak RAM instead of 900%. That fits the trade #4564 measured: fewer chunk workers, each running threads, matched the speed of many workers at a fraction of the total RAM.

Default bandpass, 384 int16 channels, best of 3. Peak RAM is above the chunk, as a percentage of its size:

chunk	variant	time	vs main	peak RAM
1 s	main (all channels)	155 ms	1.00x	903%
1 s	one channel (this PR)	234 ms	0.66x	103%
1 s	banks of 16	135 ms	1.15x	138%
1 s	banks of 16, 8 threads	29 ms	5.4x	375%
10 s	main (all channels)	2002 ms	1.00x	900%
10 s	one channel (this PR)	2803 ms	0.71x	102%
10 s	banks of 16	1617 ms	1.24x	138%
10 s	banks of 16, 8 threads	423 ms	4.7x	400%
20 s	main (all channels)	4707 ms	1.00x	900%
20 s	one channel (this PR)	6088 ms	0.77x	102%
20 s	banks of 16	4254 ms	1.11x	138%
20 s	banks of 16, 8 threads	1080 ms	4.4x	400%

Output is identical to main in every case. Single-channel banks should not be threaded: on 8 threads at 1 s they were slower than main, since the per-task overhead outweighs the work.

Threading the banks means submitting each one to a pool instead of looping over them, about ten lines. The thread count could come from max_threads_per_worker, default 1, so nothing changes unless it is set. I'd suggest that as a follow-up rather than adding it here, since how the thread count reaches a step deserves its own review: passed as an argument, it gets dropped by any step in a chain that doesn't forward it (as was discussed in #4564).

[alejoe91]

@galenlynch thanks for the additional benchmark! I agree with the 16-channel default and with following up with the threading per bank. This is a natural follow up and easier implementation than #4564 and could be implemented per processor.

@h-mayorquin can you make num_channels_per_bank an argument with 16 by default. I think the speedup and RAM usage is the best tradeoff here :)

[galenlynch]

@h-mayorquin here is the whole-pipeline benchmark you asked for: phase_shift → bandpass_filter → common_reference(operator="median"), saved to binary with 1 s chunks and the process engine. Input and output are on tmpfs, so this measures compute and leaves storage out. Median of 3 runs on 30 s of 384 int16 channels:

variant	MS/s	vs main	peak RAM
main, 8 jobs	69	1.00x	3.7 GB
one channel (this PR), 8 jobs	71	1.03x	3.7 GB
banks of 16, 8 jobs	75	1.09x	3.7 GB
main, 2 jobs	30	1.00x	1.1 GB
banks of 16, 2 jobs	31	1.03x	1.1 GB

In the full pipeline this PR costs no speed, and banks of 16 are 9% faster. Peak RAM is the same for all three, so another step sets that peak usage, likely the FFT in phase_shift.

Threading the banks helps less here than on the filter alone: 2 jobs × 4 threads gives 36 MS/s against 31, because only the bandpass uses the threads.

[h-mayorquin]

I changed the internal bank size to 16 channels as I think it is an overall improvement. The full-pipeline results address my concern about slowing down the common workflow (@galenlynch, thanks for running these!).

I also started the broader preprocessing survey. My first isolated filtering measurements show about 83% less incremental peak memory with banks of 16, but filtering took 3-8% longer for the tested 1-second and 10-second chunks. That is a different result from your measurements, so I would not claim a universal speed improvement. Complete-pipeline and storage measurements on my machine are running right now.

I am still wary of exposing num_channels_per_bank. As I argued in #4837 (preprocessing bank survey), I would like to understand the requirements across preprocessing before committing to a public interface. I think APIs should do the right thing without requiring user input. Every argument adds complexity, so I would like a clear user benefit before exposing it. What would change my mind is finding workflows that need substantially different tradeoffs between memory use and processing speed. Until then, I think we can keep improving the internal defaults.

As an example of why I don't want to get into these details yet, this parameter controls execution during get_traces(), not the filter itself. I am not sure the constructor is the right place for it. Even the name needs discussion: "bank" already refers to electrode selections in Neuropixels and collections of filters in signal processing. I would like to understand which execution controls users need across preprocessing before settling their placement and names.

[galenlynch]

On your question about which execution controls preprocessing needs: I opened #4857 with one idea for threads inside a step. It's set from the existing max_threads_per_worker, and the issue has measurements of the memory/speed trade-off. I'd be glad to hear whether it fits what the survey finds.

[alejoe91]

Can we define this at the top as NUM_CHANNELS_PER_BANK?

[alejoe91]

As an example of why I don't want to get into these details yet, this parameter controls execution during get_traces(), not the filter itself. I am not sure the constructor is the right place for it. Even the name needs discussion: "bank" already refers to electrode selections in Neuropixels and collections of filters in signal processing. I would like to understand which execution controls users need across preprocessing before settling their placement and names.

What about NUM_CHANNELS_PER_BLOCK? I agree "bank" is confusing, even more in the context of filters as filter banks are goups of filters targeting different frequency ranges