Accelerating Locality-Driven Integration in Quantum Chemistry with Block-Structured Matrix Multiplication

TL;DR

KerneLDI is a GPU framework that optimizes block-structured matrix multiplication for locality-driven integration in quantum chemistry, significantly accelerating computations while maintaining accuracy.

Contribution

It introduces a co-designed data layout and operators for block-structured matrix multiplication tailored to locality-driven quantum chemistry tasks.

Findings

Up to 10× speedup in exchange-correlation evaluation on GPUs.

Preserves numerical accuracy across diverse molecular systems.

Nearly 6× throughput improvement for ab initio molecular dynamics.

Abstract

Locality-driven integration is a pervasive computational pattern in quantum chemistry, arising whenever spatially localized basis functions interact through numerical quadrature or integral screening. The dominant matrix multiplications in these tasks exhibit dynamic, structured sparsity driven by spatial locality, posing significant challenges for both dense batched kernels and generic sparse formats on GPUs. We present KerneLDI, a GPU-oriented framework that addresses this regime by co-designing data layout, screening logic, and matrix-computation operators to realize block-structured matrix multiplication for locality-driven integration. KerneLDI reorganizes operand matrices into a unified block-filtered representation that retains only spatially relevant blocks, and executes the resulting contractions with customized dense block multipliers that adapt proven dense-matmul optimizations to retained block pairs. We develop and evaluate KerneLDI on exchange--correlation (EXC) integration in Kohn--Sham density functional theory, a representative and computationally critical instance of this pattern. Across diverse molecular systems, KerneLDI preserves numerical accuracy while delivering up to 10$\times$ speedup for EXC evaluation over a dense GPU baseline, scales favorably with increasing system size and multi-GPU parallelism, accelerates end-to-end self-consistent field calculations, and yields nearly 6$\times$ throughput improvement for ab initio molecular dynamics.