GPU plugin for PySCF

Installation

[!NOTE] The compiled binary packages support compute capability 6.0 and later (Pascal and later, such as Tesla P100, RTX 10 series and later).

Run nvcc --version in your terminal to check the installed CUDA toolkit version. Then, choose the proper package based on your CUDA toolkit version.

Compilation

One can compile the package with

git clone https://github.com/pyscf/gpu4pyscf.git
cd gpu4pyscf
cmake -S gpu4pyscf/lib -B build/temp.gpu4pyscf
cmake --build build/temp.gpu4pyscf -j 4
CURRENT_PATH=`pwd`
export PYTHONPATH="${PYTHONPATH}:${CURRENT_PATH}"

Then install cutensor for acceleration

pip3 install cutensor-cu11

The package also provides multiple dockerfiles in dockerfiles. One can use them as references to create the compilation envrionment.

Features

• Density fitting scheme and direct SCF scheme;
• SCF, analytical Gradient, and analytical Hessian calculations for Hartree-Fock and DFT;
• LDA, GGA, mGGA, hybrid, and range-separated functionals via libXC;
• Geometry optimization and transition state search via geomeTRIC;
• Dispersion corrections via DFTD3 and DFTD4;
• Nonlocal functional correction (vv10) for SCF and gradient;
• ECP is supported and calculated on CPU;
• PCM models, SMD model, their analytical gradients, and semi-analytical Hessian matrix;
• Unrestricted Hartree-Fock and Unrestricted DFT, gradient, and Hessian;
• MP2/DF-MP2 and CCSD (experimental);
• Polarizability, IR, and NMR shielding

Limitations

• Rys roots up to 9 for density fitting scheme and direct scf scheme;
• Atomic basis up to g orbitals;
• Auxiliary basis up to i orbitals;
• Density fitting scheme up to ~168 atoms with def2-tzvpd basis, bounded by CPU memory;
• Hessian is unavailable for Direct SCF yet;
• meta-GGA without density laplacian;
• Double hybrid functionals are not supported;

Examples

import pyscf
from gpu4pyscf.dft import rks

atom ='''
O       0.0000000000    -0.0000000000     0.1174000000
H      -0.7570000000    -0.0000000000    -0.4696000000
H       0.7570000000     0.0000000000    -0.4696000000
'''

mol = pyscf.M(atom=atom, basis='def2-tzvpp')
mf = rks.RKS(mol, xc='LDA').density_fit()

e_dft = mf.kernel()  # compute total energy
print(f"total energy = {e_dft}")

g = mf.nuc_grad_method()
g_dft = g.kernel()   # compute analytical gradient

h = mf.Hessian()
h_dft = h.kernel()   # compute analytical Hessian

to_gpu is supported since PySCF 2.5.0

import pyscf
from pyscf.dft import rks

atom ='''
O       0.0000000000    -0.0000000000     0.1174000000
H      -0.7570000000    -0.0000000000    -0.4696000000
H       0.7570000000     0.0000000000    -0.4696000000
'''

mol = pyscf.M(atom=atom, basis='def2-tzvpp')
mf = rks.RKS(mol, xc='LDA').density_fit().to_gpu()  # move PySCF object to GPU4PySCF object
e_dft = mf.kernel()  # compute total energy

Find more examples in gpu4pyscf/examples

Benchmarks

Speedup with GPU4PySCF v0.6.0 on A100-80G over Q-Chem 6.1 on 32-cores CPU (Desity fitting, SCF, def2-tzvpp, def2-universal-jkfit, B3LYP, (99,590))

Find more benchmarks in gpu4pyscf/benchmarks