Strong scaling of molecular dynamics simulations with <italic>ab initio</italic>accuracy

基于神经网络的第一性原理精度的分子动力学模拟方法, 目前已成为研究物质相变、材料缺陷等量子物理现象的重要工具. 然而, 尽管该方法在模拟空间尺度、计算精度等方面已取得了显著进展, 但是由于其计算过程需要进行神经网络模型推理, 包含大量复杂算子, 模拟速度仍然受限, 每天仅能实现纳秒级时间尺度的模拟. 一些复杂物理化学现象的研究, 需要微秒甚至毫秒级的模拟. 即使是当前模拟速度最快的工具也需要数周时间的计算, 这严重制约了相关领域的研究进程. 提升模拟工具的强可扩展性能, 加快模拟速度, 减少到解时间, 具有重要的研究意义. 本文对DeePMD-kit这一广泛使用的开源工具展开了一系列细粒度的优化. 我们针对其在强可扩展场景下的计算效率低下的问题, 提出了无框架代码实现、面向瘦高矩阵的GEMM算子优化, 以及大内存下Tabulate算法优化等优化方法. 我们在ARM和X86两个平台上进行了代码实现, 在每个核心只计算一个原子时, 将单个时间步所需的计算时间降低到了最低331 $\upmu$s, 在两个平台上的计算速度分别提升了34和303倍.