CPU-GPU optimization and performance of RMG linear-scaling module with optimally localized orbitals

The open-source RMG (Real-Space Multigrid) package solves the Kohn-Sham equations directly on a real-space grid using multigrid acceleration. With a careful data-structure design, RMG is massively parallel on supercomputers with and without GPU accelerators. We recently released a new RMG module that can potentially lead to linear scaling with the system size, i.e., the number of atoms or electrons. In contrast to the main module, in which the wave functions are delocalized and directly represented on real-space grids, the localized-orbitals module expands the wave functions as a linear combination of strictly-localized orbitals that are optimized variationally. For GPU acceleration, we implemented explicit memory management for multiple GPUs and CPU cores per node, which can be easily adapted to either CUDA (Nvidia) or HIP (AMD) programming environment. Timings on the new AMD testbed for the exascale Frontier show very good scalability and GPU speed-up.