[论文解读] 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 Topics被引用 4

一句话总结

本论文介绍了GPU4PySCF，这是首个针对PySCF框架的GPU加速模块，通过Rys积分法对最多包含$g$函数的收缩基集实现两电子排斥积分（ERI）的加速计算。其在单张NVIDIA A100 GPU上实现了相较于多线程CPU Hartree-Fock计算高达100倍的加速，性能与领先的GPU量子化学软件包相当，支持高效的积分直接Fock矩阵构建与核梯度计算。

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.

研究动机与目标

• 在广泛使用的基于Python的PySCF框架中实现量子化学计算的GPU加速。
• 解决Hartree-Fock及相关方法中两电子排斥积分（ERI）计算的性能瓶颈问题。
• 实现一种高效且可扩展的GPU内核，用于基于Rys积分法的四中心ERI计算，支持最多至$g$函数的收缩基集。
• 实现高GPU利用率，并达到与专用GPU量子化学软件包（如GAMESS和QUICK）相当的性能水平。
• 通过社区驱动、开源的实现方式，实现与基于Python的科学计算和机器学习工作流的无缝集成。

提出的方法

• 采用Rys积分法实现GPU加速的四中心ERI计算，通过在三个二维积分上应用高斯积分法实现精确计算。
• 设计专用的CUDA内核用于ERI计算，优化内存效率与高算术强度，支持最多至$g$函数基集。
• 采用矩阵形式的Rys积分法，以利用GPU上高度优化的密集矩阵乘法内核。
• 将GPU ERI内核集成至积分直接Hartree-Fock和核梯度计算工作流中，避免大型矩阵的中间存储。
• 利用CuPy和类似NumPy的GPU数组，加速Python生态系统内的张量收缩与线性代数运算。
• 应用屋顶性能建模以识别并缓解内存瓶颈，特别是针对高角动量积分。

实验结果

研究问题

• RQ1 能否在基于Python的PySCF框架中有效集成GPU加速的两电子排斥积分（ERI）计算，以实现高性能？
• RQ2 GPU4PySCF基于Rys积分法的ERI内核性能与现有CPU及GPU加速量子化学软件包相比如何？
• RQ3 为实现高GPU占用率与高内存效率，针对高角动量ERI（最多至$g$函数）需要哪些算法优化？
• RQ4 在速度提升与资源利用率方面，GPU加速ERI对积分直接Fock矩阵构建与核梯度计算的性能增益有多大？
• RQ5 是否能够在一个不牺牲性能或可用性的前提下，将GPU加速的量子化学框架高效地嵌入Python科学计算生态系统？

主要发现

• GPU4PySCF在单张NVIDIA A100 GPU上相较于PySCF中的多线程CPU Hartree-Fock实现，实现了高达100倍的加速。
• Fock矩阵构建内核（jk_kernel）在低阶积分（$N \leq 3$）下最高可达5 TFLOP/s，但由于内存瓶颈，$N=7$时性能下降至0.8 TFLOP/s。
• 核梯度内核（ejk_grad_kernel）在$N \leq 2$时超过3 TFLOP/s，$N=7$时仍保持0.8 TFLOP/s，由于工作负载分布更优且内存访问更少，其性能比Fock构建内核高出八倍。
• 除$\mathbf{(dp|pp)}$和$\mathbf{(ff|ff)}$等内存受限情况外，两个内核的性能均接近GPU屋顶性能，表明硬件利用率高效。
• 由于中间存储过多（每线程超过234个FP64字），导致高角动量积分（特别是$N>3$）出现寄存器溢出，造成性能下降。
• 梯度计算中采用的积分直接方法减少了全局内存访问流量，提升了缓存利用率，从而实现更好的延迟隐藏与更高的有效吞吐量。