[Paper Review] AthenaK: A Performance-Portable Version of the Athena++ AMR Framework

James M. Stone, Patrick D. Mullen+7 more|arXiv (Cornell University)|Sep 24, 2024

Medical Imaging Techniques and Applications5 citations

TL;DR

AthenaK provides a new Kokkos-based, performance-portable AMR framework built on Athena++, with diverse fluid and GR solvers, enabling exascale-scale computations on CPUs and GPUs. It demonstrates strong performance and scalability across architectures.

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.

Motivation & Objective

• Introduce AthenaK, a performance-portable AMR framework implemented with Kokkos for broad hardware compatibility (CPUs, GPUs, ARM).
• Provide a complete suite of fluid and GR solvers including hydrodynamics, MHD, and GR-radiation transport within a block-based AMR structure.
• Describe framework design choices (MeshBlockPack, device-side data, task lists, boundary handling) and open-source development model.
• Demonstrate performance and scaling across architectures, validating portability and efficiency on modern HPC systems.

Proposed method

• Rewrite the Athena++ AMR framework from scratch in the Kokkos programming model to achieve performance portability across architectures.
• Organize MeshBlock data on a device as a MeshBlockPack to enable parallel operations with reduced kernel launches.
• Store physics module data as device-resident Kokkos Views, with host controlling mesh structure and dynamic task execution via a task list.
• Implement block-based AMR with boundary communication and AMR-prolongation/restriction performed within parallel kernels.
• Provide a wide range of solvers: Newtonian, special relativistic (SR), general relativistic (GR) hydrodynamics/MHD, and GR-radiation transport; include a particle module for Lagrangian tracers and charged particles.
• Use cmake for builds; runtime-configurable options via input files; support templating to reduce register pressure on GPUs.
• Adopt non-curvilinear Cartesian coordinates; implement FOFC for stability; include IMEX integrators for stiff source terms; enable asynchronous I/O and MPI-IO-based outputs.

Figure 1: (Left) Linear wave convergence rates in 1D for both entropy (solid-line) and sound (dashed-line) waves for different methods. (Right) Density in the Shu-Osher shocktube test using 200 uniform cells and different methods. The solid black line is a reference solution computed using 2000 cell

Experimental results

Research questions

• RQ1 How does AthenaK achieve performance portability across CPUs, GPUs from multiple vendors, and ARM-based systems using Kokkos?
• RQ2 What are the performance and scalability characteristics of AthenaK’s AMR framework and solvers on diverse hardware?
• RQ3 Can AthenaK enable large-scale astrophysical simulations (e.g., AMR hydrodynamics/MHD, GR-radiation transport) at exascale-like performance?
• RQ4 How do design choices (MeshBlockPack, device-resident data, generalized task lists) impact data locality, communication, and overall efficiency?

Key findings

• AthenaK achieves over one billion cell updates per second for 3D hydrodynamics on a single NVIDIA Grace Hopper processor.
• AthenaK attains a typical parallel efficiency of about 80% on 65,536 AMD GPUs on the OLCF Frontier system.
• The framework demonstrates excellent performance and weak scaling across CPUs and GPUs from multiple vendors, enabling exascale-era applications in astrophysical fluid dynamics and numerical relativity.
• AthenaK provides a complete, open-source AMR framework with a broad suite of solvers (Newtonian, SR, GR hydrodynamics/MHD, and GR-radiation transport) and a particle module.
• The design emphasizes hardware portability via Kokkos, device-resident data, and a flexible task-based execution model to optimize performance on heterogeneous architectures.

Figure 2: Slice of temperature through the $z=0$ plane at time $t=10\,\mathrm{kyr}$ in a 3D simulation of a supernova blast wave interacting with a turbulent multiphase medium on a $2048^{3}$ grid. Zoom-ins of selected regions in left panels show nonlinear thin-shell instability (top), turbulent mix