Finite-element solvers become differentiable graph programs: one H200 runs a step up to 181 times faster than 32 CPU cores
Synopsis
A finite-element solver for conjugate heat transfer and buoyant flow is recast as fixed-weight graph programs inside an automatic-differentiation framework, with assembled operators and the algebraic-multigrid cycle executed as sparse graph operations and nothing trained. With discretization, mesh, time step and tolerances held fixed, the same problem runs up to 181 times faster per step on one H200 than on 32 CPU cores, and the pressure solve alone 18.7 times faster. Reverse-mode gradients match central differences, and implicit differentiation holds one gradient near 4.9 forward solves.
Interpretation
Conjugate heat transfer, which couples buoyant fluid flow to conduction in an adjacent solid, can now run a whole time step on GPUs end to end differentiably, without a separately written adjoint. Earlier routes either stayed with sparse linear solves on CPUs or expressed operators as convolutions on Cartesian grids, which do not naturally handle body-fitted meshes and curved solid-fluid interfaces. In a comparison holding discretization, mesh, time step, step count and pressure tolerance fixed and changing only the linear-algebra backend and device, the per-step CPU/GPU ratio for the 2D conjugate cavity rises from 1.25 to 23.9, and reaches 181 on the 3D conjugate heat sink at 1.2×107 coupled degrees of freedom.
The solver represents assembled finite-element operators and the multigrid cycle as fixed-weight graph operations, with degrees of freedom as nodes, nonzero matrix couplings as edges, and weights set by the discretization. The weights are not trained; the graph neural network describes the execution structure, and the pressure solve uses a graph U-Net built from algebraic multigrid as a preconditioner inside conjugate gradients. Manufactured solutions recover the theoretical orders in three dimensions, about 2.95 for velocity and temperature and 2.03 for pressure, within 0.06 of the theoretical {3,3,2}; the graph V-cycle matches a reference implementation to fp64 round-off.
The whole chain is differentiable in reverse mode, giving the sensitivity of the hot-wall Nusselt number to the solid-fluid conductivity ratio without a hand-derived adjoint. Differentiating the truncated iteration approaches the converged-solution sensitivity only slowly, and its memory grows with the iteration count. At the prescribed 40 forward iterations per solve, implicit differentiation is about 5×10^-4 away from the converged sensitivity at conductivity ratio 100, where the unrolled gradient is 25% away; implicit peak memory stays at 2.7 GiB against 41.6 GiB for the unrolled path at 320 iterations.
Perspective
The demonstrated scope is two- and three-dimensional Boussinesq natural convection, the conjugate cavity, a turbulent Rayleigh-Bénard cell, a branched heat sink and a single-cylinder engine geometry, all in double precision, with distributed results limited to one node of eight H200 GPUs. The authors label same-discretization comparisons (Tier B) separately from complete solver-stack comparisons (Tier C), and only the former supports causal statements about the execution device.
The authors state that the two-level coarse correction is replicated on every rank, viable on one node but reintroducing at multi-node scale the memory cost that deflation avoids; unrolled differentiation memory grows about 2.6 times per mesh doubling, so the conjugate thermal study stops at 9.4×104 degrees of freedom. The differentiable demonstration is single-GPU because in-place halo exchange breaks the computational graph, and strong scaling reaches only about 15% parallel efficiency on eight GPUs. One benchmark cell has the CPU finite-element reference 0.35% below the other three solvers, reported as an unresolved outlier.