[論文レビュー] 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 は Athena++ に基づく新しい Kokkos ベースの、CPU と GPU でのエクサスケール計算を可能にする多様な流体・GR解法を備えた性能ポータブル AMR フレームワークを提供します。さまざまなアーキテクチャで高い性能とスケーラビリティを実証します。

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 放射輸送を含む）の完全なスイートを提供。
• MeshBlock、デバイス側データ、タスクリスト、境界処理といった設計選択とオープンソース開発モデルを記述。
• アーキテクチャを跨ぐ性能とスケーリングを実証し、現代の HPC システム上でのポータビリティと効率性を検証。

提案手法

• Athena++ の AMR フレームワークをゼロから Kokkos プログラミングモデルで書き直し、アーキテクチャ横断の性能ポータビリティを達成。
• MeshBlock データをデバイス上で MeshBlockPack として整理し、カーネル起動の削減と並列操作を可能に。
• 物理モジュールデータをデバイスに居住する Kokkos Views として格納し、ホストがメッシュ構造とタスクリストによる動的タスク実行を制御。
• 境界通信を伴うブロックベースの AMR を実装し、AMR の延長/制限を並列カーネル内で実行。
• Newtonian、特殊相対論的（SR）、一般相対論的（GR）流体力学/ MHD、GR 放射輸送を含む広範な解法を提供。Lagrangian トレーサーと荷電粒子用の粒子モジュールを含む。
• cmake をビルドに使用；入力ファイルによるランタイム設定可能なオプション；GPU のレジスタ圧を減らすテンプレート化をサポート。
• 非曲率 Cartesian 座標を採用；安定性のための 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 はエクサスケールに近い性能で大規模な天体物理学的シミュレーション（例：AMR 流体力学/ MHD、GR 放射輸送）を可能にできるか。
• RQ4 設計上の選択（MeshBlockPack、デバイス居住データ、汎用タスクリスト）はデータ局所性、通信、全体的な効率性にどのように影響するか。

主な発見

• AthenaK は NVIDIA Grace Hopper プロセッサ 1 台で 3D 流体力学のセル更新を10億を超える速度で達成。
• AthenaK は Frontier システムの 65,536 台の AMD GPU 上で典型的な並列効率が約 80%。
• 本フレームワークは複数のベンダーの CPU/ GPU で優れた性能とウェイクスケールを示し、天体物理流体力学と数値相対論のエクサスケール時代のアプリケーションを可能にします。
• AthenaK は完全なオープンソースの AMR フレームワークを提供し、広範な解法（Newtonian、SR、GR 流体力学、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