GPU 上的高性能融合 FFT-GEMM-iFFT 傅里叶神经算子：TurboFNO
分布式、并行与集群计算
2025-04-17 v1
摘要
傅里叶神经算子 (FNO) 广泛用于学习偏微分方程解算算子。然而，FNO 缺乏面向体系结构的优化，其傅里叶层将 FFT、滤波、GEMM、零填充与 iFFT 执行为独立阶段，导致多次内核启动及显著的全局内存开销。我们提出 TurboFNO，首个实现 FFT-GEMM-iFFT 全链路融合的 GPU 内核，并内置 FFT 优化。我们从零开发 FFT 与 GEMM 内核，性能相当或优于闭源 SOTA 的 cuBLAS 与 cuFFT。此外，FFT 内核集成高频截断、输入零填充及裁剪功能，避免额外的内存拷贝内核。为融合 FFT 与 GEMM 负载，我们提出一种单线程块遍历隐藏维度的 FFT 变体，与 GEMM 中的 -循环对齐。我们设计两种共享内存洗牌方案，实现 FFT 输出到 GEMM 时 100% 内存 bank 利用率，并使 iFFT 能直接从共享内存中获取 GEMM 结果。在 NVIDIA A100 GPU 上的实验结果表明，TurboFNO 对 PyTorch、cuBLAS 与 cuFFT 性能提升最高可达 150%。
引用
@article{arxiv.2504.11681,
title = {TurboFNO: High-Performance Fourier Neural Operator with Fused FFT-GEMM-iFFT on GPU},
author = {Shixun Wu and Yujia Zhai and Huangliang Dai and Hairui Zhao and Yue Zhu and Haiyang Hu and Zizhong Chen},
journal= {arXiv preprint arXiv:2504.11681},
year = {2025}
}