AthenaK: A Performance-portable Version of the Athena++ Adaptive Mesh Refinement Framework

Abstract

We describe AthenaK: a new implementation of the Athena++ block-based adaptive mesh refinement framework using the Kokkos programming model. Finite volume methods for Newtonian, special relativistic, and general relativistic (GR) hydrodynamics and magnetohydrodynamics (MHD), and GR-radiation hydrodynamics and MHD, as well as a module for evolving Lagrangian tracer or charged test particles (e.g., cosmic rays) are implemented using the framework. In two companion papers, we describe (1) a new solver for the Einstein equations based on the Z4c formalism, and (2) a GRMHD solver in dynamical spacetimes also implemented using the framework, enabling new applications in numerical relativity. By adopting Kokkos, the code can be run on virtually any hardware, including CPUs, GPUs from multiple vendors, and emerging Advanced RISC Machine processors. AthenaK shows excellent performance and weak scaling, achieving over 1 billion cell updates per second for hydrodynamics in three dimensions on a single NVIDIA Grace Hopper processor. It does this with a typical parallel efficiency of 80% on 65,536 AMD GPUs on the OLCF Frontier system. Such performance portability enables AthenaK to leverage modern exascale computing systems for challenging applications in astrophysical fluid dynamics, numerical relativity, and multimessenger astrophysics.