Multinode Multi-GPU Two-Electron Integrals: Code Generation Using the Regent Language

The
computation of two-electron repulsion integrals (ERIs) is often
the most expensive step of integral-direct self-consistent field methods.
Formally it scales as O(N4), where N is the number of Gaussian basis functions
used to represent the molecular wave function. In practice, this scaling
can be reduced to O(N2) or less by neglecting small integrals with screening methods. The
contributions of the ERIs to the Fock matrix 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 the use of multiple
GPUs on a node, they did not support the use of 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 multinode multi-GPU J- and K-matrix
algorithms in TeraChem 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 multinode 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.