Aug 2026· ACM Transactions on Architecture and Code Optimization (TACO)· 0 citations· 34 references
TL;DR
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, is proposed, fully demonstrating the effectiveness of the FNO-Speed optimization strategy in improving FNO performance.
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.
Neural operators like the Fourier Neural Operator (FNO) have demonstrated remarkable success in solving partial differential equations (PDEs) but suffer from high training costs due to fast Fourier Transform (FFT) operations. Targeting the unique computational bottleneck of the FFT in FNO, FNO–MP achieves acceleration through three synergistic designs: a progressive multi-resolution and modality growth strategy that incrementally reduces FFT overhead; a stabilized early stopping mechanism adapted to FNO’s training oscillations to adaptively terminate each stage; and mixed precision as a foundational accelerator working in concert with the above strategies. On the Navier- Stokes and Darcy flow equations, FNO–MP delivers 1.92× and 2.38× speedups, respectively, while maintaining comparable or slightly improved test errors. These results demonstrate that the synergistic alignment of precision scheduling, architectural adaptation, and optimization control can effectively unlock training efficiency in neural operators, providing a practical solution for accelerating FNO without compromising predictive accuracy.
Yi-Yang Zhu· International Conference on...· 0 citations
Two general multi-stage neural operator learning frameworks applicable when the target operator can be represented by a PDE, leveraging the weak form of the PDE residual for training are introduced.
Zhiping Mao, Zhenye Wen, Yong Zhang et al.· 0 citations
Physics-Informed Neural Networks (PINNs) solve PDEs by incorporating physical constraints into neural-network training, but large-scale problems are limited by automatic-differentiation memory overhead and inefficient execution of grid-based PDE operators. We present FlashPDE, a drop-in fused operator library for grid-based scientific machine learning. FlashPDE replaces fragmented PyTorch finite-difference execution with differentiable Triton kernels. Each operator integrates fused stencil evaluation, an analytic discrete-adjoint backward pass, and boundary-gradient correction within a unified PyTorch autograd Function interface. The library provides 14 differentiable PDE operators covering 17 configurations across 1D--3D elliptic, parabolic, and Navier--Stokes systems, while remaining independent of neural architectures and training strategies. Experiments on an NVIDIA A100 GPU show that FlashPDE reduces peak memory usage by up to 37.0x compared with coordinate-based automatic differentiation and reduces CUDA kernel launches by up to 3.5x compared with eager PyTorch finite-difference implementations. Across six representative PDE benchmarks, FlashPDE achieves up to 2.30x end-to-end time-to-solution speedup and up to 19.2x kernel-level acceleration while maintaining numerical agreement with PyTorch finite-difference references. FlashPDE provides a hardware-efficient execution layer that bridges differentiable PDE solvers and GPU-optimized numerical computation within the PyTorch ecosystem.
Field temporal prediction and source identification constitute canonical problems in dynamical systems. Conventional approaches to these problems depend on a thorough understanding of the governing partial differential equations (PDEs). Recently, deep learning, as represented by neural operators, has provided a data-driven paradigm for addressing such tasks. However, most existing global neural operators for PDEs require large training datasets and many learnable parameters, with limited interpretability and generalization. We propose the local gradient neural operator (LGNO) as a lightweight and interpretable alternative for field temporal evolution prediction and source identification in typical mechanical problems. The method builds on priors from nonlinear gradient discretization and uses multilayer perceptron convolutional layers to learn translation-invariant local kernels that resemble discrete stencils. A zero consistent stencil factorization separates coefficient learning from field reconstruction, rendering the learned operators more transparent. For problems with symmetries, network folding shares equivalent components and reduces parameter counts. We evaluate the method on PDE benchmarks covering linear and nonlinear, static and dynamic, and low and high dimensional cases. Results show that LGNO maintains accuracy, parameter efficiency, and rollout stability across these tasks, and further exhibits wide applicability to mechanical problems including diffusion, flow, and quantum phenomena.
Bai-Ming Zhang, Jin-Song Tang, Ying Xu et al.· 0 citations
Edge computing and artificial intelligence have made the efficient deployment of machine vision algorithms on low-power hardware a critical challenge for integrated circuit design. Given data-intensive image pixels and deep neural network tensors, traditional von Neumann architectures inevitably encounter severe memory and power wall bottlenecks. This paper reviews low-level hardware acceleration strategies to deconstruct the mapping from algorithm logic to silicon substrates. Through systematic analysis, physical dimensionality reduction and operator substitution effectively reduce computational complexity and enhance energy efficiency. For visual preprocessing, line-buffer topologies and multiplier-less shift-and-add designs optimize bandwidth and eliminate DSP dependencies for certain edge detectors. In edge-based inference, 8-bit fixed-point quantization slashes system bus bandwidth demands by over 75% while restricting accuracy degradation to under 0.5%. Concurrently, Winograd fast convolution translates spatial convolutions into transform-domain mappings, substantially reducing core multiplier utilization and driving significant system throughput. Additionally, advanced dataflow scheduling circumvents power-hungry off-chip memory accesses by maximizing on-chip data reuse and resolving dynamic energy crises. Ultimately, these findings provide strong guidelines for the hardware-software co-design of emerging ultra-low power edge Artificial Intelligence (AI) chips.
Linenxu Zhang· MATEC Web of Conferences· 0 citations
Deep neural operators learn mappings between input functions and complete PDE solution fields, enabling forward evaluations of new problem instances orders of magnitude faster than conventional numerical solvers. Attention mechanisms have recently been introduced into neural operators, but most studies change several architectural components at once, making it difficult to identify what actually improves accuracy. This work presents a controlled and systematic study of five deep operator network (DeepONet) variants with distinct attention mechanisms, trained under both data-driven and physics-informed regimes, to isolate the effects of cross-attention, self-attention, tokenization, and attention depth. We evaluate them on a source-driven transient one-dimensional nonlinear diffusion-reaction equation, a transient one-dimensional viscous Burgers equation with variable initial conditions, and a two-dimensional Poisson heat-conduction problem with heterogeneous source fields. Per-sensor tokenization with cross-attention reduces the mean relative L_2 error of the classical DeepONet in all benchmark-training combinations by factors of 2.4-28.0, while the best attention configurations reach 3.5-32.3. Branch self-attention paired only with dot-product fusion is inconsistent, degrading the one-dimensional problems while helping the more complex two-dimensional source field; added on top of cross-attention it improves all six cases, though by less than cross-attention fusion alone. Global pre-mixing provides no consistent benefit. Increasing cross-attention depth further improves accuracy, but with diminishing returns and a substantially higher cost under physics-informed training. Overall, query-dependent cross-attention is the most reliable mechanism, whereas branch self-attention is most useful for large, spatially complex functional inputs.
Amar Alem Koric, Qi-Bang Liu, S. Koric· 0 citations
We use cookies to run the site and, with your consent, for analytics and to show ads.
See our Cookie Policy.
