发表机构
School of Mathematical Sciences, Beijing Normal University; Beijing University of Posts and Telecommunications（北京师范大学数学科学学院; 北京邮电大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
研究针对物理信息神经网络求解偏微分方程时的局限，提出FlashPDE库，用可微Triton内核取代PyTorch有限差分执行，集成多种功能。实验表明其能大幅减少内存使用和内核启动次数，实现加速，为PDE求解与GPU计算搭建高效执行层。
AI 中文摘要
物理信息神经网络（PINNs）通过将物理约束纳入神经网络训练来求解偏微分方程（PDEs），但大规模问题受自动微分内存开销和基于网格的PDE算子执行效率低的限制。我们提出了FlashPDE，一个用于基于网格的科学机器学习的即插即用融合算子库。它用可微的Triton内核取代了碎片化的PyTorch有限差分执行。每个算子在统一接口内集成了融合模板评估、解析离散伴随反向传播和边界梯度校正。该库提供14个可微PDE算子，涵盖1D - 3D椭圆、抛物和纳维 - 斯托克斯系统的17种配置，且独立于神经架构和训练策略。在NVIDIA A100 GPU上的实验表明，FlashPDE与基于坐标的自动微分相比，峰值内存使用最多可减少37.0倍，与急切的PyTorch有限差分实现相比，CUDA内核启动最多可减少3.5倍。在六个代表性PDE基准测试中，FlashPDE实现了高达2.30倍的端到端求解时间加速和高达19.2倍的内核级加速，同时与PyTorch有限差分参考保持数值一致。FlashPDE提供了一个硬件高效的执行层，在PyTorch生态系统中架起了可微PDE求解器和GPU优化数值计算之间的桥梁。
英文摘要
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.
Comments9 pages of main text, 4 figures. Code available at https://github.com/factnn/FlashPDE
