PETSc/TAO Developments for GPU-Based Early Exascale Systems

TL;DR

This paper discusses recent developments in the PETSc/TAO library to enhance GPU performance and portability for exascale systems, including new backends, asynchronous features, and communication infrastructure, evaluated on pre- and early exascale systems.

Contribution

The paper introduces new GPU backends, asynchronous support, and communication improvements in PETSc/TAO for exascale computing, addressing performance portability challenges.

Findings

GPU backends improve solver performance on exascale systems.

Asynchronous features enhance GPU kernel concurrency.

Communication infrastructure supports scalable exascale computations.

Abstract

The Portable Extensible Toolkit for Scientific Computation (PETSc) library provides scalable solvers for nonlinear time-dependent differential and algebraic equations and for numerical optimization via the Toolkit for Advanced Optimization (TAO). PETSc is used in dozens of scientific fields and is an important building block for many simulation codes. During the U.S. Department of Energy's Exascale Computing Project, the PETSc team has made substantial efforts to enable efficient utilization of the massive fine-grain parallelism present within exascale compute nodes and to enable performance portability across exascale architectures. We recap some of the challenges that designers of numerical libraries face in such an endeavor, and then discuss the many developments we have made, which include the addition of new GPU backends, features supporting efficient on-device matrix assembly, better support for asynchronicity and GPU kernel concurrency, and new communication infrastructure. We evaluate the performance of these developments on some pre-exascale systems as well the early exascale systems Frontier and Aurora, using compute kernel, communication layer, solver, and mini-application benchmark studies, and then close with a few observations drawn from our experiences on the tension between portable performance and other goals of numerical libraries.