A versatile FEM framework with native GPU scalability via globally-applied AD

Abstract

Energy-based finite-element formulations provide a unified framework for describing complex physical systems in computational mechanics. In these energy-based methods, the governing equations can be obtained directly by considering the derivatives of a single global energy functional. While Automatic Differentiation (AD) can be used to automate the generation of these derivatives, current frameworks face a clear trade-off based primarily on the scale upon which the AD method is applied. Globally applied AD offers high expressivity but cannot currently be scaled to large problems. Locally applied AD scales well through traditional assembly methods, but the variety of physics and couplings that the framework can easily represent is more limited than the global approach. Here, we introduce an energy-centric framework tatva (https://github.com/smec-ethz/tatva) that defines the physics of a problem as a single global functional and applies AD globally to generate residual and tangent operators. By leveraging Jacobian-vector products for matrix-free solvers and coloring-based sparse differentiation for materializing sparse tangent stiffness matrices when needed, our flexible design scales linearly with the problem size on GPUs. We demonstrate that our framework can handle large problems (with millions of degrees of freedom) without memory exhaustion. Additionally, it offers a unified, fully differentiable methodology that can address a wide range of problems, including multi-point constraints, mixed-dimensional coupling, and the incorporation of neural networks, while maintaining high performance and scalability on modern GPU architectures.