Porting WarpX to GPU-accelerated platforms

Andrew C Myers ,Henri Vincenti ,Rémi Lehe ,D.p Grote ,Kevin Gott ,Jean‐Luc Vay ,Olga V Shapoval ,Maxence Thévenet ,John B Bell ,Luca Fedeli ,Erchan Yang ,Michael E Rowan ,Weiqun Zhang ,Axel Huebl ,Yinjian Zhao ,L D Amorim ,Mark Hogan ,Ann S Almgren ,Lixin Ge ,Cho-Kuen Ng ,Revathi Jambunathan ,N Zaïm ,Edoardo Zoni

doi:10.1016/j.parco.2021.102833

Abstract

WarpX is a general purpose electromagnetic particle-in-cell code that was originally designed to run on many-core CPU architectures. We describe the strategy, based on the AMReX library, followed to allow WarpX to use the GPU-accelerated nodes on OLCF’s Summit supercomputer, a strategy we believe will extend to the upcoming machines Frontier and Aurora. We summarize the challenges encountered, lessons learned, and give current performance results on a series of relevant benchmark problems.

Highlights

WarpX [20] is a fully electromagnetic Particle-in-Cell (PIC) code that is being developed as part of the US Department of Energy’s Exascale Computing Project [8]
WarpX currently supports 2D, 3D and azimuthally decomposed geometries [12] and has the ability to operate in a Lorentzboosted reference frame [18]. It includes support for mesh refinement and dynamic load balancing through the AMReX library [26]
We concentrate on two areas - the scaling of the code on a uniform plasma test case and the performance on a plasma accelerator benchmark problem

Summary

Introduction

The choice was made to port the PIC kernels in PICSAR from Fortran to C++, and to offload kernels using either CUDA, HIP, or DPC++, depending on whether NVIDIA, AMD, Intel hardware is targeted. This removed any need for mixed language programming, which adds substantial complication to the codebase and defeats important compiler optimizations such as inlining. Any implementation differences between the three could be entirely hidden in a performance portability layer; in our case, in the ParallelFor routines in AMReX (see Section 2.5) In this manner, NIVIDA, AMD, and Intel GPUs could all be supported, with little to no change to the. We will present weak and strong scaling results from a uniform plasma problem setup and performance numbers obtained on a plasma acceleration stage benchmark problem

Parallelization

Particle Sorting

Performance Results

Conclusion