TERRA-NG v1.0: Extreme-Scale, GPU-accelerated Mantle Convection

Abstract

We present TERRA-NG, a portable, GPU-accelerated, matrix-free mantle-convection code. A single Kokkos C++ implementation runs at scale on NVIDIA, AMD, and Intel GPU supercomputers. TERRA-NG has a deliberately narrow design: built on a radially extruded mesh of spherical wedges, tailored to the spherical shell geometry, which enables domain-specific optimizations like single quadrature-point integral-evaluations, radial coordinate storage compression and radial shared-memory tiling. The corresponding low-order $W_1$-iso-$W_2/W_1$ wedge-based Stokes--energy discretisation is verified against the Zhong et al.(2008) spherical-shell convection benchmark suite. We showcase TERRA-NG through strong- and weak-scaling on the JUWELS Booster (NVIDIA A100), MareNostrum 5 (NVIDIA H100), LUMI-G (AMD MI250X), Hunter (AMD MI300A APU), and SuperMUC-NG Phase 2 (Intel PVC) supercomputers. Coupled mantle convection simulations at $\sim\!11$ km and $\sim\!5.6$ km radial spacing ($\sim 2.8$ B and $\sim 22$ B DoFs) can be run routinely on standard node partitions of all considered systems. Global $\sim\!1$ km-per-gridpoint mantle convection ($\sim 1.4$ T DoFs) is feasible on an extreme-scale allocation, and a sub-km hero-run at $\sim\!0.7$ km grid spacing scaling up to $\sim 11,000$ GPUs of LUMI-G ($\sim 11$ T DoFs) shows the potential of the code on future, larger machines.