Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm

We present a unified algorithmic framework for the numerical solution, constrained optimization, and physics-informed learning of PDEs with a variational structure. Our framework is based on a Galerkin discretization of the underlying variational forms, and its high efficiency stems from a novel highly-optimized and GPU-compliant TensorGalerkin framework for linear system assembly (stiffness matrices and load vectors). TensorGalerkin operates by tensorizing element-wise operations within a Python-level Map stage and then performs global reduction with a sparse matrix multiplication that performs message passing on the mesh-induced sparsity graph. The Map and Reduce stages are co-designed inside PyTorch's autograd so that the assembly graph contains $O(1)$ nodes regardless of how the number of elements and local DoFs scale. We validate this $O(1)$-graph property by deploying TensorGalerkin downstream as i) a highly-efficient numerical PDEs solver, ii) an end-to-end differentiable framework for PDE-constrained optimization, and iii) a physics-informed operator learning algorithm for PDEs. With multiple benchmarks, including 2D and 3D elliptic, parabolic, and hyperbolic PDEs on unstructured meshes, we demonstrate that the proposed framework provides significant computational efficiency and accuracy gains over a variety of baselines in all the targeted downstream applications.

(1) Problem. From designing aircraft to simulating how heat spreads through a chip, science and engineering depend on solving partial differential equations — mathematical descriptions of how physical quantities change in space and time. Classical solvers are accurate but slow, especially in design workflows that must solve the same equation thousands of times, or when training AI models on physics data. Newer AI-based methods promised relief, but in practice, they spend most of their time on bookkeeping — looping over millions of tiny mesh pieces — and leave modern GPUs heavily underused. (2) Solution. We rebuild this bookkeeping into two clean stages that a GPU can run efficiently: a single bulk tensor calculation that handles all the physics at once, and a single sparse matrix multiplication that stitches the local pieces into one global system. This replaces millions of tiny operations with a handful of large ones, and as a side effect makes gradient computation — the ingredient AI training and inverse design rely on — nearly free. (3) Impact. Built on this engine, our fast solver outperforms widely-used scientific software by up to 50×; our physics-aware AI learns to predict equation solutions without any labeled training data and generalizes better than data-hungry alternatives; and our design tool discovers efficient mechanical structures roughly 4× faster than the leading GPU baseline. Together, these speedups make scientific simulation, AI-for-science, and computational engineering substantially cheaper on modern hardware.