AMReX and pyAMReX: Looking Beyond ECP

TL;DR

This paper reviews the AMReX framework and its Python binding pyAMReX, highlighting recent developments, new functionalities, and their applications beyond the ECP, including AI/ML integration and in situ analysis.

Contribution

It introduces new features and performance improvements in AMReX and pyAMReX, expanding their capabilities for data science, AI/ML, and non-ECP applications.

Findings

Enhanced GPU data access for AI/ML applications

New functionalities and performance optimizations in AMReX

Broader adoption beyond ECP projects

Abstract

AMReX is a software framework for the development of block-structured mesh applications with adaptive mesh refinement (AMR). AMReX was initially developed and supported by the AMReX Co-Design Center as part of the U.S. DOE Exascale Computing Project, and is continuing to grow post-ECP. In addition to adding new functionality and performance improvements to the core AMReX framework, we have also developed a Python binding, pyAMReX, that provides a bridge between AMReX-based application codes and the data science ecosystem. pyAMReX provides zero-copy application GPU data access for AI/ML, in situ analysis and application coupling, and enables rapid, massively parallel prototyping. In this paper we review the overall functionality of AMReX and pyAMReX, focusing on new developments, new functionality, and optimizations of key operations. We also summarize capabilities of ECP projects that used AMReX and provide an overview of new, non-ECP applications.