Accelerating Locality-Driven Integration in Quantum Chemistry with Block-Structured Matrix Multiplication arXiv:2605.10363v1 [physics.comp-ph] 11 May 2026
Xinran Wei∗†¶ , Yan Pan∗ , Fusong Ju¶∗∥ , Zehao Zhou¶ , Yihong Zhang¶ , Lin Huang§ , Jianwei Zhu¶ , Jia Zhang§ , Huanhuan Xia∗ , Bin Shao∗† , Tao Qin∗ ∗ Zhongguancun Academy, Beijing, China † Zhongguancun Institute of Artificial Intelligence, Beijing, China § IQuest Research, Beijing, China
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 codesigning 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 densematmul 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× 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× throughput improvement for ab initio molecular dynamics. Index Terms—GPU computing, density functional theory, locality-driven integration, structured sparsity, block-structured matrix multiplication.
I. I NTRODUCTION Matrix operations are a foundational abstraction for modern high-performance computing. Dense matrix multiplication and tensor contraction kernels sit at the core of numerical simulation, machine learning, and scientific computing because they expose regular data reuse and map well to accelerator hardware [1]–[3]. Over the past decade, GPU software stacks and hardware-aware libraries have made dense matrix workloads increasingly efficient through hierarchical tiling, sharedmemory staging, register-level accumulation, batching, kernel fusion, and hardware-specialized memory pipelines [4]. In parallel, the HPC community has also made substantial progress on sparse and irregular matrix computation through compressed sparse formats, block-sparse representations, graph ∗ Equal contribution. Correspondence to: [email protected].
and hypergraph reorderings, task-based scheduling, and specialized sparse kernels such as SpMV and SpMM [5]–[9]. These advances have greatly improved the tractability of many sparse linear-algebra workloads, but they do not eliminate the gap between dense-kernel efficiency and irregular scientific matrices. That gap persists because generic sparse methods are typically most effective when sparsity is sufficiently high, the indexing overhead can be amortized, and the access pattern is stable enough to benefit from reusable data layouts [7], [10]. Conversely, forcing irregular data back into dense tiles or batches can recover some locality, but usually at the cost of padding, wasted memory traffic, and strong dependence on preprocessing heuristics [11], [12]. A difficult middle ground remains: matrices that are neither dense nor extremely sparse, whose structure is irregular but not arbitrary, and whose performance depends on preserving locality without paying the full cost of generic sparse indirection. Locality-driven integration is a typical workload that falls into this difficult middle ground in quantum chemistry. Exchange–correlation integration in density functional theory, Coulomb density fitting via the resolution-of-identity (RI) approximation [13], [14], Schwarz-screened exact-exchange construction [15], and related workloads all exhibit the shared pattern of locality-driven integration: spatially localized basis functions or grid partitions interact through numerical quadrature or integral screening. Although these tasks are often introduced as grid-based integration procedures, their dominant cost can be reorganized as matrix multiplications of locality-governed retained contributions, typically involving basis-function evaluation, density-dependent weights, and Fock-like accumulation. The resulting workloads are therefore neither dense nor extremely sparse: contributions are retained dynamically according to locality, geometry, screening thresholds, and grid organization, placing them squarely in the middle ground described above. Current GPU implementations of this class of localitydriven integration workloads often regularize their irregular retained patterns into dense batched GEMM kernels. That regularization improves nominal occupancy, but it also intro-
therefore present KerneLDI, a GPU-oriented framework for block-structured matrix multiplication in locality-driven integration. It reorganizes operand matrices into a unified blockfiltered format and executes the corresponding contractions with customized dense block multipliers. The representation preserves spatial locality while allowing different grid batches to retain different basis blocks without being forced into a common dense shape. The execution model screens block pairs dynamically, improves memory coalescing and arithmetic intensity, and naturally adapts proven dense-matmul optimizations to KerneLDI.
Among typical locality-driven integration tasks, exchange– correlation (EXC) integration in Kohn–Sham density functional theory is both a representative and a scientifically critical case [20]–[22]. It lies on the critical path of self-consistent field iterations, analytic gradients, and ab initio molecular dynamics, and its cost typically scales as O(n3 ) because both the basis size and grid size grow approximately linearly with system size [16], [17]. As a result, EXC directly limits the reachable system sizes and simulated timescales of DFT workloads on modern accelerators [23]–[25]. We therefore take EXC as the workload for developing and evaluating our approach in this paper, although the block-filtered representation and execution model are designed to be applicable to locality-driven integration tasks more broadly.
In summary, the main contributions of this work are as follows:
Taking EXC integration as a representative case, we formulate locality-driven integration as a sequence of dynamically filtered matrix multiplications, making the underlying sparsity explicit and amenable to blockstructured GPU execution. • We introduce a block-filtered transformation that unifies the major matrix operations of locality-driven integration under a common block-structured representation, enabling matrix multiplication over retained blocks with efficient indexing. • We design GPU dense block multipliers that exploit the block-filtered representation for high-locality, highintensity execution on retained block pairs, leveraging Tensor Core acceleration and other proven dense-matmul optimizations on the filtered workload. • We develop a dynamic multi-GPU execution model that decomposes filtered EXC workloads into heterogeneous grid-batch tasks and schedules them on demand to mitigate load imbalance across devices. • We demonstrate on EXC that KerneLDI preserves numerical accuracy while delivering up to 10× 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× throughput improvement for ab initio molecular dynamics. •
Fig. 1: Overview of the KerneLDI workflow. EXC operands generated from basis functions and integration grids are reordered, partitioned, and filtered to produce a compact blockstructured representation. Retained blocks are organized in row- and column-oriented BCS layouts and evaluated in a block-pair fashion. Dynamic scheduling distributes tasks across multiple GPUs for parallel execution. duces padding and wasted memory traffic when heterogeneous basis support is forced into a common batch shape [16]– [19]. Generic fine-grained sparse formats are also a poor fit, because their metadata overhead and irregular access patterns make it difficult to preserve locality and arithmetic intensity on GPUs [5], [7]. Together, these limitations indicate that locality-driven integration requires customized data structures and matrix-computation operators, rather than a fallback to either dense batching or generic sparse execution. This observation points to a co-design problem across data layout, screening logic, and computational operators. Rather than treating sparsity only as a preprocessing step before reverting to dense kernels, the representation and execution model should be designed together to preserve the arithmetic structure of locality-driven integration while matching the memory/computation hierarchy of modern accelerators. We
2
II. W ORKLOAD F ORMULATION AND C HALLENGES
significantly to the numerical quadrature. This locality carries over to the matrix products in Eqs. (4) and (5), where many entries are numerically negligible even though the matrices are formally dense. Importantly, this sparsity is not random. Because retained basis contributions are induced by spatial proximity, significant entries tend to appear in clustered regions associated with nearby atoms, neighboring basis groups, or spatially compact grid partitions. Moreover, the size and location of these retained regions vary across grid batches: different batches may involve substantially different active basis supports depending on molecular geometry, pruning, and screening thresholds. The resulting structure is therefore both clustered and dynamic, rather than uniformly sparse or fixed throughout the calculation. This regime is poorly matched to standard GPU execution strategies. Dense batched implementations regularize heterogeneous retained supports into common matrix shapes, which introduces padding, wasted memory traffic, and unnecessary arithmetic when different grid batches activate different subsets of basis functions. At the other extreme, generic finegrained sparse formats incur substantial metadata overhead and irregular memory access, making it difficult to preserve arithmetic intensity and data locality on modern GPUs. EXC integration therefore falls into an intermediate regime: its dominant matrix products are neither efficiently dense nor well served by generic sparse kernels. These observations motivate the design of KerneLDI. Rather than reverting to dense batching or element-wise sparse execution, KerneLDI preserves locality at block granularity, filters numerically negligible regions before multiplication, and exposes enough regular structure to support efficient GPU execution.
A. EXC as Matrix Products The exchange–correlation (EXC) contribution is a dominant component of Kohn–Sham density functional theory (KSDFT) workloads because it lies on the critical path of selfconsistent field (SCF) iterations, analytic gradients, and ab initio molecular dynamics. In practical Gaussian-basis KSDFT implementations, its dominant cost arises from numerical quadrature over atom-centered integration grids and can be reorganized into matrix products whose structure is determined by basis locality and grid support. Ngrid Let {(ri , wi )}i=1 denote the integration grid points and quadrature weights. The EXC contribution to the Fock matrix is given by Z ∂Exc [ρ(r)] xc χν (r) dr Fµν = χµ (r) ∂ρ(r) 3 R (1) Ngrid X xc ≈b wi vi χµ (ri )χν (ri ), i=1
where χµ (r) and χν (r) are basis functions, vixc denotes the exchange–correlation potential evaluated on grid point ri , and b is a constant prefactor determined by the spin treatment (b = 1 for restricted, b = 2 for unrestricted calculations; see the Appendix). Although Eq. (1) is written as a weighted summation over grid points, its dominant cost can be recast into a small number of matrix operations. We therefore introduce the basis-on-grid matrix Φ ∈ RNbasis ×Ngrid ,
Φµi = χµ (ri ),
(2)
together with the diagonal grid-weight matrix W = diag(wi vixc ) ∈ RNgrid ×Ngrid .
(3)
C. Design Implications
Let D ∈ RNbasis ×Nbasis denote the density matrix. Then one of the dominant intermediates in EXC evaluation can be written as R = DΦ, (4)
The locality-driven EXC workload imposes three requirements on an effective GPU implementation. First, the representation must preserve heterogeneous retained supports without forcing them into a common dense shape. Second, irregularity must be amortized at a coarser granularity than element-wise sparsity so that metadata and indirection do not dominate useful arithmetic. Third, the resulting computation must still expose enough regular structure to map efficiently onto GPU memory hierarchies, matrix units, and multi-GPU execution. KerneLDI is designed around these requirements. It reorganizes EXC operands into fixed-size retained blocks, executes only compatible retained block pairs with dense blocklevel operators, and distributes heterogeneous grid-batch tasks dynamically across devices. The next section details this representation, execution model, and runtime design.
and the EXC contribution itself becomes Fxc = b ΦWΦ⊤ .
(5)
Equations (4) and (5) show that the dominant EXC workload can be expressed as a sequence of matrix products of the generic form C = AB, (6) which we take as the computational abstraction for the remainder of this work. B. Locality-Induced Structure Although the EXC workload can be expressed as matrix products, the resulting operands are not dense in practice. In Gaussian-basis electronic structure calculations, basis functions are spatially localized, so their values decay rapidly away from their associated atomic centers. As a result, for any given grid region, only a subset of basis functions contributes
III. K ERNE LDI A. Overview and Design Goals Given the locality-driven and batch-dependent structure of the EXC matrix products described in Section II, KerneLDI is designed to execute only numerically relevant interactions
3
while avoiding both dense padding and fine-grained sparse indirection. The central idea is to preserve the clustered structure induced by spatial locality, but to do so at block granularity so that the resulting computation remains compatible with highthroughput GPU execution. Figure 1 summarizes the KerneLDI workflow. The framework proceeds in three stages. First, it transforms the original EXC operands into a block-filtered representation through locality-preserving reordering, fixed-size block partitioning, and block-level filtering. This stage converts formally dense but locality-structured matrices into a compact representation that retains only numerically significant blocks. Second, it executes the retained computation using dense block-level multipliers that operate on block pairs rather than on padded dense batches or element-wise sparse entries. Third, it extends the same representation and execution model to multi-GPU platforms by decomposing the filtered EXC workload into heterogeneous grid-batch tasks that are scheduled dynamically across devices. These design choices directly target the requirements identified in Section II. The block-filtered representation preserves locality without enforcing a common dense shape across heterogeneous grid batches; the block-level execution model amortizes sparse irregularity while recovering structured data reuse; and the task-based runtime improves scalability when the retained work varies substantially across batches. The remainder of this section details these three components.
locality can be filtered and exploited at the level of retained matrix blocks. 2) Locality-Preserving Reordering: Block partitioning alone is not sufficient to produce an efficient representation, because the quality of the resulting blocks depends strongly on the ordering of rows and columns before partitioning. If spatially related entries are scattered in the original operand layout, fixed-size blocking will fragment numerically significant regions across many blocks and reduce the effectiveness of subsequent filtering. KerneLDI therefore reorders both grid points and basis functions prior to blockization so that spatially correlated entries become more compact in memory. For grid points, KerneLDI applies a locality-preserving ordering based on atomic partitioning followed by Morton sorting within each atomic region. Grid points are first grouped according to their associated atoms, and the points within each group are then ordered by a Z-order space-filling curve. Given a grid point with discretized coordinates (x, y, z), its Morton code is constructed by bit-interleaving the coordinate components into a one-dimensional key, which is then used for sorting. This procedure maps nearby three-dimensional coordinates to nearby one-dimensional positions while largely preserving spatial proximity, thereby reducing fragmentation when the basis-on-grid matrices are later partitioned into fixedsize blocks. For basis functions, KerneLDI applies a similarity-based ordering derived from the overlap structure of the basis set. The key observation is that block partitioning is performed over matrix indices rather than over physical space. If basis indices with similar spatial support are scattered throughout the operand layout, then numerically significant contributions associated with a given grid region will be split across many blocks, weakening the effect of block filtering. KerneLDI therefore seeks an ordering in which basis functions with similar spatial interaction patterns appear close to one another in memory. To this end, KerneLDI uses the overlap matrix Z Sij = ϕi (r)ϕj (r) dr (8)
B. Block-Filtered Representation 1) Blockization and Notation: KerneLDI operates on fixedsize blocks rather than on individual matrix entries. Let A ∈ RM ×K and B ∈ RK×N denote a pair of operand matrices arising from the EXC formulations in Section II, and let C = AB be the corresponding output. We partition each operand into square blocks of size d × d, which serve as the fundamental unit of storage, filtering, and execution throughout the framework. Under this partitioning, matrix A is divided into Mr = ⌈M/d⌉ block rows and Kc = ⌈K/d⌉ block columns, yielding block entries Apq ∈ Rd×d for 0 ≤ p < Mr and 0 ≤ q < Kc . Similarly, B is divided into Kc block rows and Nc = ⌈N/d⌉ block columns, with block entries Bqr ∈ Rd×d . The output matrix C is therefore partitioned into blocks Cpr ∈ Rd×d , where each output block is formed by accumulating products over compatible inner-dimension block indices: X Cpr = Apq Bqr . (7)
as a practical proxy for spatial and chemical locality. For each basis function i, we define an overlap signature si = (Si1 , Si2 , . . . , SiNbasis ),
(9)
q
and measure the similarity between basis functions i and j using cosine similarity, si · s j . (10) sim(i, j) = ∥si ∥ ∥sj ∥
This blockization serves two purposes. First, it defines a coarser-grained execution unit that is large enough to amortize sparse metadata and indexing overhead. Second, it exposes computation in a form that can be mapped naturally to sharedmemory staging and Tensor-Core-based dense block multiplication on GPUs. KerneLDI therefore does not attempt to preserve fine-grained sparsity explicitly; instead, it reorganizes the EXC workload into a block-structured form in which
Basis functions with similar overlap signatures tend to occupy nearby spatial regions or share similar local chemical environments, and are therefore more likely to be activated together on the same grid batches. KerneLDI places such basis functions close to one another in the reordered layout using similaritybased clustering, so that the active basis support associated with a grid region is concentrated into a smaller number of contiguous index ranges.
4
The purpose of this reordering step is not merely to permute the operands, but to improve block compactness before filtering. After reordering, numerically significant entries are concentrated into fewer and denser blocks, boundary fragmentation is reduced, and the retained block pattern becomes more amenable to efficient block-level execution on GPUs. Implementation details of the Morton ordering, overlap-signature construction, and basis-function clustering are provided in Appendix A-B. 3) Block-Level Filtering: After reordering and block partitioning, KerneLDI applies block-level filtering to discard numerically negligible regions before multiplication. The goal of this step is to reduce both storage and computation by retaining only those blocks whose contributions are potentially significant, while avoiding the metadata and access overhead of fine-grained sparse representations. For each block Apq , KerneLDI computes a block summary based on its maximum absolute entry, apq = ∥Apq ∥max = max |(Apq )uv |, u,v
4) Dual Block-Compressed Layouts: After block-level filtering, KerneLDI stores the retained blocks in blockcompressed layouts tailored to the role of each operand in multiplication. The left operand A is stored in a row-oriented block-compressed format, denoted BCS(R), while the right operand B is stored in a column-oriented format, denoted BCS(C). This dual-layout design is chosen to make the retained block intersections that contribute to each output block directly enumerable during multiplication. For the left operand, BCS(R) groups all retained blocks that share the same block row. Concretely, for each block row p, the layout stores the list of retained column indices q such that (p, q) ∈ SA , together with the corresponding dense block data Apq . Symmetrically, for the right operand, BCS(C) groups all retained blocks that share the same block column. For each block column r, the layout stores the list of retained row indices q such that (q, r) ∈ SB , together with the corresponding dense block data Bqr . In practice, different EXC matrix products may involve non-transposed or transposed operands (e.g., NN or TN forms), and KerneLDI chooses the corresponding block-compressed orientation so that the retained inner-dimension intersections can still be traversed efficiently. This organization is important for efficient multiplication. For a given output block Cpr , KerneLDI needs to identify the compatible inner-dimension block indices q for which both Apq and Bqr are retained. Under the dual block-compressed layouts, this becomes the intersection
(11)
and retains the block only if apq ≥ tb ,
(12)
where tb is the block-retention threshold. The retained block set of A is therefore SA = {(p, q) | ∥Apq ∥max ≥ tb }.
(13)
I(p, r) = { q | (p, q) ∈ SA , (q, r) ∈ SB },
The same procedure is applied to B, yielding SB = {(q, r) | ∥Bqr ∥max ≥ tb }.
and the output block can be written as X Cpr = Apq Bqr .
(14)
This filtering step serves two purposes. First, it removes blocks whose entries are uniformly too small to contribute meaningfully to the subsequent multiplication. Second, it shrinks the candidate search space for block-pair execution, since only products involving retained blocks need to be considered downstream. Because the preceding reordering concentrates significant entries into spatially compact regions, the retained blocks after filtering are typically denser and more clustered than they would be under the original ordering. Importantly, block-level filtering does not yet determine which block pairs are multiplied. It only defines a compact operand representation by removing negligible blocks independently in each matrix. The finer-grained selection of executable block pairs is performed during multiplication, as described in the next subsection. In the present implementation, reordering and block-level filtering are performed once for a given molecular geometry and the resulting block structure is reused across repeated EXC evaluations. Under a fixed geometry, the locality pattern that determines block retention remains unchanged, while only the numerical values associated with the retained blocks are updated. This design amortizes the preprocessing overhead rather than incurring it at every evaluation.
(15)
(16)
q∈I(p,r)
The dual layouts therefore align the operand representation with the block-pair execution pattern, allowing KerneLDI to enumerate only those retained block pairs that can contribute to a given output block. Importantly, this design is not intended merely to compress storage. Its main purpose is to preserve the locality exposed by filtering while supporting efficient traversal of retained block pairs during GPU execution. By storing left operands contiguously by block row and right operands contiguously by block column, KerneLDI reduces indexing overhead, improves memory locality when forming block intersections, and provides the direct input structure for the executable block-pair selection described next. C. Block-Pair Execution Operator 1) Executable Pair Selection: Block-level filtering produces compact operand representations, but it does not by itself determine which retained block pairs should be multiplied. For a given output block Cpr , KerneLDI only needs to consider inner-dimension block indices q that lie in the retained intersection set I(p, r) defined in Eq. (15). Even within this
5
intersection, however, some retained block pairs may still be numerically negligible and need not be executed. KerneLDI therefore applies a second filtering stage at the level of block pairs. For each retained block Apq and Bqr , we reuse the block summaries introduced in Eqs. (11)–(14) and form the pairwise screening criterion apq bqr ≥ td ,
(Appendix A-D) and then assigned to thread blocks in blockpair units, while partial products targeting the same output tile are accumulated before global write-back. This execution strategy is designed to recover the main performance advantages of dense GPU kernels at block granularity. Shared-memory staging reduces repeated global-memory traffic for retained blocks, block-level multiply–accumulate exposes regular data reuse, and warp-level dense execution avoids the control divergence and irregular access patterns typical of fine-grained sparse kernels. At the same time, because only executable retained pairs are processed, the kernel cost scales with the filtered block-pair workload rather than with the full dense product space.
(17)
where apq = ∥Apq ∥max , bqr = ∥Bqr ∥max , and td is the blockpair execution threshold. Only those retained block pairs that satisfy Eq. (17) are dispatched to the multiplication kernel. This two-level filtering serves distinct purposes. The blockretention threshold tb removes negligible blocks independently in each operand and defines the compact block-filtered representation. The block-pair threshold td , in contrast, is applied only after compatible retained intersections have been identified, and further reduces the executed work by screening out block pairs whose joint contribution is too small to justify multiplication. As a result, the set of executed block pairs can be substantially smaller than the Cartesian product of retained blocks in the two operands. For each output block Cpr , the executed computation is therefore X Cpr = Apq Bqr , (18)
In practice, multiple executable pairs may contribute to the same output block Cpr . KerneLDI therefore organizes the kernel around block-pair accumulation: each executed pair produces a partial d × d contribution, and these partial results are accumulated into the corresponding output block before write-back. The precise thread-block mapping and accumulation strategy are implementation dependent, but the governing abstraction remains the same: dense computation is performed only on those block pairs that survive locality-aware filtering. 3) Why This Regime Favors Block-Dense Execution: The execution regime targeted by KerneLDI is not well served by either standard dense batching or generic sparse kernels. Dense batched execution requires heterogeneous retained supports to be regularized into common shapes, which introduces padding and unnecessary arithmetic when different grid batches activate different subsets of basis functions. Generic sparse kernels avoid this dense overhead, but they typically operate at too fine a granularity to make effective use of GPU shared memory, warp-level matrix units, and structured data reuse.
q∈E(p,r)
where E(p, r) = { q ∈ I(p, r) | apq bqr ≥ td }
(19)
denotes the set of executable retained block pairs. This executable-pair formulation provides the immediate interface between the block-filtered representation and the GPU multiplication kernel described next. 2) GPU Kernel Mapping: Once the executable block pairs have been identified, KerneLDI evaluates them using a dense block-level multiplication kernel specialized to the filtered workload. The key design choice is to treat each retained block pair as a unit of dense computation rather than reverting either to padded dense batches or to element-wise sparse execution. This allows the kernel to preserve the arithmetic regularity of dense matrix multiplication while restricting work to numerically relevant regions only. For each executable pair (Apq , Bqr ), the corresponding dense d × d blocks are loaded from the block-compressed layouts into shared memory buffers. GPU threads then perform block-level multiply–accumulate operations on these sharedmemory tiles, with the resulting partial products accumulated into the corresponding output block Cpr . In the current implementation, these block multiplications are mapped to warp-level matrix operations through the WMMA interface, allowing the retained computation to benefit from Tensor Core acceleration while maintaining the block-filtered execution model. For fused evaluations of the electron density and its derivatives, KerneLDI additionally exploits symmetry in the underlying AO-space contraction to reduce redundant block-pair computation and accumulation. Retained pairs are compactly enumerated via parallel prefix-sum compaction
KerneLDI therefore adopts an intermediate strategy: it deliberately gives up fine-grained sparsity in order to recover dense execution regularity at the level of retained blocks. This tradeoff is appropriate because the EXC workload is neither fully dense nor extremely sparse; its significant contributions are clustered by locality and thus can be compacted into dense block tiles after reordering and filtering. By operating on these retained tiles, KerneLDI amortizes sparse indexing overhead over substantially more useful arithmetic than element-wise sparse execution would allow. This block-dense formulation yields three practical advantages. First, it aligns naturally with GPU memory hierarchies, since retained blocks can be staged contiguously and reused within shared memory. Second, it exposes compute tiles that are compatible with high-throughput dense matrix instructions, including Tensor Core operations. Third, it reduces the amount of wasted work relative to dense batching because only retained and executable block pairs are processed. For localitydriven EXC matrix products, this combination makes blockdense execution a better match to the workload structure than either padded dense kernels or generic sparse formats.
6
D. Complexity and Cost Model The cost of KerneLDI is governed by the size of the filtered block workload rather than by the full dense product space. For a given matrix product C = AB, let nA = |SA | and nB = |SB | denote the numbers of retained blocks in the two operands after block-level filtering, and let nP denote the total number of executable block pairs that survive the pairwise screening in Eq. (17). Under this notation, the cost of KerneLDI can be separated into two components: a preprocessing stage and an execution stage. The preprocessing stage includes reordering, block partitioning, block summary construction, filtering, and block-compressed layout generation. The execution stage consists of dense block multiplications over the nP executable block pairs together with the associated block-level accumulation into output tiles. Consequently, the dominant arithmetic cost is governed by the filtered workload size and scales approximately with nP d3 , rather than with the full dense product space alone. This separation explains the performance regime targeted by KerneLDI. On smaller systems, or in cases where locality is not yet strong enough to eliminate a substantial fraction of the block pairs, the one-time preprocessing overhead can account for a noticeable portion of the total runtime. As system size increases, however, the spatial locality of the EXC workload becomes more pronounced, the fraction of retained and executable block pairs decreases relative to the dense block space, and the reduction in multiplication work increasingly outweighs the preprocessing cost. In the present implementation, the block structure is constructed once for a fixed molecular geometry and reused across repeated EXC evaluations, so the preprocessing overhead is amortized rather than paid at every evaluation. KerneLDI therefore exhibits a characteristic crossover behavior: the gains are moderate when the filtered workload is still close to dense, but become substantially larger once locality-induced block sparsity dominates the computation.
Fig. 2: Single-GPU speedup relative to dense GPU execution. Each point denotes one molecule from the dataset [26] and is colored by basis size. The horizontal axis reports the mean absolute error in energies (left) or gradients (right), while the vertical axis reports KerneLDI speedup over EXC-cuBLAS. execution operator described above, producing partial EXC contributions locally. After all tasks have been processed, the partial results are combined through a final reduction step to recover the complete EXC output. This runtime design improves load balance without changing the underlying numerical formulation: communication is confined to task coordination and final reduction, while the bulk of the filtered block computation remains device-local. F. Numerical Considerations and Scope The numerical approximation introduced by KerneLDI is controlled primarily by the block-retention threshold tb and the block-pair execution threshold td , which determine which operand blocks and block products are discarded prior to multiplication. Numerical errors for both energies and gradients are evaluated empirically in Section IV. The current implementation targets local and semilocal EXC functionals, where the dominant computation admits the locality-driven matrix-product form described in Section II; extensions to nonlocal correlation terms remain future work.
E. Multi-GPU Task Decomposition and Dynamic Scheduling IV. R ESULTS
KerneLDI extends naturally to multi-GPU execution by decomposing the filtered EXC workload into grid-batch tasks. Each task corresponds to a filtered grid batch together with its associated retained basis support and block-pair workload, and produces a partial contribution to the EXC matrices or related quantities. Because the retained structure is determined by locality and screening, different grid batches generally contain different numbers of retained blocks and executable block pairs. The resulting task costs are therefore heterogeneous even within the same calculation, making uniform static partitioning ineffective once the workload is distributed across multiple devices. To address this heterogeneity, KerneLDI adopts a pullbased dynamic scheduling strategy. Filtered grid-batch tasks are maintained in a shared task pool, and GPU workers acquire new tasks on demand as soon as they complete their current work. Each GPU evaluates its assigned tasks independently using the same block-filtered representation and block-pair
We evaluate KerneLDI along five axes: correctness and sensitivity of the EXC approximation, scaling performance with increasing system size, end-to-end SCF evaluation, multiGPU scaling, and impact on ab initio molecular dynamics. a) Baselines.: For standalone EXC component benchmarks, the baselines are a CPU dense implementation using OpenBLAS (EXC-OpenBLAS) and a GPU dense implementation using cuBLAS (EXC-cuBLAS), with all remaining electronic-structure components following MADFT [26]. For the end-to-end SCF evaluation (Section IV-C), we additionally compare against GPU4PySCF [18], a widely used GPUaccelerated DFT implementation based on dense batched execution. b) Hardware platform.: Single-GPU experiments use one NVIDIA Tesla V100-PCIE-16GB (5120 CUDA cores, 640 Tensor Cores, 16 GB HBM2, 900 GB/s memory bandwidth) driven by CUDA Toolkit 12.2 and cuBLAS 12.2. CPU
7
To further examine the approximation introduced by pair filtering, we sweep the block-pair execution threshold td on three representative systems while keeping all other settings fixed. Figure 3 reports the resulting mean absolute error (MAE) in EXC energies for taxol, chignolin, and trapcage, complementing the dataset-level accuracy results above by showing how the error varies around the default operating point. Across all three systems, the energy error remains very small in the vicinity of the default threshold td = 10−12 . As td is increased, the screening criterion in Eq. (17) becomes more selective, fewer block pairs are executed, and the resulting approximation error increases. The trend is monotonic for all three systems, and the default setting remains in the low-error regime in each case. Together with the dataset-level accuracy results in Fig. 2, this shows that block-pair filtering remains well controlled in the operating regime used throughout the paper.
Energy MAE (Hartree)
10−7
10−8
taxol chignolin trapcage
default td
10−9
10−10
10−11
10−12 10−10
10−11
10−12
10−13
10−14
Block-pair threshold td
Fig. 3: Sensitivity of EXC energy accuracy to the blockpair threshold td . The vertical dashed line marks the default setting used in the main experiments. Error remains low near the default threshold and increases systematically as td is increased. baseline runs use an AMD EPYC 7452 (32 cores, 2.35 GHz base / 3.35 GHz boost, 128 MB L3 cache) with OpenBLAS 0.3.21. Multi-GPU scaling experiments are conducted on nodes equipped with multiple V100 GPUs connected by NVLink, with inter-node communication over InfiniBand and coordination through OpenMPI. c) Numerical settings.: The block size is d = 32, and both filtering thresholds are set to tb = td = 10−12 in all experiments. EXC numerical integration uses a grid equivalent to the PySCF level-4 setting. The ERI screening tolerance is ϵERI = 10−12 Hartree and the SCF energy convergence threshold is 10−8 Hartree. All experiments use the M06-2X functional [27]. The end-to-end SCF (Section IV-C) and multiGPU scaling (Section IV-D) experiments use the def2-TZVP basis set [28]; all other experiments use the def2-SVP basis set [28]. All reported timing results are averaged over three independent runs.
B. Scaling Performance of KerneLDI To examine how the method behaves as sparsity becomes more pronounced, we evaluate six representative molecules whose atom counts increase progressively: porphy (85 atoms), taxol (110 atoms), chignolin (166 atoms), trapcage (272 atoms), olestra (453 atoms), and crambin (642 atoms). This setup is intended to expose the crossover from regimes where sparse data structures provide limited benefit to regimes where locality-driven sparsity dominates the computation. Figures 4a and 4b reveal a clear crossover behavior. The dense openBLAS and cuBLAS baselines continue to grow much more rapidly with system size because they execute dense matrix operations regardless of how much of the underlying EXC structure is numerically negligible. KerneLDI, by contrast, increasingly benefits from screening and blockstructured execution as the systems become larger and more weakly coupled. The widening gap shows that the method is best matched to the large-system regime where padded dense batching becomes least efficient. For the smaller molecules, the overhead of block construction, filtering, and sparse-aware execution partially offsets the reduction in arithmetic work, so the performance gap remains moderate. Once the workload reaches roughly 200– 300 atoms, corresponding to about 1500–2000 basis functions in this setup, the reduction in screened multiplication work becomes the dominant factor in the overall speedup. Beyond that crossover point, the KerneLDI curves grow much more slowly than the dense baselines for both energy and gradient evaluation. The key point is therefore not simply that KerneLDI is faster, but that its advantage strengthens exactly in the large-system regime where dense batching becomes least well matched to the workload.
A. Correctness, Threshold Sensitivity, and Single-GPU Performance To assess correctness and single-GPU throughput across chemically diverse workloads, we compare KerneLDI against the dense cuBLAS baseline on the 329-molecule dataset proposed in [26]. The dataset contains large molecular systems, including transition-metal complexes, spanning elements from the first through fourth rows of the periodic table; each molecule contains at least 100 atoms and at least four distinct atom types. Figure 2 shows that KerneLDI preserves numerical accuracy while delivering substantial acceleration over dense GPU execution. For both energies and gradients, the mean absolute error remains in the range of roughly 10−11 to 10−9 Hartree, while the speedup increases with system size and reaches up to about one order of magnitude on the largest systems. This trend is consistent with the locality-driven sparsity characteristics described in Section II: as the basis and grid grow, the retained structure becomes more exploitable and the cost of padded dense batching becomes increasingly pronounced.
C. End-to-End SCF Performance Evaluation To determine how the kernel-level acceleration translates into application-level benefit, we next evaluate full selfconsistent field (SCF) calculations. As a practical baseline, we
8
Speedup (relative to 8 GPUs)
8
6
Efficiency
6.33
5.15
4
3.51 3.28 1.85
2
1.00
(a) EXC energy runtime across six molecular systems.
KerneLDI KerneLDI w/o dynamic scheduling Ideal scaling
1.00 1.00 1.00 1.00
0.75
1.79
0.92
0.88
0.90
0.79
0.82 0.64
0.50 0.25 0.00
8
16
32
64
Number of GPUs
Fig. 6: Multi-GPU scaling of KerneLDI on ubiquitin from 8 to 64 GPUs. The top panel reports the speedup normalized to the 8-GPU configuration, and the bottom panel reports the corresponding parallel efficiency. We compare the full KerneLDI implementation against an ablated version without dynamic scheduling. speedups over GPU4PySCF range from 1.52× to 2.70×, demonstrating that the advantages of the block-filtered EXC execution model are preserved at the level of complete SCF calculations rather than being confined to an isolated kernel benchmark. For the individual systems, the speedup is 2.70× for porphy, 2.47× for taxol, 1.52× for chignolin, 1.65× for trapcage, 1.89× for olestra, and 2.57× for crambin. This consistency is important because it shows that the advantage of KerneLDI is retained at the level of complete SCF calculations rather than appearing only in an isolated EXC kernel benchmark. These results show that preserving locality-driven structure in the EXC pathway yields a clear end-to-end performance advantage in complete SCF calculations. Overall, these results indicate that KerneLDI improves the throughput of full DFT calculations in a practically meaningful way. Rather than only accelerating a standalone EXC kernel, the proposed representation and execution model reduces the time-to-solution of complete SCF workflows across a chemically diverse set of systems.
(b) EXC gradient runtime for the same systems.
Fig. 4: Crossover behavior of KerneLDI under increasing system size. Each panel shows calculation time (ms) on the vertical axis for three methods plotted against six molecular systems of growing size on the horizontal axis (porphy (85) through crambin (642), where the trailing number denotes the number of atoms). The annotated × markers indicate the acceleration ratio of KerneLDI relative to EXC-openBLAS at each system size. End-to-end SCF time (s)
104
GPU4PySCF KerneLDI 103
102
Speedup
3.0
2.7×
2.6×
2.5×
2.5
1.9×
2.0
1.7×
1.5×
1.5
D. Multi-GPU Scaling Performance
1.0 porphy (85)
taxol (110)
chignolin (166)
trapcage (272)
olestra (453)
crambin (642)
To evaluate the scalability of KerneLDI on larger accelerator configurations, we perform multi-GPU experiments on ubiquitin using 8, 16, 32, and 64 GPUs. All runs use the same molecular system, basis set, integration settings, and screening thresholds, so that the measured differences reflect only the effect of GPU count and scheduling strategy. We use the 8-GPU configuration as the reference point for normalized speedup and parallel efficiency. Figure 6 reports the resulting scaling behavior. The full KerneLDI implementation exhibits strong scaling across the entire 8–64 GPU range, reaching speedups of 1.85×, 3.51×, and 6.33× at 16, 32, and 64 GPUs, respectively, relative to the 8-GPU baseline. The corresponding parallel efficiencies
System
Fig. 5: End-to-end SCF comparison across six molecular systems. The top panel reports the total SCF wall-clock time for GPU4PySCF and KerneLDI, and the bottom panel reports the corresponding speedup of KerneLDI over GPU4PySCF. compare against GPU4PySCF, a widely used GPU-accelerated implementation based on dense batched execution. All comparisons are performed on the same hardware under identical numerical settings and convergence criteria. Figure 5 shows that KerneLDI consistently improves end-toend SCF performance across all six test systems. The observed
9
are 92.4%, 87.8%, and 79.2%. These results indicate that the filtered EXC workload decomposition of KerneLDI exposes substantial task parallelism and maintains high utilization as the number of GPUs increases. To assess the contribution of the scheduler, we compare against an ablated version of KerneLDI without dynamic scheduling. This variant reaches speedups of 1.80×, 3.28×, and 5.15× at 16, 32, and 64 GPUs, with corresponding parallel efficiencies of 89.7%, 82.1%, and 64.4%. Although the ablated version still scales, its efficiency drops more rapidly as the GPU count increases. The difference is modest at 16 GPUs but becomes increasingly pronounced at 32 and 64 GPUs, where the full KerneLDI implementation maintains noticeably better scaling. This behavior is consistent with the workload characteristics of locality-driven integration. After block filtering, different grid batches retain different numbers of basis blocks and block-pair interactions, so their computational costs are heterogeneous. Without dynamic scheduling, this heterogeneity leads to increasing load imbalance as the workload is distributed across more GPUs. By assigning tasks on demand, KerneLDI reduces this imbalance and preserves higher parallel efficiency at larger scale. Overall, the ubiquitin results show that KerneLDI not only accelerates the EXC workload on a single GPU configuration but also scales efficiently across a substantial multi-GPU range. The persistent gap between the full implementation and the version without dynamic scheduling shows that the scheduling strategy is an important part of making the blockfiltered execution model effective at scale.
Fig. 7: Simulated AIMD trajectory length achievable within a fixed wall-clock budget of 800 hours across three methods. Annotated values indicate the total simulated time achieved by each method within the same computational budget. The molecular snapshots of the AceAla15Lys system alongside each bar illustrate the corresponding structural evolution. allows the simulation to advance to about 35 ps within the identical runtime. The atomic configurations obtained at the end of these runs therefore correspond to different stages of the dynamical evolution. The final structures reached by the three approaches are shown in Fig. 7, demonstrating that the improved EXC performance of KerneLDI enables AIMD simulations to explore a much longer physical timescale within the same computational budget. From a systems perspective, this result is important because it converts kernel-level acceleration into directly usable scientific throughput. KerneLDI increases the amount of first-principles dynamics that can be simulated within a fixed wall-clock budget. V. R ELATED W ORK a) Structured and block-sparse matrix computation on GPUs.: Generic sparse libraries such as cuSPARSE [6] provide element-wise formats (CSR, COO) optimized for highly sparse regimes, while block-sparse representations have been explored in sparse tensor processing [29] and density-matrix methods [30]. Hardware-level structured sparsity [31] doubles throughput but requires a fixed fine-grained pattern. KerneLDI addresses the intermediate regime where block occupancy varies dynamically and the sparsity pattern is determined at runtime by spatial locality. b) GPU-accelerated numerical integration in DFT.: GPU acceleration of EXC integration has been demonstrated via batched dense GEMM in GauXC [17], QUICK [16], and more recent pipelines [18], [19]; Tensor Core acceleration has also been applied in this context [32]–[34]. All of these approaches regularize EXC into uniformly shaped dense workloads. KerneLDI takes a different path by preserving and exploiting the block-level sparsity that batched dense strategies discard. c) Dynamic task scheduling for irregular GPU workloads.: Task-based scheduling and work-stealing strategies have been proposed for irregular GPU workloads in sparse linear algebra [35], molecular dynamics [36], and general GPU task management [37]. KerneLDI adopts a pull-based dynamic scheduling strategy that assigns filtered grid batches to GPU
E. Accelerating Molecular Dynamics with KerneLDI To illustrate the downstream consequence of faster EXC evaluation, we next examine ab initio molecular dynamics for AceAla15Lys, a capped polypeptide comprising 15 alanine residues followed by a C-terminal lysine. This system is a standard folding-oriented test case widely used in vacuum simulations. The initial structure is an extended linear chain. All simulations share the same initial configuration and identical computational settings; the only difference is the EXC evaluation method. Rather than fixing the simulated physical time, we compare different methods under the same wall-clock runtime. Since each MD step requires a full electronic structure evaluation, faster EXC computation enables a larger number of MD steps to be completed within the same runtime. Consequently, the simulated physical time (measured in picoseconds) that can be reached differs across methods. Starting from the same initial configuration and using identical simulation parameters, the CPU baseline EXC-openBLAS is able to propagate the trajectory by only 0.088 ps within the fixed runtime. The GPU baseline EXC-cuBLAS implementation significantly improves the throughput, reaching approximately 6 ps of simulated time under the same conditions. With KerneLDI, the substantially reduced EXC computational cost
10
devices on demand, addressing the load imbalance inherent in locality-driven integration without requiring a priori cost estimates.
[11] R. W. Vuduc and H.-J. Moon, “Fast sparse matrix-vector multiplication by exploiting variable block structure,” in High Performance Computing and Communications: First International Conference, HPCC 2005, Sorrento, Italy, September 21-23, 2005. Proceedings 1. Springer, 2005, pp. 807–816. [Online]. Available: https://doi.org/10.1007/11557654 91 [12] R. Eberhardt and M. Hoemmen, “Optimization of block sparse matrixvector multiplication on shared-memory parallel architectures,” in 2016 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW). IEEE, 2016, pp. 663–672. [Online]. Available: https://doi.org/10.1109/IPDPSW.2016.42 [13] O. Vahtras, J. Almlöf, and M. W. Feyereisen, “Integral approximations for lcao-scf calculations,” Chemical Physics Letters, vol. 213, no. 5–6, pp. 514–518, 1993. [Online]. Available: https://doi.org/10.1016/ 0009-2614(93)89151-7 [14] J. Kussmann, H. Laqua, and C. Ochsenfeld, “Highly efficient resolution-of-identity density functional theory calculations on central and graphics processing units,” Journal of Chemical Theory and Computation, vol. 17, no. 3, pp. 1512–1521, 2021. [Online]. Available: https://doi.org/10.1021/acs.jctc.0c01252 [15] H. Laqua, T. H. Thompson, J. Kussmann, and C. Ochsenfeld, “Highly efficient, linear-scaling seminumerical exact-exchange method for graphic processing units,” Journal of Chemical Theory and Computation, vol. 16, no. 3, pp. 1456–1468, 2020. [Online]. Available: https://doi.org/10.1021/acs.jctc.9b00860 [16] M. Manathunga, Y. Miao, D. Mu, A. W. G”otz, and K. M. Merz, Jr., “Parallel implementation of density functional theory methods in the quantum interaction computational kernel program,” Journal of Chemical Theory and Computation, vol. 16, no. 7, pp. 4315–4326, 2020. [Online]. Available: https://doi.org/10.1021/acs.jctc.0c00290 [17] D. B. Williams-Young, A. Bagusetty, W. A. de Jong, D. Doerfler, H. J. J. van Dam, Á. Vázquez-Mayagoitia, T. L. Windus, and C. Yang, “Achieving performance portability in gaussian basis set density functional theory on accelerator based architectures in nwchemex,” Parallel Computing, vol. 108, p. 102829, 2021. [Online]. Available: https://doi.org/10.1016/j.parco.2021.102829 [18] X. Wu, Q. Sun, Z. Pu, T. Zheng, W. Ma, W. Yan, Y. Xia, Z. Wu, M. Huo, X. Li et al., “Enhancing gpu-acceleration in the python-based simulations of chemistry frameworks,” Wiley Interdisciplinary Reviews: Computational Molecular Science, vol. 15, no. 2, p. e70008, 2025. [Online]. Available: https://doi.org/10.1002/wcms.70008 [19] R. Stocks and G. M. Barca, “Efficient algorithms for gpu accelerated evaluation of the dft exchange-correlation functional,” Journal of Chemical Theory and Computation, vol. 21, no. 20, pp. 10 263–10 280, 2025. [Online]. Available: https://doi.org/10.1021/acs.jctc.5c01229 [20] A. D. Becke, “A multicenter numerical integration scheme for polyatomic molecules,” The Journal of Chemical Physics, vol. 88, no. 4, pp. 2547–2553, 1988. [Online]. Available: https://doi.org/10. 1063/1.454033 [21] O. Treutler and R. Ahlrichs, “Efficient molecular numerical integration schemes,” The Journal of Chemical Physics, vol. 102, no. 1, pp. 346–354, 1995. [Online]. Available: https://doi.org/10.1063/1.469408 [22] P. M. W. Gill, B. G. Johnson, and J. A. Pople, “A standard grid for density functional calculations,” Chemical Physics Letters, vol. 209, no. 5–6, pp. 506–512, 1993. [Online]. Available: https: //doi.org/10.1016/0009-2614(93)80125-9 [23] D. Marx and J. Hutter, Ab Initio Molecular Dynamics: Basic Theory and Advanced Methods. Cambridge University Press, 2009. [24] G. Kresse and J. Furthmüller, “Efficient iterative schemes for ab initio total-energy calculations using a plane-wave basis set,” Physical Review B, vol. 54, no. 16, pp. 11 169–11 186, 1996. [Online]. Available: https://doi.org/10.1103/PhysRevB.54.11169 [25] J. Hutter, M. Iannuzzi, F. Schiffmann, and J. VandeVondele, “Cp2k: atomistic simulations of condensed matter systems,” WIREs Computational Molecular Science, vol. 4, no. 1, pp. 15–25, 2014. [Online]. Available: https://doi.org/10.1002/wcms.1159 [26] F. Ju, X. Wei, L. Huang, A. J. Jenkins, L. Xia, J. Zhang, J. Zhu, H. Yang, B. Shao, P. Dai, D. B. Williams-Young, A. Mayya, Z. Hooshmand, A. Efimovskaya, N. A. Baker, M. Troyer, and H. Liu, “Acceleration without disruption: Dft software as a service,” Journal of Chemical Theory and Computation, 2024. [Online]. Available: https://doi.org/10.1021/acs.jctc.4c00940 [27] Y. Zhao and D. G. Truhlar, “The m06 suite of density functionals for main group thermochemistry, thermochemical kinetics, noncovalent interactions, excited states, and transition elements:
VI. F UTURE W ORK At present, KerneLDI is implemented for local and semilocal exchange–correlation functionals. Extensions to nonlocal correlation terms remain an important direction for future work. More broadly, because the block-filtered representation and dense block multiplier are designed around the general structure of locality-driven matrix multiplications rather than EXC-specific properties, we expect the framework to be applicable to other quantum-chemistry workloads that share this computational pattern, including Coulomb density fitting and Schwarz-screened exact-exchange construction. VII. ACKNOWLEDGMENT This project was supported by the Zhongguancun Academy under the Internal Research Grant No. C20250501. Part of the early-stage work was conducted by X.W., Y.P., and F.J. during internship or affiliation at Microsoft Research. We sincerely thank Derk Pieter Kooi, David Williams-Young and Hongbin Liu for their helpful suggestions during the preparation of the initial version of this manuscripts, and Paola Gori Giorgi, Sebastian Ehlert, Jan Hermann and Zun Wang for their valuable insights into density functional theory and related computational software. R EFERENCES [1] J. Nickolls and W. J. Dally, “GPU computing,” Proceedings of the IEEE, vol. 98, no. 8, pp. 1479–1492, 2010. [2] V. Volkov and J. W. Demmel, “Benchmarking GPUs to tune dense linear algebra,” in SC ’08: Proceedings of the 2008 ACM/IEEE Conference on Supercomputing. IEEE, 2008, pp. 1–11. [Online]. Available: https://doi.org/10.1109/SC.2008.5214359 [3] “cublas: The nvidia cuda basic linear algebra subroutine library,” https: //docs.nvidia.com/cuda/cublas/, 2024, accessed: 2024-11-11. [4] NVIDIA Corporation, CUDA C Programming Guide, 2019, https://docs. nvidia.com/cuda/cuda-c-programming-guide/. [5] N. Bell and M. Garland, “Implementing sparse matrix-vector multiplication on throughput-oriented processors,” in Proceedings of the Conference on High Performance Computing Networking, Storage and Analysis (SC ’09), 2009, pp. 1–11. [Online]. Available: https://doi.org/10.1145/1654059.1654078 [6] M. Naumov, L. Chien, P. Vandermersch, and U. Kapasi, “Cusparse library,” in GPU Technology Conference, vol. 12, 2010. [7] S. Filippone, V. Cardellini, D. Barbieri, and A. Luque, “Sparse matrix-vector multiplication on GPGPUs,” ACM Transactions on Mathematical Software, vol. 43, no. 4, pp. 1–49, 2017. [Online]. Available: https://doi.org/10.1145/3017994 [8] A. Buluç and J. R. Gilbert, “Parallel sparse matrix-matrix multiplication and indexing: Implementation and experiments,” SIAM Journal on Scientific Computing, vol. 34, no. 4, pp. C170–C191, 2012. [Online]. Available: https://doi.org/10.1137/110848244 [9] C. Yang, A. Buluç, and J. D. Owens, “Design principles for sparse matrix multiplication on the GPU,” in Euro-Par 2018: Parallel Processing. Springer, 2018, pp. 672–687. [Online]. Available: https://doi.org/10.1007/978-3-319-96983-1 48 [10] T. A. Davis and Y. Hu, “The University of Florida sparse matrix collection,” ACM Transactions on Mathematical Software, vol. 38, no. 1, pp. 1–25, 2011. [Online]. Available: https://doi.org/10.1145/ 2049662.2049663
11
two new functionals and systematic testing of four m06class functionals and 12 other functionals,” Theoretical chemistry accounts, vol. 120, no. 1, pp. 215–241, 2008. [Online]. Available: https://doi.org/10.1007/s00214-007-0310-x [28] F. Weigend and R. Ahlrichs, “Balanced basis sets of split valence, triple zeta valence and quadruple zeta valence quality for h to rn: Design and assessment of accuracy,” Physical Chemistry Chemical Physics, vol. 7, no. 18, pp. 3297–3305, 2005. [Online]. Available: https://doi.org/10.1039/b508541a [29] E. Qin, G. Jeong, W. Won, S.-C. Kao, H. Kwon, S. Srinivasan, D. Das, G. E. Moon, S. Rajamanickam, and T. Krishna, “Extending sparse tensor accelerators to support multiple compression formats,” in 2021 IEEE International Parallel and Distributed Processing Symposium (IPDPS). IEEE, 2021, pp. 1014–1024. [Online]. Available: https://doi.org/10.1109/IPDPS49936.2021.00110 [30] A. G. Artemov and E. H. Rubensson, “Sparse approximate matrixmatrix multiplication for density matrix purification with error control,” Journal of Computational Physics, vol. 438, p. 110354, 2021. [Online]. Available: https://doi.org/10.1016/j.jcp.2021.110354 [31] J. Choquette, W. Gandhi, O. Giroux, N. Stam, and R. Krashinsky, “Nvidia a100 tensor core gpu: Performance and innovation,” IEEE Micro, vol. 41, no. 2, pp. 29–35, 2021. [Online]. Available: https://doi.org/10.1109/MM.2021.3061394 [32] G. M. J. Barca, C. Bertoni, L. Carrington, D. Datta, N. De Silva, J. E. Deustua, D. G. Fedorov, J. R. Gour, A. O. Gunber, E. Guidez et al., “Recent developments in the general atomic and molecular electronic structure system,” The Journal of Chemical Physics, vol. 152, no. 15, p. 154102, 2020. [Online]. Available: https://doi.org/10.1063/5.0005188 [33] S. Markidis, S. W. Der Chien, E. Laure, I. B. Peng, and J. S. Vetter, “Nvidia tensor core programmability, performance & precision,” in 2018 IEEE international parallel and distributed processing symposium workshops (IPDPSW). IEEE, 2018, pp. 522–531. [Online]. Available: https://doi.org/10.1109/IPDPSW.2018.00091 [34] A. Haidar, S. Tomov, J. Dongarra, and N. J. Higham, “Harnessing GPU tensor cores for fast FP16 arithmetic to speed up mixedprecision iterative refinement solvers,” Proceedings of the International Conference for High Performance Computing, Networking, Storage, and Analysis (SC ’18), pp. 1–11, 2018. [Online]. Available: https://doi.org/10.1109/SC.2018.00050 [35] H. Anzt, T. Cojean, G. Flegar, F. Göbel, T. Grützmacher, P. Nayak, T. Ribizel, Y. M. Tsai, and E. S. Quintana-Ortı́, “Ginkgo: A modern linear operator algebra framework for high performance computing,” ACM Transactions on Mathematical Software, vol. 48, no. 1, pp. 1–33, 2022. [Online]. Available: https://doi.org/10.1145/3480935 [36] A. W. Götz, M. J. Williamson, D. Xu, D. Poole, S. Le Grand, and R. C. Walker, “Routine microsecond molecular dynamics simulations with AMBER on GPUs. 1. generalized born,” Journal of Chemical Theory and Computation, vol. 8, no. 5, pp. 1542–1555, 2012. [Online]. Available: https://doi.org/10.1021/ct200909j [37] M. Steinberger, M. Kenzel, P. Boechat, B. Kerber, M. Dokter, and D. Schmalstieg, “Whippletree: task-based scheduling of dynamic workloads on the GPU,” in ACM Transactions on Graphics (TOG), vol. 33, no. 6, 2014, pp. 1–11. [Online]. Available: https://doi.org/10.1145/2661229.2661250 [38] W. Kohn and L. J. Sham, “Self-consistent equations including exchange and correlation effects,” Physical Review, vol. 140, no. 4A, pp. A1133–A1138, 1965. [Online]. Available: https://doi.org/10.1103/ PhysRev.140.A1133 [39] V. I. Lebedev, “Quadratures on a sphere,” USSR Computational Mathematics and Mathematical Physics, vol. 16, no. 2, pp. 10–24, 1976. [Online]. Available: https://doi.org/10.1016/0041-5553(76)90100-2 [40] I. S. Ufimtsev and T. J. Martinez, “Quantum chemistry on graphical processing units. 2. direct self-consistent-field implementation,” Journal of Chemical Theory and Computation, vol. 5, no. 10, pp. 2619–2628, 2009. [Online]. Available: https://doi.org/10.1021/ct800526s [41] K. Yasuda, “Accelerating density functional calculations with graphics processing unit,” Journal of Chemical Theory and Computation, vol. 4, no. 8, pp. 1230–1236, 2008. [Online]. Available: https: //doi.org/10.1021/ct8001046 [42] J. L. Gálvez Vallejo, G. M. J. Barca, and M. S. Gordon, “High-performance gpu-accelerated evaluation of electron repulsion integrals,” Molecular Physics, 2022. [Online]. Available: https: //doi.org/10.1080/00268976.2022.2112987
[43] M. Zhou and J. Wu, “A gpu implementation of classical density functional theory for rapid prediction of gas adsorption in nanoporous materials,” The Journal of Chemical Physics, vol. 153, no. 7, 2020. [Online]. Available: https://doi.org/10.1063/5.0020797
12
A PPENDIX A S UPPLEMENTARY T ECHNICAL D ETAILS
gible values at grid points far from their center. Grid pruning discards basis-function/grid-point pairs (µ, i) for which |χµ (ri )| < ϵgrid , typically with ϵgrid ≈ 10−10 to 10−14 . This element-wise pruning determines which entries of the basison-grid matrices Φ and Ψ are nonzero. KerneLDI operates at a coarser granularity: after the pruned matrices are partitioned into d × d blocks, the block-level filtering thresholds tb and td further discard blocks and block pairs whose contributions are negligible (see Sections III-B and IV). f) The constant b in Eq. (5).: The factor b in the EXC matrix formula accounts for the spin multiplicity: b = 1 for restricted (closed-shell) calculations and b = 2 for unrestricted (open-shell) calculations in which α and β spin contributions are handled separately.
A. DFT Primer for the HPC Audience This appendix provides a concise introduction to the quantum-chemistry concepts referenced in the main text, aimed at readers familiar with GPU computing but not with electronic-structure theory. a) The Kohn–Sham framework.: Density Functional Theory (DFT) computes the ground-state electronic properties of a molecular or solid-state system by solving for the electron density ρ(r) rather than the full many-body wavefunction. The Kohn–Sham (KS) formulation [38] replaces the intractable many-electron Schrödinger equation with a set of singleparticle equations. The solutions of which, called the KS orbitals ψk (r), reproduce the exact ground-state density. The orbitals are expanded in a finite set of basis functions {χµ }, giving the matrix eigenvalue problem FC = SCϵ stated in Section II. The Fock matrix F depends on the orbitals through the density, so the equation is solved iteratively in a selfconsistent field (SCF) loop until the input and output densities converge. b) Gaussian basis functions.: In the molecular DFT codes, each basis function χµ (r) is a contracted Gaussian: a fixed linear combination of Gaussian primitives g(r) = xl y m z n exp(−α|r−R|2 ) centered on atom R. Two properties are computationally important: (i) products of two Gaussians centered on different atoms are themselves Gaussians, making integral evaluation tractable; and (ii) each Gaussian decays exponentially with distance from its center, so its value on a grid point far from R is negligibly small. Property (ii) is the physical origin of the locality-driven sparsity exploited by KerneLDI. c) The density matrix.: The density matrix D is constructed from the occupied KS orbitals: Dµν = P k∈occ Cµk Cνk , where Cµk is the expansion coefficient of orbital k in basis function µ. It fully P determines the electron density on any grid point via ρ(ri ) = µν Dµν χµ (ri )χν (ri ). d) Numerical integration grids.: The exchange– correlation energy Exc [ρ] is a functional of the density whose analytic integral is generally unknown, so it must be evaluated by numerical quadrature over a three-dimensional grid. Following Becke [20], molecular grids are constructed as a superposition of atom-centered grids, each consisting of a radial component (e.g., Gauss–Chebyshev or Euler– Maclaurin) and an angular component (Lebedev spherical quadrature [39]). A Becke partition function assigns each grid point to an atom so that the molecular integral decomposes into a sum of atomic contributions. The “Level-4 grid” referenced in this paper corresponds to PySCF’s grid convention with up to ∼590 angular points per radial shell per atom, yielding on the order of one to several million grid points for a 100-atom system depending on the pruning scheme. e) Grid pruning.: Because Gaussian basis functions decay rapidly with distance, most basis functions have negli-
B. Reordering Implementation Details This subsection provides the implementation details of the locality-preserving reordering described in Section III-B. 1) Grid-Point Morton Ordering: Grid points are reordered using a Z-order (Morton) space-filling curve to concentrate spatially nearby points into contiguous memory ranges. Each grid point with Cartesian coordinates (x, y, z) is first scaled to an integer lattice by multiplying by a resolution factor (128 in our implementation) and truncating to integer values. A 60-bit Morton code is then constructed by bit-interleaving the three 20-bit integer coordinates: m=
19 X
zk · 23k + yk · 23k+1 + xk · 23k+2 ,
(20)
k=0
where xk , yk , zk denote the k-th bits of the discretized coordinates. Grid points are then sorted by their Morton codes using a stable sort, which preserves the relative order of points sharing the same code. This mapping places threedimensionally adjacent grid points at nearby one-dimensional positions, reducing fragmentation when the basis-on-grid matrices are subsequently partitioned into fixed-size blocks. 2) Overlap-Signature Construction and Basis-Function Clustering: Basis functions are reordered by clustering their overlap signatures, as outlined in Section III-B. The overlap matrix S (Eq. (8)) is already available from the Kohn– Sham setup, so no additional integral evaluation is required. For each basis function i, the overlap signature si = (Si1 , Si2 , . . . , SiNbasis ) (Eq. (9)) is extracted as the i-th row of S. A pairwise cosine-distance matrix is then formed: si · s j , (21) d(i, j) = 1 − ∥si ∥ ∥sj ∥ which equals zero for identical signatures and approaches one for orthogonal ones. Agglomerative (hierarchical) clustering with average linkage is applied to the distance matrix. Starting from Nbasis singleton clusters, the algorithm repeatedly merges the pair with the smallest average inter-cluster distance until a single cluster remains, producing a binary dendrogram. At each merge the two sub-trees are oriented so that their mutually
13
closest leaves are adjacent, yielding a leaf ordering that places basis functions with similar overlap signatures—and hence similar spatial support—at neighbouring index positions. The resulting leaf order is read off as the basis-function permutation π, and the operand matrices are reindexed accordingly. Because the overlap matrix is symmetric and typically has O(Nbasis ) significant entries per row (due to Gaussian 2 decay), the clustering runs in O(Nbasis log Nbasis ) time and is executed once per molecular geometry as part of the preprocessing stage described in Section III-D.
serialize this step; KerneLDI instead uses a parallel prefix sum to perform the compaction entirely within shared memory in O(log n) steps, which is why we detail the procedure here. In the algorithm, the input is a 0/1 array M of length equal to the total number of threads, stored in shared memory, with M[tid ] = 1 if thread tid ’s block pair passed the filter. A parallel exclusive prefix sum over M yields an array P of length thread num + 1 whose i-th entry equals the number of retained pairs in positions 0 through i − 1. For any position where M[i] = 1, the corresponding prefix sum P[i] directly gives that pair’s compact output index, allowing all retained pairs to be written to a dense output array I without gaps or atomic operations.
C. BCS Data Layout Specification The Block Compressed Sparse (BCS) layouts used by KerneLDI are block-granularity analogues of the well-known Compressed Sparse Row (CSR)/ Compressed Sparse Column (CSC) element-wise sparse formats: a) BCS(R) — Block Compressed Sparse Row.: The retained blocks of matrix A are grouped by block-row. The format consists of three arrays: • row_ptr[0 . . . Mr ]: the i-th entry gives the offset into col_idx (and the data array) where block-row i begins. • col_idx[0 . . . RA −1]: for each retained block, its block-column index. 2 • data[0 . . . RA · d −1]: the dense d × d entries of each retained block, stored contiguously in row-major order within each block. Blocks sharing the same block-row are stored adjacent in memory, so iterating over all blocks in a given row requires a single contiguous read of length (row_ptr[i+1] − row_ptr[i]) × d2 elements. b) BCS(C) — Block Compressed Sparse Column.: The retained blocks of matrix B are grouped by block-column using col_ptr, row_idx, and data arrays with mirrored semantics. This ensures that the multiplier can access all blocks contributing to a single output block-column with coalesced reads. c) Comparison with element-wise formats.: In standard CSR, every nonzero carries one column index (typically 4 bytes), yielding metadata overhead ≈ 4 × nnz bytes. In BCS, each retained block carries one column index, so the metadata overhead is ≈ 4 × R bytes for R retained blocks, while the data volume is R × d2 × 4 bytes (single precision). For the typical block size d = 32, the metadata-to-data ratio is 1 : d2 = 1 : 1024, which is negligible compared to elementwise formats.
Algorithm 1 Parallel Addressing 1: Parameters: Number of threads thread num 2: Input: Mask array after filtering M ∈ [0, 1]thread num 3: Output: Total count S, index array I ∈ NS 4: Initialize array P ∈ Nthread num+1 in shared memory 5: Copy M to P in each thread ▷ Thread-level assignment 6: Synchronize threads 7: Update P using prefix sum algorithm ▷ Thread-level
computation 8: Synchronize threads 9: S ← P[thread num] 10: v ← P[thread id + 1] 11: if P[thread id] ̸= v then 12: I[v − 1] ← thread id 13: end if
E. Tensor Core Hardware Mapping KerneLDI invokes Tensor Cores via the WMMA API [33] with FP16 input fragments and FP32 accumulation. Each warp loads a 16 × 16 sub-tile from shared memory into WMMA fragments, executes a hardware matrix multiplyaccumulate, and writes the FP32 result back to a sharedmemory accumulator. For a 32 × 32 block pair, four such tiles are computed and accumulated. Because Tensor Cores deliver ∼125 TFLOPS (FP16) on V100 versus ∼15.7 TFLOPS for FP32 CUDA cores, this mapping yields a significant throughput gain on the dense sub-problems, while the FP32 accumulation preserves the numerical accuracy demonstrated in Section IV. A PPENDIX B E XTENDED R ELATED W ORK
D. Parallel Addressing The parallel addressing scheme described here is invoked in the dense block multiplier (Section III-C) immediately after the block-pair filtering step (Eq. (17)). At that point, each CUDA thread has independently evaluated whether its assigned block pair passes the significance threshold, producing a per-thread Boolean result. The multiplier must then compactly enumerate only the retained pairs so that subsequent shared-memory loading and Tensor Core computation operate on a contiguous, gap-free work list. A naı̈ve serial scan would
This supplement provides a more detailed discussion of related work summarized in the main text. a) Structured and block-sparse matrix computation on GPUs.: Generic sparse libraries such as cuSPARSE [6] provide element-wise formats (CSR, COO) optimized for highly sparse regimes, and autotuning frameworks have further improved SpMV and SpMM throughput [11], [12]. For workloads with coarser structure, block-sparse representations
14
have been explored in accelerator-oriented sparse tensor processing [29] and in quantum-chemistry density-matrix methods [30]. NVIDIA’s Ampere architecture introduced hardwarelevel 2:4 structured sparsity [31], which doubles effective throughput for qualifying weight matrices but requires a fixed fine-grained pattern. These approaches target either very high sparsity or predetermined structures. KerneLDI addresses the intermediate regime where block occupancy varies dynamically across grid batches, and the sparsity pattern is determined at runtime by spatial locality and screening thresholds. b) GPU-accelerated numerical integration in DFT.: GPU acceleration of Kohn–Sham DFT has progressed from early GPU-based electron repulsion integral evaluation [40], [41] to density-fitting approaches for Coulomb and exchange terms [42]. Parallel GPU implementations have also been demonstrated for classical density functional theory [43]. For EXC integration specifically, Williams-Young et al. demonstrated efficient GPU execution within the GauXC framework by grouping grid batches into dense sub-matrices and dispatching them via batched GEMM [17], and Manathunga et al. adopted a similar batched dense strategy in the QUICK package [16]. More recent work by Wu et al. [18] and Stocks et al. [19] has further improved batched-GEMM-based EXC pipelines through refined grid partitioning and memory management. All of these approaches regularize EXC into uniformly shaped dense workloads, discarding the block-level sparsity that KerneLDI preserves. c) Tensor Core utilization in scientific computing.: NVIDIA Tensor Cores, originally designed for deep-learning workloads [33], have been adopted in scientific computing for mixed-precision dense linear algebra, iterative solvers, and FFT-based computations [34]. Barca et al. [32] demonstrated the use of Tensor Cores for accelerating integral evaluation in quantum chemistry. KerneLDI leverages Tensor Cores within its dense block multiplier to accelerate retained block-pair products, adapting hardware-accelerated warp-level matrix multiply-accumulate to the block-filtered execution model. d) Dynamic task scheduling for irregular GPU workloads.: Load imbalance arising from irregular or datadependent work distributions is a well-known challenge in GPU computing. Task-based runtime systems and workstealing schedulers have been proposed for GPU workload management [37], molecular dynamics force decomposition [36], and sparse linear algebra [35]. Within DFT, grid-based numerical integration is inherently heterogeneous because different spatial regions retain different numbers of significant basis functions. KerneLDI addresses this heterogeneity through a pull-based dynamic scheduling strategy that assigns filtered grid batches to GPU devices on demand, ensuring balanced utilization across multi-GPU platforms without requiring a priori cost estimates.
15