🔬 condensed matter

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

The paper introduces tatva, a versatile GPU-scalable finite element framework that applies global Automatic Differentiation to a single energy functional, enabling the efficient solution of large-scale, complex multi-physics problems through matrix-free solvers and sparse differentiation techniques.

Original authors: Mohit Pundir, Flavio Lorez, David S. Kammer

Published 2026-02-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mohit Pundir, Flavio Lorez, David S. Kammer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build a massive, complex Lego city. In the world of computer simulations (specifically Finite Element Analysis or FEM), this city represents a physical object like a bridge, a human heart, or a piece of fabric. To understand how this city behaves when you push, pull, or heat it, you need to solve a giant math puzzle.

For decades, engineers have had to choose between two ways of building this puzzle, and both had major flaws:

The "Local" Approach (The Assembly Line): Imagine building the city one tiny Lego brick at a time. You calculate the physics for one brick, then another, then another, and finally try to glue them all together.
- The Problem: Gluing millions of bricks together is messy. On modern super-fast computers (GPUs), this "gluing" process is slow and causes traffic jams because the computer has to constantly jump around memory to connect the pieces. It's efficient for simple cities, but if you want to add a weird, complex feature (like a bridge that talks to a river), the assembly line breaks down.
The "Global" Approach (The Master Blueprint): Imagine looking at the entire city as one single, giant energy blob. You calculate the physics for the whole thing at once.
- The Problem: This is incredibly flexible and easy to describe, but the math gets so huge that your computer runs out of memory instantly. It's like trying to hold a picture of the entire universe in your head; it's too big to fit.

The Breakthrough: Tatva

The authors of this paper (from ETH Zurich) have built a new framework called Tatva that solves this dilemma. They managed to have their cake and eat it too: the flexibility of the "Master Blueprint" with the speed of the "Assembly Line."

Here is how they did it, using some everyday analogies:

1. The "Batched" Kitchen (Solving the Memory Problem)

Imagine you are a chef trying to cook a meal for 10,000 people.

Old Way: You try to chop all 10,000 onions at once on one cutting board. Your board explodes, and you can't do it.
Tatva's Way: You chop the onions in small batches. You chop 50 onions, put them in a bowl, chop the next 50, and so on.
The Magic: Because they process the math in these small "batches" (groups of elements), the computer never runs out of memory, even for massive problems. It keeps the "Master Blueprint" view but only looks at a small slice of it at a time.

2. The "Smart Color-Coding" (Solving the Speed Problem)

When you need to know how the whole city reacts to a push, you usually need to build a giant "stiffness matrix" (a massive spreadsheet of how every part connects to every other part).

The Problem: Most of this spreadsheet is empty (zeros). Building the whole thing is a waste of time and space.
Tatva's Way: They use a trick called Graph Coloring. Imagine you have a map of the city. You color the houses so that no two houses of the same color touch each other.
- Because houses of the same color don't touch, you can calculate their effects all at the same time without them interfering with each other.
- Instead of calculating the whole spreadsheet one by one, they calculate all the "Red" houses, then all the "Blue" houses, etc. This turns a slow, linear task into a super-fast, parallel task that modern computers (GPUs) love.

3. The "Universal Translator" (Solving the Complexity Problem)

Usually, if you want to add a new rule to your simulation (like "this bridge is made of neural networks" or "these two parts are glued together"), you have to rewrite the entire math code.

Tatva's Way: Because they treat the whole problem as a single "Energy Score," they don't need to rewrite the math. They just add a new term to the score.
- Analogy: Imagine you are playing a video game where your goal is to maximize your "Score."
- If you want to add a new power-up (like a jetpack), you just add "Jetpack Points" to your total score formula. You don't need to rebuild the game engine.
- This allows them to mix and match wildly different physics (heat, electricity, neural networks, contact) without the computer getting confused.

Why This Matters

It's Fast: It runs at full speed on modern graphics cards (GPUs), handling millions of variables without slowing down.
It's Flexible: You can simulate weird, complex things like a bridge interacting with a neural network, or a crack spreading through a material, without writing complex, error-prone code.
It's "Future-Proof": As computers get faster and AI gets smarter, this framework can easily plug in new AI models or complex physics without needing a total overhaul.

In short: Tatva is like a new way of organizing a massive construction site. Instead of building brick-by-brick (slow and rigid) or trying to hold the whole building in your head (impossible), they build it in smart, manageable chunks that fit perfectly into the super-fast processors of today, while still allowing architects to design incredibly complex and creative structures.

1. Problem Statement

Current Automatic Differentiation (AD) frameworks for Finite Element Methods (FEM) face a fundamental trade-off between expressiveness and scalability:

Globally Applied AD: Differentiates a single global energy functional. This offers high expressivity, allowing for seamless integration of complex physics (non-local constraints, mixed-dimensional couplings, multiphysics) and a direct link to variational formulations. However, it traditionally produces dense derivative operators, leading to $O(N^2)$ memory and computational costs, making it infeasible for large-scale problems.
Locally Applied AD: Restricts differentiation to the element or quadrature level, followed by standard "scatter-add" assembly. This preserves sparsity and scales well to large problems but inherits rigid element-centric abstractions. It struggles to naturally represent global constraints, non-local interactions, and complex couplings without intrusive modifications to the assembly logic. Furthermore, the scatter-add assembly pattern relies on irregular memory access, which creates bottlenecks on modern GPU architectures.

