[논문 리뷰] Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework

Rui Li, Qiming Sun외 2명|arXiv (Cornell University)|2024. 07. 12.

Various Chemistry Research Topics인용 수 4

한 줄 요약

이 논문은 계약 기저 함수까지 $g$-함수를 지원하는 Rys 적분을 사용하여 이차 전자 반발력 상호작용(ERI) 평가를 GPU로 가속하는 PySCF의 첫 번째 GPU 가속 모듈인 GPU4PySCF를 소개한다. 다중 스레드 CPU 해틀리-폭 계산 대비 최대 100배의 성능 향상을 달성했으며, 단일 NVIDIA A100 GPU에서 최첨단 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.

연구 동기 및 목표

• 광범위하게 사용되는 파이썬 기반의 PySCF 프레임워크 내에서 양자화학 계산의 GPU 가속을 가능하게 하기 위해.
• 해틀리-폭 및 관련 방법에서 두 전자 반발력 상호작용(ERI) 평가의 성능 저하 요인을 해결하기 위해.
• 계약 기저 함수까지 $g$-함수를 지원하는 Rys 적분을 사용하여 네 중심 ERI에 대한 효율적이고 확장 가능한 GPU 커널을 구현하기 위해.
• 고성능 GPU 활용도를 확보하고, GAMESS 및 QUICK와 같은 전용 GPU 양자화학 소프트웨어와 유사한 성능을 달성하기 위해.
• 커뮤니티 주도의 오픈소스 구현을 통해 파이썬 기반 과학 계산 및 머신러닝 워크플로우와의 원활한 통합을 가능하게 하기 위해.

제안 방법

• Rys 적분을 사용한 네 중심 ERI 평가를 GPU로 가속화하여, 3개의 2차원 적분에 대한 가우시안 적분을 통해 정확한 평가를 가능하게 하였다.
• 메모리 효율성과 높은 산술 집약도를 고려해 최적화된 커스텀 CUDA 커널을 설계하여, 최대 $g$-함수 기저 함수까지 지원하였다.
• Rys 적분 방법을 행렬 기반으로 재구성하여 GPU에서 고도로 최적화된 밀집 행렬 곱셈 커널을 활용하였다.
• GPU ERI 커널을 적분-직접 해틀리-폭 및 원자핵 기울기 워크플로우에 통합하여 큰 행렬의 중간 저장을 피하였다.
• CuPy 및 NumPy 유사 GPU 배열을 활용하여 파이썬 생태계 내에서 텐서 결합 및 선형 대수 연산을 가속화하였다.
• 루프라이닝 성능 모델링을 적용하여 메모리 병목 현상을 식별하고, 특히 고각운동량 ERI에 대해 대응하였다.

실험 결과

연구 질문

• RQ1 GPU 가속화된 두 전자 반발력 상호작용(ERI) 평가가 파이썬 기반의 PySCF 프레임워크에 효과적으로 통합되어 고성능을 달성할 수 있는가?
• RQ2 GPU4PySCF의 Rys 적분 기반 ERI 커널 성능이 기존 CPU 및 GPU 가속화된 양자화학 소프트웨어와 비교해 어떻게 되는가?
• RQ3 고각운동량 ERI(최대 $g$-함수까지)에 대해 고성능 GPU 점유율과 메모리 효율성을 확보하기 위해 어떤 알고리즘 최적화가 필요한가?
• RQ4 적분-직접 Fock 행렬 구성 및 원자핵 기울기 계산이 GPU 가속 ERI를 통해 얼마나 빠르게 개선되고, 자원 활용도가 향상되는가?
• RQ5 성능나 비용을 포기하지 않고도 파이썬 과학 계산 생태계 내에 GPU 가속 양자화학 프레임워크를 효율적으로 통합할 수 있는가?

주요 결과

• GPU4PySCF는 단일 NVIDIA A100 GPU에서 다중 스레드 CPU 해틀리-폭 구현 대비 최대 100배의 성능 향상을 달성하였다.
• Fock 행렬 구성 커널(jk_kernel)은 저차수 ERI($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 행렬 구성 커널보다 8배 뛰어난 성능을 보였다. 이는 더 나은 워크로드 분배와 감소된 메모리 접근으로 인한 것이다.
• 모든 커널의 성능는 GPU 루프라이닝에 매우 가까운 수준을 유지하여 하드웨어 활용도가 높은 것으로 나타났다. 다만 $\mathbf{(dp|pp)}$ 및 $\mathbf{(ff|ff)}$와 같은 메모리 병목 케이스를 제외하고는 그러하다.
• 고각운동량 ERI에서 과도한 중간 저장(스레드당 234개 이상의 FP64 단어)으로 인한 레지스터 스플링 현상으로 성능 저하가 발생했으며, 이는 $N>3$에서 특히 두드러졌다.
• 기울기 계산에서 적분-직접 접근은 글로벌 메모리 트래픽을 감소시키고 캐시 활용도를 향상시켜 더 나은 지연 숨기기와 높은 효과적 처리량을 가능하게 하였다.