Quantum Chemistry at Scale: Multi-Node Multi-GPU Two-Electron Integral Code Generation
K. Grace Johnson, Stanford University
As computational power increases, quantum chemistry becomes an increasingly important tool for understanding the behavior of known chemical systems as well as for chemical discovery and design. Increased computational power is often achieved by using increased heterogeneity in computer architectures. This presents new challenges for the quantum chemistry community, as code will need to be adapted to new architectures and scaled to larger systems. For a popular family of quantum chemistry methods—Hartree-Fock (HF) and density functional theory (DFT)—computing two-electron repulsion integrals (ERIs) is often the most computationally intensive step. ERIs scale formally as O(N), where N is the number of atom-centered Gaussian basis functions used to represent the chemical system. The contributions of the ERIs are of Coulomb (J) and exchange (K) type, and require separate algorithms to compute matrix elements efficiently. We previously implemented highly efficient GPU-accelerated J-matrix and K-matrix algorithms in the electronic structure code TeraChem. Although these implementations supported using multiple GPUs on a node, they did not support multiple nodes. This presents a key bottleneck to cutting-edge ab initio simulations of large systems, e.g. excited state dynamics of photoactive proteins. We present our implementation of multi-node multi-GPU J and K-matrix algorithms using the Regent programming language. Regent directly supports distributed computation in a task-based model and can generate code for a variety of architectures, including NVIDIA GPUs. We demonstrate multi-node scaling up to 45 GPUs (3 nodes) and benchmark against hand-coded TeraChem integral code. We also outline our metaprogrammed Regent implementation, which enables flexible code generation for integrals of different angular momenta. In addition, we briefly highlight work on (a) a parallel, scalable exciton model to study the dynamics of biological light-harvesting complexes and (b) quantum circuit simulation using tensor network methods.
Abstract Author(s): K. Grace Johnson, Alex Aiken, Todd J. Martinez