无需矩阵的算法，用于在分布式CPU架构上利用有限元离散化手段快速进行从头算（ab initio）计算

《Computer Physics Communications》：Matrix-free algorithms for fast ab initio calculations on distributed CPU architectures using finite-element discretization

　　摘要
有限元（FE）离散化方法已成为高效进行大规模Kohn–Sham密度泛函理论（DFT）计算的有力实空间替代方案，但其性能往往受到迭代特征值求解器中重复应用离散化哈密顿量到大量多向量限制。传统的全局稀疏矩阵-向量方法和FE单元-矩阵方法在DFT中常用的高阶有限元多项式时，越来越遇到内存限制和高数据移动开销问题。本研究开发了基于有限元的DFT的无矩阵算法，显著加速了在现代分布式CPU架构上的哈密顿量计算。特别是，为FE离散化的Kohn–Sham算子的所有组成部分（包括动能项、局域势能项、GGA梯度项、周期性布洛赫位移项）推导出了无矩阵表达式，并结合了使用一维基函数和积分数据的结构化张量收缩的非局域赝势项。引入了一种统一的多级批量数据布局，以高效处理实数和复数值算子，并最大化AVX2、AVX-512和SVE系统的缓存重用率和SIMD利用率。其他策略，包括奇偶分解和混合精度内联函数，进一步减少了浮点运算和数据移动成本。在Frontier、Param Pravega和Fugaku上的广泛基准测试表明，对于涉及大量多向量的赝势DFT计算，速度提升了高达3.6倍；对于全电子DFT计算，速度提升了高达5.8倍。当与能够容忍矩阵-向量乘法近似的Chebyshev滤波子空间迭代特征值求解器结合使用时，无矩阵方法在保持能量和力的计算精度的同时，实现了显著的性能提升。结果表明，利用FE基的张量积结构的无矩阵算法为在当今CPU超级计算机上进行的小到中等规模从头算模拟提供了一种实用的方法。

引言
基于量子力学理论的从头算计算长期以来一直是材料属性理论研究的核心，为实验和计算研究提供了基本见解。其中，Kohn–Sham密度泛函理论（DFT）[1][2]已成为计算材料科学中最广泛使用和最具影响力的方法之一。在过去几十年中，DFT计算使得能够预测各种材料类别的结构、热学、电子学、光学和力学性质，并且在许多科学高性能计算设施中占据了最大的计算时间。尽管DFT在降低多电子问题的复杂性方面取得了显著成功，但对于涉及数百个以上原子、复杂化学环境或扩展的时空尺度的系统来说，DFT模拟仍然计算密集。因此，仍然需要能够在不牺牲鲁棒性或精度的情况下，使DFT计算在现代并行架构上更快运行的数值算法和实现策略，因为这些发展直接扩展了DFT模拟可处理的系统规模和时间尺度。
在过去几十年中，开发了许多使用DFT进行从头算材料模拟的代码，其中平面波（PW）[3][4][5][6][7]和原子轨道（AO）基组[8][9][10][11][12]被最广泛使用。虽然PW方法对周期性系统具有系统收敛性，但它们仅限于周期性边界条件，并且对于非周期性或半周期性系统需要人工周期性，这会增加计算成本。AO基组虽然对局部化系统高效，但缺乏系统收敛性，并且在不同的化学环境中难以实现精确性。此外，PW和AO方法在现代高性能计算（HPC）架构上都存在可扩展性限制，这些架构需要大规模并行性和分布式内存效率。这些限制促使了实空间离散化技术的发展，如有限差分（FD）[13][14][15][16]、有限元（FE）[17][18][19][20][21][22][23]、小波[25]、psinc基[26]以及其他降阶或自适应基组公式[27][28]，这些技术本质上是灵活且可扩展的。其中，有限元方法是实空间DFT公式领域中相对较新但发展迅速的方法。与其他广泛使用的DFT离散化方案相比，FE基提供了几个显著的优势。特别是，它为收敛性提供了一个系统可改进的基础，同时自然适应包括完全周期性、半周期性和开放边界在内的通用边界条件。FE基函数的局部支持也使该方法特别适合大规模并行计算。FE框架的另一个强大特性是其能够结合自适应空间分辨率，在电子结构快速变化的区域（如原子核心附近或键合区域）局部细化基函数。这种灵活性使得在同一有限元框架内实现全电子和赝势DFT模拟成为可能。最近的研究表明，高阶有限元公式（阶数≥5）在规范保守的赝势DFT计算[22][24]和投影增强波DFT计算[23]中，对于超过5,000个电子的系统，可以显著优于平面波（PW）方法，大幅降低计算成本。开源代码DFT-FE作为2023年ACM Gordon Bell奖的计算引擎[29]，通过其系统可改进的FE离散化和高度可扩展且高效的Kohn–Sham方程求解器体现了这些优势。

