TensorKMC: 新一代神威超级计算机的动力学蒙特卡洛模拟
最近和本文的作者合作进行神威上的相关工作，学习一下老师的上一篇工作。
TensorKMC: Kinetic Monte Carlo Simulation of 50 Trillion Atoms Driven by Deep Learning on a New Generation of Sunway Supercomputer
文章提出了一种三编码算法和空置缓存机制，高效地将第一性原理势神经网络和原子的动力学蒙特卡洛模拟（AKMC）结合。通过实现了快速的特征运算符、NNP的大型算子融合和内存优化，TensorKMC算法在SW26010-pro上模拟了54trillions原子，并最高可以扩展至27456000个核心。
INTRODUCTION
动力学蒙特卡洛（KMC）是一种中尺度的模拟方法，用于研究材料的宏观失效，例如断裂、磨损、腐蚀，其中，原子的动力学蒙特卡洛方法（AKMC）将跳跃事件对应至晶格点状态的改变。AKMC与第一性原理的方法DFT结合，提高了精确度，但受限于较高的计算开销。分子动力学模拟（MD）领域通过结合DFT与神经网络取得了较大进展，但由于较小的时间步，使得MD无法模拟较大的时间尺度，不能替代KMC方法。
OpenKMC是一种AKMC的实现，用于研究锕系合金中空位缺陷的动态演化。文章结合OpenKMC和TensorAlloy，在超级计算机上模拟核反应压力容器上合金的微结构演化。主要的挑战包括：将适用于MD的TensorAlloy，高效地与AKMC结合；内存占用过大；如何针对新神威架构优化。
BACKGROUND
AKMC模型的基本假设是，空位扩散通过一系列的最近邻居跳跃完成。AMKC首先计算所有可能发生的空位跳跃的概率，概率取决于不同状态的能量变化，随后按照概率随机发生空位跳跃，最后计算时间间隔。
并行AKMC框架使用MPI进行数据并行，由于晶体中原子时刻处于晶格点上，这势必面临跳跃冲突的问题，即不同的MPI进程中的原子跳跃到相同的位置。解决方案是通过对每个MPI进程划分子晶格分别计算来避免这个问题。
超级计算机使用了搭载新众核CPU SW26010-pro的新一代神威超级计算机。
AKMC与MD有着许多不同，主要体现在：MD中原子位置不固定，但AKMC中原子始终处于晶格点上；MD中通过原子能量计算受力，AKMC中通过原子能量驱动演化；AKMC若使用传统的数据并行算法会发生跳跃冲突的问题；MD的每个时间步需要计算所有原子，AKMC中只有空位跳跃附近的原子会被影响。而根据MD的原则构建的OpenKMC，遇到了许多挑战：存储了所有原子的性质，导致内存占用过大；对所有原子计算能量，导致冗余的计算；计算所需的内存超过LDM。
INNOVATION AND OPTIMIZATION
Triple-Encoding
每个空位跳跃只由空位以及周围的最近邻居参与，影响范围被限制在截断距离之内。又由于晶体中空位浓度很低，每次计算只需要考虑很少部分的原子，这样通过对每个空位及周边原子进行编码，可以节省内存并便于CPE参与计算。文章提出了Triple-Encoding方法，使用三个张量：相对坐标编码表（CET）、近邻列表编码表（NET）、空位编码表（VET），来描述一个空位系统。
Vacancy Cache Mechanism
TensorKMC中不再将所有原子信息缓存到内存之中，而是只缓存空位系统的编码表，并在每次计算完成后更新。
Memory Usage
TensorKMC通过以算代存，将根据空间坐标访问晶格点性质的过程由2次内存访问减少为1次，同时减少了内存占用。
Fast Feature Operator
TensorKMC使用的特征函数如下，其中p、q为超参数。又由于晶体中距离r是离散分布的，因此可以对函数的指数部分预先建立关于r、p、q的表，便于使用for循环计算，并利用CPE并行加速。
Big-Fusion Operator
卷积神经网络是内存密集的任务，通过改进算法，隐藏内存访问和数据交换时间，将这个过程变为计算密集的。
a. 将1*1卷积计算转化为矩阵乘，ci指输入通道，co指输出通道。
b. 将矩阵乘、偏置和ReLU合并为一个kernel。
c. 将卷积神经网络的所有layers合并为一个kernel，减少内存访问次数。
d. 模型分布式存储在CPEs上，通过RMA共享参数。
e. 使用2个buffer，流水线式计算所有state的能量。
f. CPEs内部的流水线计算。
EVALUATION
这部分就略过了……
参考文献
[1] SHANG H, CHEN X, GAO X, 等. TensorKMC: kinetic Monte Carlo simulation of 50 trillion atoms driven by deep learning on a new generation of Sunway supercomputer[C/OL]//Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. St. Louis Missouri: ACM, 2021: 1-14[2023-02-26]. https://dl.acm.org/doi/10.1145/3458817.3476174.
