FNO-Speed: Accelerate Fourier Neural Operators Training for Solving PDEs

Abstract

Deep learning for solving partial differential equations (PDEs) has become increasingly prominent. The Fourier Neural Operator (FNO) architecture has been proven to be an efficient and high-precision method that is widely used in scientific research. However, FNO incurs significant overhead by increasing the scale and dimensionality of practical problems. The insufficient utilization of hardware resources in its key operations reduces the computational efficiency of FNO solvers in high-resolution and time-sensitive problem scenarios, and cannot provide effective solution capabilities. To address the latency induced by low computational resource utilization and large-scale data access and computation, we propose FNO-Speed, an integrated solution incorporating the multi-level parallel FNO-aware mapping and tiling GEMM optimization strategy and the custom-sized high-frequency signal filtering scheme. FNO-Speed effectively leverages the data characteristics of FNO layers and the GPU hierarchical structure to adopt a data tiling and partitioning strategy, implementing matrix multiplication based on vector outer products and operator fusion to replace convolution. It also adopts a data reorganization scheme and computation restructuring to address fragmented memory access operations and serial einsum in frequency-domain. The FNO-Speed optimization strategy enhances the utilization of device memory bandwidth and computational efficiency and achieves significant acceleration in both 2D and 3D problem scenarios while maintaining nearly identical accuracy. The model achieves up to 1.4 × end-to-end training speedup, and the parallel efficiency achieves around 70% on 4 GPUs, fully demonstrating the effectiveness of the FNO-Speed optimization strategy in improving FNO performance.