Electron-phonon physics at the exascale: A hybrid MPI-GPU-OpenMP framework for scalable Wannier interpolation

TL;DR

This paper presents a GPU-accelerated implementation of Wannier interpolation for electron-phonon interactions, significantly improving computational efficiency and scalability on exascale supercomputers.

Contribution

A novel hybrid MPI-GPU-OpenMP framework for Wannier interpolation that is portable across multiple GPU architectures and achieves substantial speedups and scalability.

Findings

Up to 29-fold speedup on supercomputers with GPUs.

Seamless portability across NVIDIA, AMD, and Intel GPUs.

Successful large-scale application to topological nanoribbons.

Abstract

We demonstrate a highly efficient GPU implementation of the Wannier interpolation of electron-phonon matrix elements in the EPW code. Building on a systematic analysis of the computational complexity of the algorithm for electron-phonon interpolation, we designed a GPU porting strategy that integrates naturally into the current EPW implementation, and is seamlessly portable to NVIDIA, AMD, and Intel GPUs. We demonstrate this development via extensive benchmarks on conventional semiconductors such as silicon and monolayer MoS$_2$, as well as a large-scale application to topological stanene nanoribbons of width as large as 20nm, which was intractable with previous implementations. Compared to the single MPI parallelization scheme of EPW v5.9, the resulting hybrid MPI-GPU-OpenMP scheme achieves up to 29-fold speedup on leadership-class supercomputers equipped with NVIDIA and Intel accelerators, namely Vista at the Texas Advanced Computing Center, Perlmutter at the National Energy Research Scientific Computing Center, and Aurora at the Argonne Leadership Computing Facility. This framework also achieves nearly ideal scalability up to thousands of GPU nodes on the Aurora supercomputer. With this development, EPW is ready to support electron-phonon physics calculations on exascale platforms.