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

2025-04-16

AI Summary

Fourier Neural Operators (FNOs) suffer from high memory bandwidth pressure and frequent kernel launch overhead on GPUs due to the sequential, non-fused execution of FFT–GEMM–iFFT stages in PDE solving. Method: This work introduces the first end-to-end fused GPU kernel for FNOs, featuring an architecture-aware, fully integrated FFT–GEMM–iFFT computational flow. It incorporates high-frequency truncation, zero-padding, and channel pruning, while achieving 100% shared memory bank utilization via swizzling and thread-block remapping. Custom high-performance FFT and GEMM micro-kernels are designed with optimized global memory access patterns. Results: On NVIDIA A100, our implementation delivers up to 150× speedup over PyTorch with cuBLAS and cuFFT, significantly reducing global memory traffic and kernel launch overhead.

Technology Category

• Machine Learning: Kernel Methods
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Non-convex Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

Fourier Neural Operators (FNO) are widely used for learning partial differential equation solution operators. However, FNO lacks architecture-aware optimizations,with its Fourier layers executing FFT, filtering, GEMM, zero padding, and iFFT as separate stages, incurring multiple kernel launches and significant global memory traffic. We propose TurboFNO, the first fully fused FFT-GEMM-iFFT GPU kernel with built-in FFT optimizations. We first develop FFT and GEMM kernels from scratch, achieving performance comparable to or faster than the closed-source SOTA cuBLAS and cuFFT. Additionally, our FFT kernel integrates a built-in high-frequency truncation, input zero-padding, and pruning feature to avoid additional memory copy kernels. To fuse the FFT and GEMM workloads, we propose an FFT variant in which a single thread block iterates over the hidden dimension, aligning with the $k$-loop in GEMM. Additionally, we design two shared memory swizzling patterns to achieve 100% memory bank utilization when forwarding FFT output to GEMM and enabling the iFFT to retrieve GEMM results directly from shared memory.Experimental result on an NVIDIA A100 GPU shows TurboFNO outperforms PyTorch, cuBLAS, and cuFFT by up to 150%.

Problem

• Optimize Fourier Neural Operator for GPU performance
• Fuse FFT-GEMM-iFFT operations to reduce kernel launches
• Minimize global memory traffic in FNO computations

Innovation

• Fully fused FFT-GEMM-iFFT GPU kernel
• Built-in FFT optimizations and truncation
• Shared memory swizzling for 100% utilization