傅里叶神经算子（Fourier Neural Operator，简称 FNO）是一种被广泛应用于学习偏微分方程解算子的方法。不过，现有的 FNO 实现存在一些效率问题，主要体现在其傅里叶层将快速傅里叶变换（FFT）、滤波、通用矩阵乘法（GEMM）、零填充和逆快速傅里叶变换（iFFT）作为独立阶段分开执行，这就导致了频繁的内核启动和大量的全局内存访问开销，直接影响到计算效率。 针对这一问题，研究团队提出了 TurboFNO，这是一种针对 GPU 架构优化的高性能傅里叶神经算子。TurboFNO 的核心创新在于它是首个实现了完全融合的 FFT-GEMM-iFFT 内核的方案，大幅减少了内核启动次数和全局内存访问。具体来说，TurboFNO 重新设计了 FFT 和 GEMM 内核，使其性能可以与封闭源代码的 cuBLAS 和 cuFFT 相匹敌甚至更优。特别是 FFT 内核集成了高频截断、输入零填充和剪枝等功能，通过避免额外的内存拷贝内核来进一步提高效率。 为了实现 FFT 和 GEMM 工作负载的有效融合，TurboFNO 设计了一种特殊的 FFT 变体，这种变体允许单个线程块在隐藏维度上迭代，与 GEMM 中使用的 $k$-loop 对齐。此外，团队还设计了两种共享内存重排模式，确保 FFT 的输出能够高效地传递给 GEMM，并使 iFFT 可以直接从共享内存中获取 GEMM 的结果。这不仅提高了数据传输速度，还最大限度地减少了内存开销。 实验结果显示，在 NVIDIA A100 GPU 上运行时，TurboFNO 的性能比传统的 PyTorch、cuBLAS 和 cuFFT 框架高出最多 150%。这意味着在处理复杂计算任务时，TurboFNO 能够显著减少计算时间和资源消耗，提升整体计算效率。 业内人士对 TurboFNO 的评价普遍积极。他们认为，TurboFNO 不仅在理论上有突破性的进展，还在实际应用中展现出了明显的优势，特别是在大规模数值模拟和科学计算领域。此外，TurboFNO 的高效优化对于推动 GPU 在深度学习领域的应用具有重要意义。开发者团队来自国际知名的科研机构，拥有丰富的 GPU 优化经验和深厚的数学背景，为 TurboFNO 的成功提供了坚实的技术支撑。