Efficient Computing for Reduced Density Matrix on GPUs.

The reduced density matrix (RDM) plays an important role in quantum physics. With the RDM, we can accurately describe the energy of the quantum system. However, solving the RDM is a computationally intensive task where matrix multiplication accounts for more than 95% of the computation time, and exponential complexity, large memory consumption, and tensor sparsity make it challenging to achieve efficient computation. This paper proposes an efficient, scalable and multi-GPU solver for computing the RDM. The complexity of computing RDM is reduced to polynomial by using a tensor network approach in the form of matrix product state (MPS). The sparsity blocks are composed of contiguous memory locations using the list algorithm. The precomputation approach uses a small memory overhead to replace some redundant computations. Furthermore, the sparsity blocks are computed on the GPU using variable-size batch general matrix multiplication (GEMM). We benchmark the performance using a water molecule, and our results show that we achieve a 22.4 × speedup when using four MI50 GPUs compared to a 32-core CPU. Additionally, we demonstrate excellent strong scaling, achieving over 92% parallel efficiency when comparing the performance on 256 nodes versus 16 nodes.