AMReX and pyAMReX: Looking Beyond ECP

2024-03-18

The international journal of high performance computing applications

AI Summary

This work addresses key challenges in integrating adaptive mesh refinement (AMR) simulation with data science ecosystems—including AI/ML, in situ analysis, and multiphysics coupling—as well as limitations in GPU zero-copy memory access and inefficient parallel prototyping. To this end, we propose a deep extension of the AMReX framework: (1) a lightweight, high-performance Python binding—pyAMReX—built on pybind11 to uniformly expose AMReX’s C++17/20 core and its CUDA, HIP, MPI, and OpenMP backends; (2) the first implementation of cross-device zero-copy GPU memory sharing, eliminating CPU–GPU data movement overhead; and (3) an extensible AMR algorithm abstraction layer that significantly accelerates heterogeneous parallel prototyping. The framework has been deployed across multiple Exascale Computing Project (ECP) initiatives and extended to cutting-edge simulation domains including plasma physics, combustion, and astrophysics, markedly improving development agility and extreme-scale scalability of AMR applications.

Technology Category

• Machine Learning: Hardware-aware ML
• Multiagent Systems: Agent Communication
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Data management and stream processing for Web, mobile and wireless applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Economics, Online Markets and Human Computation: The sharing economy

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 (ECP), 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.

Problem

• Develops AMReX for block-structured mesh applications with AMR.
• Introduces pyAMReX for Python-based GPU data access and AI/ML integration.
• Enables rapid prototyping and in situ analysis for parallel applications.

Innovation

• AMReX enables adaptive mesh refinement applications.
• pyAMReX bridges AMReX with Python for data science.
• pyAMReX supports zero-copy GPU data for AI/ML.