Multithread many FFT operations in Python / NUMBA?

user_8858111 0
Viewed 234

To to a convolution / cross-correlation of different kernels on a 3D NumPy array, I want to calculate many smaller FFTs in parallel. As I found out the @njit(parallel = True) tag of NUMBA does not support the FFT / IFFT functions of SciPy or NumPy.

Is there any chance to calculate several 3D FFTs multi-threaded with NUMBA without having to implement the FFT algorithm myself? Or does the NUMBA parallel = True tag work without the @njit tag? I don't care too much about code compilation, the multithreading part is what I am really interested in.

I know that I could always use Python's build-in modules for multithreading / multiprocessing - but I am wondering if there is a more elegant solution using NUMBA for that purpose?

Tank you in advance for your help and all the best,

Valentin

numba fft numpy multithreading python
1 Answers

You cannot parallelize a code (using multiple threads like Numba does) that use any pure-Python type because of the GIL (Global Interpreter Lock). Rewriting your own FFT algorithm will likely be pretty inefficient. Indeed, FFT libraries (typically used by Python libraries) are often very optimized.

The most famous and one of the fastest is the FFTW. It generate an algorithm (possibly at runtime or ahead of time) by assembling small portions of codes regarding the parameters of the algorithm. It beats almost all carefully-optimized human implementations often by a large margin. FFTW support the computation of parallel multidimensional FFTs. Hopefully, there are Python wrappers of the library you can use.

Alternatively, if no Python wrappers are correct, you can write a simple C/C++ function calling the FFTW internally which is itself called from Python. Cython can help to do that quite easily. Note that it seems Numba @njit functions can be mixed with Cython code. This can be useful if your FFT is computed in the middle of a complex Numba @njit code.

user_129395570
Related