[Paper Review] Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework

Rui Li, Qiming Sun+2 more|arXiv (Cornell University)|Jul 12, 2024

Various Chemistry Research Topics4 citations

TL;DR

This paper introduces GPU4PySCF, the first GPU-accelerated module for PySCF that accelerates two-electron repulsion integral (ERI) evaluation using Rys quadrature for contracted basis sets up to $g$-functions. It achieves up to a 100× speedup over multi-threaded CPU Hartree-Fock and performance comparable to leading GPU quantum chemistry packages on a single NVIDIA A100 GPU, enabling efficient integral-direct Fock builds and nuclear gradient calculations.

ABSTRACT

We introduce the first version of GPU4PySCF, a module that provides GPU acceleration of methods in PySCF. As a core functionality, this provides a GPU implementation of two-electron repulsion integrals (ERIs) for contracted basis sets comprising up to g functions using Rys quadrature. As an illustration of how this can accelerate a quantum chemistry workflow, we describe how to use the ERIs efficiently in the integral-direct Hartree-Fock Fock build and nuclear gradient construction. Benchmark calculations show a significant speedup of two orders of magnitude with respect to the multi-threaded CPU Hartree-Fock code of PySCF, and performance comparable to other GPU-accelerated quantum chemical packages including GAMESS and QUICK on a single NVIDIA A100 GPU.

Motivation & Objective

• To enable GPU acceleration of quantum chemistry calculations within the widely used Python-based PySCF framework.
• To address the performance bottleneck of two-electron repulsion integral (ERI) evaluation in Hartree-Fock and related methods.
• To implement an efficient, scalable GPU kernel for four-center ERIs using Rys quadrature, supporting contracted basis sets up to $g$-functions.
• To achieve high GPU utilization and performance comparable to specialized GPU quantum chemistry packages like GAMESS and QUICK.
• To enable seamless integration with Python-based scientific and machine learning workflows through a community-driven, open-source implementation.

Proposed method

• Implemented GPU-accelerated four-center ERI evaluation using Rys quadrature, which enables exact evaluation via Gaussian quadrature on three 2D integrals.
• Designed custom CUDA kernels for ERI computation, optimized for memory efficiency and high arithmetic intensity, with support for up to $g$-function basis sets.
• Employed a matrix-based formulation of the Rys quadrature method to leverage highly optimized dense matrix multiplication kernels on the GPU.
• Integrated the GPU ERI kernels into integral-direct Hartree-Fock and nuclear gradient workflows, avoiding intermediate storage of large matrices.
• Utilized CuPy and NumPy-like GPU arrays to accelerate tensor contractions and linear algebra operations within the Python ecosystem.
• Applied roofline performance modeling to identify and mitigate memory bottlenecks, particularly for high-angular-momentum integrals.

Experimental results

Research questions

• RQ1 Can GPU acceleration of two-electron repulsion integrals (ERIs) be effectively integrated into the Python-based PySCF framework to achieve high performance?
• RQ2 How does the performance of GPU4PySCF's Rys quadrature-based ERI kernel compare to existing CPU and GPU-accelerated quantum chemistry packages?
• RQ3 What algorithmic optimizations are required to achieve high GPU occupancy and memory efficiency for high-angular-momentum ERIs (up to $g$-functions)?
• RQ4 To what extent can integral-direct Fock build and nuclear gradient calculations benefit from GPU-accelerated ERIs in terms of speedup and resource utilization?
• RQ5 Can a GPU-accelerated quantum chemistry framework be efficiently embedded within the Python scientific computing ecosystem without sacrificing performance or usability?

Key findings

• GPU4PySCF achieves a 100× speedup over the multi-threaded CPU Hartree-Fock implementation in PySCF on a single NVIDIA A100 GPU.
• The Fock build kernel (jk_kernel) achieves up to 5 TFLOP/s for low-order integrals ($N \leq 3$), with performance dropping to 0.8 TFLOP/s for $N=7$ due to memory bottlenecks.
• The nuclear gradient kernel (ejk_grad_kernel) achieves over 3 TFLOP/s for $N \leq 2$, and maintains 0.8 TFLOP/s at $N=7$, outperforming the Fock build kernel by a factor of eight due to better workload distribution and reduced memory access.
• The performance of both kernels remains close to the GPU roofline, indicating efficient hardware utilization, except for memory-bound cases like $\mathbf{(dp|pp)}$ and $\mathbf{(ff|ff)}$.
• Register spilling due to excessive intermediate storage (over 234 FP64 words per thread) causes performance degradation in high-angular-momentum integrals, particularly for $N>3$.
• The integral-direct approach in gradient calculations reduces global memory traffic and improves cache utilization, enabling better latency hiding and higher effective throughput.