AthenaK: A Performance-Portable Version of the Athena++ AMR Framework

TL;DR

AthenaK is a new performance-portable AMR framework based on Athena++, utilizing Kokkos to run efficiently across diverse hardware, enabling advanced astrophysical simulations with high scalability and performance.

Contribution

This work introduces AthenaK, a Kokkos-based, performance-portable implementation of the Athena++ AMR framework supporting diverse physics and hardware architectures.

Findings

Achieves over one billion cell updates per second on a single GPU.

Maintains about 80% parallel efficiency on 65,536 GPUs.

Demonstrates excellent scalability and performance portability across hardware.

Abstract

We describe AthenaK: a new implementation of the Athena++ block-based adaptive mesh refinement (AMR) framework using the Kokkos programming model. Finite volume methods for Newtonian, special relativistic (SR), 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 ARM processors. AthenaK shows excellent performance and weak scaling, achieving over one billion cell updates per second for hydrodynamics in three-dimensions on a single NVIDIA Grace Hopper processor and with a typical parallel efficiency of 80% on 65536 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.