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

TL;DR

This paper introduces TensorGalerkin, a GPU-optimized framework for efficient Galerkin assembly in PDE solutions, enabling fast, accurate, and differentiable PDE solving, optimization, and learning across various PDE types.

Contribution

The paper presents a novel TensorGalerkin framework that unifies PDE solving, constrained optimization, and physics-informed learning with high efficiency on GPUs.

Findings

Significant computational efficiency gains over baselines.

Accurate solutions for 2D and 3D PDEs across multiple types.

Seamless integration into PDE solving, optimization, and learning workflows.

Abstract

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. It can be seamlessly employed 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.