Designing quantum chemistry algorithms with just-in-time compilation

2025/07/13 by Wu, Xiaojie, Sun, Qiming, Wang, Yuanheng
#Computational Physics (physics.comp-ph) #FOS: Mathematics #FOS: Physical sciences #Numerical Analysis (math.NA)

paper · doi:10.48550/arxiv.2507.09772

Abstract

We introduce just-in-time (JIT) compilation to the integral kernels for Gaussian-type orbitals (GTOs) to enhance the efficiency of electron repulsion integral computations. For Coulomb and exchange (JK) matrices, JIT-based algorithms yield a 2x speedup for the small 6-31G* basis set over GPU4PySCF v1.4 on an NVIDIA A100-80G GPU. By incorporating a novel algorithm designed for orbitals with high angular momentum, the efficiency of JK evaluations with the large def2-TZVPP basis set is improved by up to 4x. The core CUDA implementation is compact, comprising only ~1,000 lines of code, including support for single-precision arithmetic. Furthermore, the single-precision implementation achieves a 3x speedup over the previous state-of-the-art.

Citations

• Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework
• Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework
• Software update: The ORCA program system—Version 5.0
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Libcint: An efficient general integral library for Gaussian basis functions

Related

• Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework
• ByteQC: GPU-Accelerated Quantum Chemistry Package for Large-Scale Systems
• Enhancing PySCF-based Quantum Chemistry Simulations with Modern Hardware, Algorithms, and Python Tools
• Computing and Compressing Electron Repulsion Integrals on FPGAs
• Reducing Numerical Precision Requirements in Quantum Chemistry Calculations