GPU acceleration of plane-wave density functional theory calculations in Abinit

TL;DR

This paper presents the GPU porting of the Abinit code for plane-wave DFT calculations, optimizing algorithms for GPU efficiency and comparing performance on multi-GPU architectures.

Contribution

It introduces algorithmic revisions and implementation strategies for GPU acceleration of Abinit, focusing on diagonalization methods and performance optimization.

Findings

GPU porting significantly improves performance over CPU nodes.

Chebyshev polynomial filtering shows better GPU efficiency than LOBPCG.

Performance results demonstrate effective scaling on multi-GPU systems.

Abstract

We report on the GPU porting of the Abinit high-performance simulation code for plane-wave DFT calculations. Large-scale electronic structure calculations require computing the electronic wave function by solving the Kohn-Sham problem discretized over a large number of plane-wave basis functions. Porting such calculations over hundreds of GPU nodes relies not only on extensive usage of vendor libraries from a development perspective, but also on algorithmic revisions of the iterative diagonalization procedure in the resolution of the Kohn-Sham problem to identify GPU-efficient mathematical operations (linear algebra, FFTs) applied to wave functions distributed in memory. The present contribution discusses the Abinit implementation on multi-GPU architectures, providing detailed performance results to compare CPU nodes versus heterogeneous CPU-GPU nodes. Particular attention is given in the comparison of two different diagonalization algorithms, that is Locally Optimal Block Preconditioned Conjugate Gradient and Chebyshev polynomial filtering, in terms of their GPU efficiency.