[논문 리뷰] AthenaK: A Performance-Portable Version of the Athena++ AMR Framework

James M. Stone, Patrick D. Mullen외 7명|arXiv (Cornell University)|2024. 09. 24.

Medical Imaging Techniques and Applications인용 수 5

한 줄 요약

AthenaK는 Athena++ 위에 구축된 Kokkos 기반의 성능 포터블 AMR 프레임워크로, 다양한 유체 및 GR 솔버를 갖추어 CPU와 GPU에서 엑사스케일 규모의 계산을 가능하게 한다. 여러 아키텍처에서 강력한 성능과 확장성을 보여준다.

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.

연구 동기 및 목표

• Kokkos로 구현된 폭넓은 하드웨어 호환성(CPUs, GPUs, ARM)을 가진 성능 포터블 AMR 프레임워크인 AthenaK를 소개한다.
• 블록 기반 AMR 구조 내에서 유체 및 GR 솔버의 완전한 모음을 제공한다(유체역학, MHD, GR-방사선 운반 포함).
• 프레임워크 설계 선택(메시블록팩, 디바이스 측 데이터, 태스크 리스트, 경계 처리) 및 오픈 소스 개발 모델을 기술한다.
• 현대 HPC 시스템에서 이식성과 효율성을 검증하며 아키텍처 간 성능 및 확장성을 시연한다.

제안 방법

• Kokkos 프로그래밍 모델에서 Athena++ AMR 프레임워크를 처음부터 재작성하여 아키텍처 간 성능 포터빌리티를 달성한다.
• 디바이스에서 MeshBlock 데이터를 MeshBlockPack으로 구성하여 커널 실행 수를 줄인 병렬 연산을 가능하게 한다.
• 물리 모듈 데이터를 디바이스 거주 Kokkos View로 저장하고, 호스트가 메쉬 구조와 동적 태스크 실행을 태스크 리스트를 통해 제어한다.
• 경계 통신 및 AMR 보간/제한은 병렬 커널 내에서 수행하도록 블록 기반 AMR을 구현한다.
• 뉴트로니안, 특수상대론 SR, 일반상 relatistic GR 유체역학/MHD, GR-방사선 운반 등 다양한 솔버를 제공하고, Lagrangian 추적자 및 전하 입자용 입자 모듈을 포함한다.
• 빌드는 cmake를 사용하고 실행 시 구성 가능한 옵션을 입력 파일로 제공하며, GPU의 레지스터 압력을 줄이기 위한 템플레이팅을 지원한다.
• 비구 곡선 좌표를 채택하고 안정성을 위한 FOFC를 구현하며, 뻣뻣한 소스항에 대해 IMEX 적분기를 포함하고 비동기 IO 및 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 메시블록팩, 디바이스 거주 데이터, 일반화 태스크 리스트와 같은 설계 선택이 데이터 현지성, 통신 및 전반적 효율성에 어떤 영향을 미치는가?

주요 결과

• AthenaK는 단일 NVIDIA Grace Hopper 프로세서에서 3D 유체역학의 10억 개 셀 업데이트를 초과 달성한다.
• AthenaK는 OLCF Frontier 시스템의 65,536개의 AMD GPU에서 약 80%의 일반적인 병렬 효율을 달성한다.
• 프레임워크는 여러 벤더의 CPU 및 GPU에서 우수한 성능과 약한 스케일링을 보여주어 천체물리 유체역학 및 수치 상대론의 엑사스케일 시대 애플리케이션에 기여한다.
• AthenaK는 완전한 오픈 소스 AMR 프레임워크로서 넓은 솔버 세트(N약 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