Locality-Aware Automatic Differentiation on the GPU for Mesh-Based Computations

Fig. 1. We introduce a GPU system for efficient automatic differentiation of computations defined on triangle meshes that exploits locality and sparsity in mesh-based workloads. Using our system, users specify only the energy terms of their application while our system computes gradients, sparse Hessians, and Jacobians automatically and efficiently on the GPU. Here, we use a Newton solver for large-scale elastic shell simulation to simulate ≈ 700 Spot cows (≈2.1M vertices in total) falling to the ground where derivative computation accounts for only 12.2% of the total runtime.

We present a GPU-based system for automatic differentiation (AD) of functions defined on triangle meshes, designed to exploit the locality and sparsity in mesh-based computation. Our system evaluates derivatives using perelement forward-mode AD, confining all computation to registers and shared memory and assembling global gradients, sparse Jacobians, and sparse Hessians directly on the GPU. By avoiding global computation graphs, intermediate buffers, and device-host synchronization, our approach minimizes memory traffic and enables efficient differentiation under both static and dynamically changing sparsity. Our programming model lets users express energy terms over mesh neighborhoods, while our system automatically manages parallel execution, derivative propagation, sparse assembly, and matrix-free operations such as Hessian-vector products. Our system supports both scalar-and vector-valued objectives, dynamic interaction-driven sparsity updates, and seamless integration with external GPU sparse linear solvers. We evaluate our system on applications including elastic and cloth simulation, surface parameterization, mesh smoothing, frame field design, ARAP deformation, and spherical manifold optimization. Across these tasks, our system consistently outperforms state-of-the-art differentiation frameworks, including PyTorch, JAX, Warp, Dr.JIT, EnzymeAD, and Thallo. We