Matrix-free algorithms for fast ab initio calculations on distributed CPU architectures using finite-element discretization

TL;DR

This paper introduces matrix-free algorithms for finite-element discretized DFT calculations that significantly improve computational efficiency and scalability on distributed CPU architectures, enabling faster large-scale quantum simulations.

Contribution

The work develops on-the-fly matrix-free algorithms utilizing structured tensor contractions and a unified data layout, achieving substantial speedups over traditional FE-based DFT methods.

Findings

Achieves 1.5-4x speedup over cell-matrix approaches in pseudopotential DFT

Up to 5.8x performance gain in all-electron DFT calculations

Reduces end-to-end time-to-solution in large-scale DFT simulations

Abstract

Finite-element (FE) discretisations have emerged as a powerful real-space alternative to large-scale Kohn-Sham density functional theory (DFT) calculations, offering systematic convergence, excellent parallel scalability, while accommodating generic boundary conditions. However, the dominant computational bottleneck in FE-based DFT arises from the repeated application of the discretised sparse Hamiltonian to large blocks of trial vectors during iterations in an iterative eigensolver. Traditional sparse matrix-vector multiplications and FE cell-matrix approaches encounter memory limitations and high data-movement overheads, particularly at higher polynomial orders, typically used in DFT calculations. To overcome these challenges, this work develops matrix-free algorithms for FE-discretised DFT that substantially accelerate these products by doing on-the-fly operations that utilize structured tensor contractions over 1D basis functions and quadrature data. A unified multilevel batched data layout that handles both real and complex-valued operators is introduced to maximise cache reuse and SIMD utilisation on Frontier (AVX2), Param Pravega (AVX512) and Fugaku (SVE). We also combine terms for optimal cache reuse, even-odd decomposition to reduce FLOP, and mixed-precision intrinsics. Extensive benchmarks show that for large multivector pseudopotential DFT calculations, the matrix-free kernels deliver 1.5-4x speedups over the state-of-the-art cell-matrix approach baselines. For all-electron DFT calculations, the matrix-free operator achieves gains of up to 5.8x due to its efficient implementation and superior arithmetic intensity. When integrated with an error-tolerant Chebyshev-filtered subspace iteration eigensolver, the matrix-free formalism yields substantial reductions in end-to-end time-to-solution using FE meshes that deliver desired accuracies in ground-state properties.