Kohn–Sham DFT问题的有限元离散化导致了一个大型稀疏广义特征值问题，其中核心计算挑战是解决这个稀疏特征值问题中最小的Nv个特征对，其中Nv与系统中的电子数量成正比。对于涉及多原子物种的实际问题，Nv通常在10到10,000以上，具体取决于所模拟的材料系统大小。这些特征值问题受益于完全迭代的特征值求解器[30]，这些求解器依赖于迭代正交投影方法。这些方法包括将底层稀疏矩阵正交投影到一个精心构建的、富含所需特征向量的较小子空间（Rayleigh-Ritz步骤），然后对投影矩阵进行子空间对角化，并进行子空间旋转步骤以恢复原始稀疏厄米矩阵的所需正交特征向量估计。常用的DFT特征值求解方法包括Davidson[31][32]、Generalized Davidson[33]、Chebyshev滤波子空间迭代（ChFSI）[34]、LOBPCG[35]、PPCG[36]、Lanczos[37]和块Krylov方法[30]。最近，提出了R-ChFSI[38]，这是一种对矩阵-向量乘法近似具有容忍性的ChFSI变体，用于加速DFT计算中出现的大型稀疏特征值问题。这些迭代方法大多依赖于分块方法，而10,000–15,000电子范围内的关键计算瓶颈在于构建富含所需特征向量的子空间。这需要在分布式环境中重复评估有限元离散化的稀疏矩阵-多向量乘积，高效地完成这一过程对于整体计算性能至关重要。

传统上，有限元离散化的矩阵-向量乘积是使用稀疏矩阵-向量乘法（SpMV）算法执行的。然而，对于高阶FE基函数，SpMV受不规则内存访问和高数据移动成本的影响，先前的研究[39][40]表明，通过在多线程架构上用FE单元级别的密集矩阵-向量乘法替换全局SpMV操作，并随后组装结果元素（单元）级别的贡献，可以实现更好的性能。DFT-FE[22][41]代码在多节点CPU和GPU系统上采用这种策略，证明了在评估涉及数百个向量的FE离散化矩阵-向量乘积时的出色吞吐量，从而有效地解决了基于DFT的材料量子模拟中的大规模非线性特征值问题。

最近，无矩阵有限元算法在计算固体[42]和流体力学[43][44][45]领域成为有前途的替代方案，其中矩阵-向量乘积是在不显式形成或存储元素级别矩阵的情况下即时计算的。通过利用FE基的张量积结构，这些方法将FE离散化算子在试验向量上的作用所涉及的三维积分重新表示为一系列张量收缩，从而显著降低了高阶有限元的算术复杂性，提高了缓存局部性，并消除了显式存储的单元矩阵或全局组装的稀疏矩阵的内存占用。尽管这些方法很有前景，但现有的开源实现尚未针对有效评估大量密集向量上的FE离散化算子的作用进行优化或直接扩展。虽然最近的发展，如使用libCEED库进行多组件向量的无矩阵算子评估[46]和向稀疏多向量扩展的努力[47]在这方面取得了进展，但仍然缺乏一个高效且可扩展的无矩阵FE矩阵-向量乘法算法。直到最近[48]，才提出了用于此目的的批量无矩阵策略，但其演示迄今为止仅限于亥姆霍兹型算子的有限元离散化，而不能直接扩展到包含局部和非局部贡献的实数和复数值FE离散化DFT算子。

受这些挑战的启发，本研究开发了使用六面体元素的FE基离散化的Kohn–Sham DFT问题的无矩阵算法。实现程序旨在现代分布式CPU架构上实现高效的矩阵-多向量乘法，建立在我们最近的发展[48]基础上。我们首先推导出DFT问题的无矩阵方法表达式，其中FE离散化Kohn–Sham算子的局部部分（包括动能项（拉普拉斯算子）、交换-相关项（梯度）、局域势能和周期性布洛赫位移项）完全通过一维FE形状函数及其梯度的结构化张量收缩来表示。通过利用高阶六面体有限元的张量积特性，FE单元级别的局部算子作用可以在飞行中评估，利用积分点消除了构建和存储元素矩阵的需要，从而减少了内存占用并显著提高了数据局部性。我们进一步概述了实数和复数值算子的处理方法，将各种项与相同的积分结合起来，并在每一步张量收缩中重用数据结构。接下来，我们描述了这种方法的数值实现，提出了一种多层次批量布局，以在高效的单一框架中处理实数和复数值的无矩阵作用。我们调整了所提出布局的批量大小，以便张量收缩能够在AVX2、AVX-512和SVE架构上高效映射到SIMD通道，同时利用奇偶分解通过利用FE基函数的对称性来减少浮点运算。这种布局，加上对临时张量的仔细管理，确保了在一系列收缩过程中大部分数据保留在缓存中。我们提出的无矩阵算法通过适当重组循环来重用缓存中的数据，高效地处理试验多向量上的局部和非局部作用，从而减少了数据移动。我们还讨论了在多层次批量布局的背景下，通过非阻塞MPI通信交换处理器边界自由度的算法策略，以及在容忍矩阵-向量乘法近似的特征值求解器中选择性使用统一积分规则和低精度内联函数。最后，我们对所得实现的性能进行了广泛的评估，涵盖了代表性的赝势和全电子材料系统，测试了实数和复数值DFT算子的无矩阵作用。该评估涵盖了DFT算子的局部和非局部组成部分，以及仅局部组成部分的情况。基准测试涵盖了三种代表性的基于CPU的超级计算平台：Frontier（AVX2）、Param Pravega（AVX-512）和Fugaku（SVE）。我们首先评估了使用无矩阵方法与单元-矩阵基线相比，对于不同阶数的FE插值多项式，FE离散化DFT算子在多向量上的作用的速度提升。还提出了一个顶线分析，提供了关于实现的算术强度、持续性能和预期可伸缩性的架构视角。我们进一步讨论了在本工作中考虑的基准系统内，无矩阵方法在Chebyshev滤波子空间迭代特征值求解器中的端到端改进。在所有系统中，无矩阵方法相较于最先进的DFT-FE单元-矩阵方法实现了明显且一致的改进，对于实数和复数赝势算子，加速比为3.6倍；对于全电子算子，加速比为5.8倍，同时保持了总能量和力的化学精度。这些发展共同带来了一个高效且可扩展的无矩阵引擎，用于有限元离散化的Kohn–Sham密度泛函理论（DFT）问题，显著减少了在不同材料系统中求解所需的时间，并提供了一种通用策略，该策略可以扩展到现代超级计算架构上的其他张量积有限元离散化的偏微分方程（PDE）应用。本文的其余部分将逐步开发、分析和验证所提出的Kohn–Sham DFT问题的无矩阵有限元框架。第2节回顾了实空间Kohn–Sham DFT的控制方程，并概述了本工作中采用的FE离散化策略。第3节和第4节介绍了所提出的无矩阵方法，详细说明了基于张量积的局部和非局部算子作用的构建、多层次批量数据布局以及相关的硬件感知优化。第5节提出了一套全面的基准研究，评估了在三种不同CPU架构上的赝势和全电子系统的性能和可扩展性，并进行了屋顶线分析，以解释相对于硬件极限所实现的性能。最后，第6节总结了主要发现，并讨论了将所提出技术扩展到异构系统和GPU加速平台的未来方向。

