Python / NUMBA中的多线程许多FFT操作？

Question

为了对 3D NumPy 数组上的不同内核进行卷积/互相关，我想并行计算许多较小的 FFT。我发现 NUMBA 的 @njit(parallel = True) 标签不支持 SciPy 或 NumPy 的 FFT / IFFT 函数。

有没有机会用 NUMBA 计算多个多线程的 3D FFT，而不必自己实现 FFT 算法？或者 NUMBA parallel = True 标签在没有 @njit 标签的情况下是否有效？我不太关心代码编译，多线程部分是我真正感兴趣的。

我知道我总是可以使用 Python 的内置模块进行多线程/多处理 - 但我想知道是否有更优雅的解决方案使用 NUMBA 来实现此目的？

提前感谢您的帮助，祝您一切顺利，

瓦伦丁

Answer 1

由于 GIL（全局解释器锁），您无法并行化使用任何纯 Python 类型的代码（使用像 Numba 那样的多个线程）。重写你自己的 FFT 算法可能效率很低。事实上，FFT 库（通常由 Python 库使用）通常非常优化。

最著名和最快的之一是 FFTW。它通过组装关于算法参数的一小部分代码来生成算法（可能在运行时或提前）。它经常以很大的优势击败几乎所有精心优化的人工实现。 FFTW 支持并行多维 FFT 的计算。希望有 Python wrappers 的库可供您使用。

或者，如果没有正确的 Python 包装器，您可以编写一个简单的 C/C++ 函数，在内部调用 FFTW，该函数本身是从 Python 调用的。 Cython 可以很容易地帮助做到这一点。请注意，似乎 Numba @njit 功能 can be mixed with Cython code。如果您的 FFT 是在复杂的 Numba @njit 代码中间计算的，这将很有用。