The authors aim to resolve this trade-off by creating a framework that retains the expressiveness of global energy formulations while achieving linear $O(N)$ scalability on GPUs.

2. Methodology

The authors propose tatva, a monolithic, energy-centric framework built on the JAX ecosystem. The core methodology involves three key pillars:

A. Monolithic Global Formulation

Instead of assembling element-wise contributions, the physical system is defined as a single scalar potential energy functional $\Psi(u)$ .

Batched Computation: To prevent memory exhaustion, the global mesh is partitioned into contiguous batches. The framework evaluates the energy density $\psi$ simultaneously across all elements in a batch using SIMD (Single Instruction, Multiple Data) parallelization.
Differentiation: The residual vector $r = \nabla \Psi$ and the tangent stiffness operator $K = \nabla^2 \Psi$ are derived directly from this global scalar graph using AD, bypassing traditional element-level assembly loops entirely.

B. Dual Strategy for Tangent Operators

To handle the second-order derivatives without incurring $O(N^2)$ costs, the framework employs two complementary strategies depending on the solver requirements:

Matrix-Free Jacobian-Vector Products (JVP):
- Used for iterative solvers (e.g., Newton-Krylov) where the explicit matrix is not needed.
- Computes $K \cdot \delta u$ via forward-mode AD on the residual function.
- Cost scales linearly with $N$ and avoids materializing the matrix.
Sparse Differentiation via Graph Coloring:
- Used when an explicit sparse stiffness matrix is required (e.g., for direct solvers or eigenvalue problems).
- Utilizes distance-2 graph coloring on the mesh connectivity to partition degrees of freedom (DoFs) into independent sets (colors).
- A single JVP with a composite perturbation vector (representing a color) recovers multiple non-overlapping columns of the stiffness matrix simultaneously.
- The number of JVP evaluations required is bounded by a small constant (dependent on element type, not mesh size), maintaining $O(N)$ complexity.

C. Versatility and Generalization

Non-Variational Systems: The framework extends to non-variational problems (e.g., fluid transport) by formulating them as a scalar "virtual work" functional $W(u, v)$ , which is then differentiated.
AI Integration: Neural networks (for constitutive modeling or domain-level operators) are treated as native additive terms in the global energy sum. AD automatically handles the differentiation through the neural layers, eliminating the need for custom bridge code.

3. Key Contributions

Unified Framework: Introduces tatva, a library that unifies the expressiveness of global variational formulations with the performance of modern hardware.
Linear Scalability: Demonstrates ideal $O(N)$ scaling on GPUs for problems with millions of DoFs, overcoming the quadratic bottleneck of naive global differentiation.
Physics-Agnostic Assembly: Replaces rigid element-wise assembly with a physics-agnostic global assembly process. Complex problems (contact, mixed-dimensions, MPCs) are implemented simply by adding terms to the global energy functional.
Hardware Optimization: Leverages batched evaluation and graph coloring to eliminate atomic contention and unstructured memory access, achieving high GPU throughput.
Seamless AI Coupling: Provides a "future-proof" architecture where data-driven models (NCMs, Neural Operators) are integrated as first-class citizens within the variational framework.

4. Results

The framework was benchmarked on an NVIDIA A100 GPU against traditional scatter-add assembly methods:

Scaling: Both matrix-free JVP and sparse differentiation via coloring exhibited strict linear $O(N)$ scaling up to 10 million DoFs.
Throughput:
- Matrix-free JVP achieved peak throughputs of ~47 million DoFs/s (2D) and ~13 million DoFs/s (3D).
- Sparse differentiation maintained stable throughput (~8.7M DoFs/s in 2D), whereas traditional scatter-add assembly showed a significant drop in throughput as DoFs increased due to memory bandwidth bottlenecks.
Validation: The framework successfully solved diverse problems with high accuracy, including:
- Linear and hyperelasticity with essential boundary conditions.
- Multi-point constraints (rigid body, periodic homogenization via Lagrange multipliers).
- Contact problems (Hertzian contact) and cohesive fracture.
- Mixed-dimensional coupling (embedded fibers, cohesive interfaces).
- Multiphysics (thermal shock fracture).
- Non-variational advection-diffusion on curved manifolds.
- Integration of Neural Constitutive Models (NCM) and Neural Operator Element Methods (NOEM).

5. Significance

This work fundamentally shifts the paradigm of computational mechanics implementation:

Decoupling Physics from Implementation: By delegating derivative construction to the compiler (via AD) at the global level, researchers can focus on defining physical laws as energy functionals without worrying about assembly logic, sparsity patterns, or solver-specific constraints.
Scalability vs. Expressivity: It proves that high-level expressiveness and large-scale performance are not mutually exclusive. The framework enables the simulation of complex, coupled, and data-driven systems that were previously too difficult to formulate or too expensive to solve.
Hardware Alignment: The approach aligns naturally with modern accelerator architectures (GPUs/TPUs) by utilizing batched, vectorized operations and avoiding the irregular memory access patterns that plague traditional FEM codes.
Open Source: The release of tatva as an open-source library provides a practical tool for the community to adopt these advanced methodologies, facilitating rapid prototyping of novel physical models and AI-integrated mechanics.

In conclusion, the paper presents a robust solution to the long-standing scalability bottleneck in AD-based FEM, offering a unified, differentiable, and highly scalable platform for next-generation computational mechanics.