Reproduce “TurboFNO: High-Performance Fourier Neural Operator with Fused FFT-GEMM-iFFT on GPU” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Shixun Wu, Yujia Zhai, Huangliang Dai, Hairui Zhao, Yue Zhu, Haiyang Hu, Zizhong Chen
TurboFNO fuses FFT, complex matrix multiplication, and inverse FFT in Fourier Neural Operator (FNO) layers in one GPU kernel, keeping filtered data local. Kernels incorporate truncation and padding, while shared-memory layouts connect the stages. On NVIDIA A100 tests, authors report 67% average performance improvement and up to 150% speedup over PyTorch for 2D FNO configurations; this measures kernel performance, not PDE solving.
Key insights
- Kernel fusion significantly reduces memory traffic and kernel launch overhead
- 2D FNO implementations show greater benefits due to being more memory-bound
- Architecture-specific optimizations like shared memory swizzling are crucial for performance
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.