**Kohn–Sham密度泛函理论的控制方程**
在Kohn–Sham密度泛函理论中，由Na原子核和Ne电子组成的材料系统的基态性质是通过求解非线性特征值问题中的Nv（≥Ne/2）个最小特征对来计算的[1]。为了完整性，我们在这里讨论了通过将相关的静电能重构为适合有限元离散化的局部变分形式而获得的DFT控制方程（见[41]），这种形式可以统一处理各种边界条件。

**无矩阵方法用于DFT问题**
首先，我们将建立必要的符号来描述所提出的Kohn–Sham DFT问题的多层次批量无矩阵算法。本小节中为DFT问题推导出的数学方法遵循Kronbichler和Kormann [44]的精神。在这项工作中，通过将计算域Ω分解为E个不重叠的六面体FE单元Ω(e)来进行有限元离散化，即Ω=?????=1Ω(??)。为了实现高效的分布式并行性，...

**无矩阵算法的数值实现**
我们现在讨论在分布式多节点CPU架构上实现第3节中讨论的多层批量算法策略，以对多向量执行有限元离散化DFT算子的无矩阵操作。关键步骤包括：(i) 提取阶段，使用子域到FE单元级别的映射提取FE单元级别的多向量??(??,??)；(ii) 在无矩阵框架内进行FE单元级别的矩阵-多向量乘法计算，这通过张量收缩来完成。

**性能基准测试**
我们现在检查所提出的无矩阵算法的性能特征，以评估其数值效率和可扩展性。基准研究是在选定的代表性材料系统上进行的，这些系统涵盖了三类不同的有限元离散化DFT算子??=??loc+??nloc：(i) 实数算子，这些算子出现在非周期性系统的赝势DFT模拟中，其中FE算子表示为（??(??)loc=??(??)+??(??)+??(??)）；(ii) 复数算子...

**结论与未来工作**
这项工作提出了一个全面的无矩阵有限元框架，用于在现代分布式CPU架构上加速Kohn–Sham密度泛函理论计算。从头算模拟提出的一个重大计算挑战是在迭代特征值求解器中对数百或数千个试验向量重复应用FE离散化的Kohn–Sham DFT算子。随着材料系统规模的增加和FE多项式的阶数提高，基于全局稀疏性的传统方法...

**在写作过程中声明生成式AI和AI辅助技术**
在准备这项工作的过程中，作者使用了ChatGPT、Gemini和Writefull进行校对。使用这些工具/服务后，作者根据需要审查和编辑了内容，并对出版物的内容承担全部责任。

**未引用的引用**
图1、图10C

**作者贡献声明**
Gourab Panigrahi：撰写——审稿与编辑、撰写——原始草稿、可视化、验证、软件、方法论、调查、形式分析。
Phani Motamarri：撰写——审稿与编辑、监督、资源管理、项目管理、方法论、调查、形式分析、概念化。

**利益冲突声明**
作者声明他们没有已知的利益冲突或个人关系可能影响本文报告的工作。

**作者信息**
Gourab Panigrahi
印度科学研究所计算与数据科学系
CV Raman Road, Bengaluru, 560012, Karnataka, India