Reuse Cache Plan

[seanlaw]

@mreineck Thank you for this package. @NimaSarajpoor and I have been using scipy.fft to compute thousands of repeated rfft/ifft calculations for (power of 2) 1-dimensional inputs in the range of 2^10 to 2^20. Essentially, I am performing a convolution (or sliding window dot product). In the aforementioned range of inputs, PocketFFT seems to be performing ~2x slower than FFTW and so I was hoping to regain some of that time back by caching/reusing the plan. However, it appears that plan caching has been disabled for scipy. So, I wanted to ask if you might be able to offer a possible scipy workaround that doesn't require recompiling Pocketfft (as this is for an existing package that cannot have additional package dependencies added)?

Fictitious Example:

import numpy as np
from scipy.fft import rfft, ifft, next_fast_len

def scipy_convolution(Q, T):
    n = len(T)
    m = len(Q)
    shape = next_fast_len(n + m - 1, real=True)

    Qraf = rfft(np.ascontiguousarray(Q[::-1]), n=shape)
    Taf = rfft(T, n=shape)
    QT = irfft(np.multiply(Qraf, Taf), n=shape)

    return QT.real[m-1:n]

if __name__ == "__main__":
    Q = np.random.rand(2**11)  # Q will always be a power of 2 and always real inputs
    T = np.random.rand(2**19)  # T will always be a power of 2 and always real inputs
    for i in range(1_000):
        scipy_convolution(Q, T)

Perhaps there is an alternative approach beyond plan caching that could help improve the performance? I've also tried np.fft but, unsurprisingly, the performance is comparable to scipy.fft. Thanks in advance!

[mreineck]

Hi @seanlaw, thanks for reaching out!

I fear that the performance discrepancy you see compared to FFTW is genuine ... pocketfft won't ever be competitive with FFTW on 1D transforms of power-of-2 lengths, because that would require implementing dedicated FFT subpasses with length 16, 32, 64 etc., which are

• very large (probably much longer than pocketfft's current code size), and
• nonportable (i.e. you need slightly different implementations on different hardware for optimal performance).

FFTW can afford this because of dynamic planning (it tries various algorithms to find the best one for the current situation); pocketfft doesn't do that, it only computes twiddle factor tables during the planning stage.
As far as I know, scipy currently does use cached plans for pocketfft, and the size of the cache is 10, so that would fit your needs perfectly. That's one more reason why I expect that the slowdown you see actually comes from the implemenation itself.

I'm not sure whether the convolutions you are carrying out in your code have to happen in a particular order, but if not, then I recommend to collect transforms with length n in an np.array of shape (ntrans,n) and then FFT all of those together with a single scipy FFT call that has the appropriate axes argument. If you present the problem to scipy in this way, pocketfft can make use of vector instructions to compute several transforms together, and you might even be able to use the workers argument for some parallel execution.

[mreineck]

Looking at the example again, I wonder whether scipy.signal.oaconvolve might be useful here?

[seanlaw]

Thank you for your prompt response @mreineck!

As far as I know, scipy currently does use cached plans for pocketfft, and the size of the cache is 10, so that would fit your needs perfectly. That's one more reason why I expect that the slowdown you see actually comes from the implemenation itself.

Okay, I had come across this closed PR but it wasn't clear from the SciPy code what value the POCKETFFT_CACHE_SIZE was set to as I could no longer track down the location of the pocketfft_hdronly.h (within SciPy) where this variable would normally be defined.

I'm not sure whether the convolutions you are carrying out in your code have to happen in a particular order, but if not, then I recommend to collect transforms with length n in an np.array of shape (ntrans,n) and then FFT all of those together with a single scipy FFT call that has the appropriate axes argument. If you present the problem to scipy in this way, pocketfft can make use of vector instructions to compute several transforms together, and you might even be able to use the workers argument for some parallel execution.

Ahh, gotcha! I understand what you mean and I will need to think about this more. @NimaSarajpoor maybe we can use this trick when you do your DAMP+MASS chunks (i.e., instead of doubling the length each time, you were chunking them up to avoid re-computing dot products).

Looking at the example again, I wonder whether scipy.signal.oaconvolve might be useful here?

Yes, we tried oaconvolve and it was around 2x slower at the aforementioned sizes (which isn't horrible) and only performed better than FFTW at longer array lengths. Additionally, @NimaSarajpoor discovered that oaconvolve falls back to fftconvolve when the two input arrays are within a length of 2^4 of each other (i.e., when the inputs are close-ish in length).

[NimaSarajpoor]

@mreineck
Thanks for sharing your input!!

@seanlaw

the convolutions you are carrying out in your code have to happen in a particular order, but if not, then I recommend to collect transforms with length n in an np.array of shape (ntrans,n) and then FFT all of those together with a single scipy FFT call that has the appropriate axes argument

I will think about it and discuss it with you probably via a different channel to not mix things here. This trick showed its advantage when we used Scipy's rfft on a 2D array recently 👍 So, I will keep this in mind.

[mreineck]

I just tested the "many-FFTs -at-once" trick, and it does give some speedup, but quite a bit less than I expected. Most likely that's because the binary scipy wheels are't compiled with AVX support to keep them portable. That's a bit of a pity.