Machine Learning Operator Lexicon for Partial Differential Equation Solvers on Irregular and Sparse Domains

Abstract

Solvers for grid-based discretizations of Partial Differential Equations (PDEs) are ubiquitous in computer graphics. However, efficient implementation of such solvers is contingent on specific design choices for grid data structures (especially for sparse domains), hand-tuned kernels for specific platforms (CPU vs GPU), and typically relies on the lexicon of sparse linear algebra to express its algorithmic building blocks. Meanwhile, modern hardware and software stacks are increasingly optimized around machine learning (ML) workloads, where tensor operators such as convolution, masking, and pooling form the dominant computational vocabulary. We introduce an ML-operator abstraction for discrete PDE solvers on sparse and irregular domains. Our formulation separates solver computation into uniform stencil actions, pointwise coefficient operations, masking, reductions, and localized boundary corrections, yielding a compact operator pipeline that supports practical boundary conditions and matrix-free execution. Unlike neural PDE methods, our approach does not learn a solver. It uses ML operators to express and execute established numerical discretizations. We demonstrate that this abstraction accommodates multigrid, multigrid-preconditioned Krylov solvers, and adjoint differentiation on practical 3D problems, including sparse, irregular, and inverse settings. The same operator-level view also exposes concrete backend optimization targets. Through a sparse domain optimization case study, we show that targeted improvements to a key operator can substantially improve performance over an existing baseline.