[论文解读] 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 Applications被引用 5

一句话总结

AthenaK 提供了一个基于 Kokkos 的、面向性能可移植的 AMR 框架，建立在 Athena++ 之上，具备多样的流体和 GR 求解器，能够在 CPU 和 GPU 上执行规模达到 exascale 的计算。它在多架构上展现出强健的性能和可扩展性。

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.

研究动机与目标

• 介绍 AthenaK，一个使用 Kokkos 实现、面向广泛硬件兼容性的性能可移植 AMR 框架（CPU、GPU、ARM）。
• 在基于块的 AMR 结构中，提供完整的流体和 GR 求解器套件，包括流体力学、MHD 以及 GR-辐射输运。
• 描述框架设计选择（MeshBlockPack、设备端数据、任务列表、边界处理）以及开源开发模型。
• 展示跨架构的性能和可扩展性，验证在现代 HPC 系统上的可移植性与效率。

提出的方法

• 从头用 Kokkos 编程模型改写 Athena++ 的 AMR 框架，以实现跨架构的性能可移植性。
• 将 MeshBlock 数据在设备上整理为 MeshBlockPack，以实现并行操作并减少内核启动次数。
• 将物理模块数据存储为设备驻留的 Kokkos Views，由主机控制网格结构并通过任务列表进行动态任务执行。
• 实现基于块的 AMR，边界通信，以及 AMR-延拓/限制 在并行内核中完成。
• 提供广泛的求解器：Newtonian、特殊相对论（SR）、广义相对论（GR）流体力学/磁流体力学，以及 GR 辐射输运；包含用于拉格朗日示踪和带电粒子的粒子模块。
• 使用 cmake 构建；通过输入文件进行运行时可配置选项；支持模板化以降低 GPU 的寄存器压力。
• 采用非曲线坐标的笛卡尔坐标系；实现 FOFC 以提高稳定性；包含用于刚性源项的 IMEX 积分器；支持异步 I/O 与基于 MPI-IO 的输出。

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

实验结果

研究问题

• RQ1 AthenaK 如何利用 Kokkos 在 CPU、来自不同供应商的 GPU，以及 ARM 基于系统上实现性能可移植性？
• RQ2 在不同硬件上，AthenaK 的 AMR 框架和求解器的性能与可扩展性特性是什么？
• RQ3 AthenaK 是否能够在接近 exascale 的性能下实现大规模天体物理模拟（如 AMR 流体力学/MHD、GR 辐射输运）？
• RQ4 设计选择（MeshBlockPack、设备驻留数据、广义任务列表）如何影响数据局部性、通信以及整体效率？

主要发现

• AthenaK 在单个 NVIDIA Grace Hopper 处理器上实现了 3D 流体力学超过十亿单元每秒的更新。
• AthenaK 在 OLCF Frontier 系统上对 65,536 张 AMD GPU 常规并行效率约为 80%。
• 该框架在来自多家厂商的 CPU 和 GPU 上表现出色的性能和弱尺度性，推动天体物理流体动力学与数值相对论等领域在 exascale 时代的应用。
• AthenaK 提供一个完整的、开源的 AMR 框架，配有广泛的求解器（Newtonian、SR、GR 流体力学/MHD，以及 GR 辐射输运）和粒子模块。
• 该设计通过 Kokkos、设备驻留数据以及灵活的基于任务的执行模型，强调硬件可移植性，以在异构架构上优化性能。

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