MCE SIG发布4款MD模拟流程核心算子,补齐分子动力学模拟主循环
1. 引言
CANN开源社区Material Chemical Engineering SIG(材料化学工程特别兴趣小组,以下简称"MCE SIG")发布4款分子动力学(MD)核心算子——GAFF2力场算子、PME长程静电算子、SHAKE键长约束算子、Velocity Verlet时间积分算子。4款算子分别对应MD主循环中的力场计算、长程静电处理、键长约束与时间积分环节,与此前已发布的LJForceFused融合算子、DPD介观模拟算子一起,可直接服务于以下典型分子动力学模拟场景,帮助开发者快速构建面向化工、材料、能源领域的NPU加速模拟流程:
应用场景 | 核心算子 |
有机分子体系力场计算 | GAFF2 |
带电周期性体系的长程静电 | PME |
含氢键体系的键长约束 | SHAKE |
常规NVE/NPT时间推进 | Velocity Verlet |
2. 背景:MD算子化迁移的主要工程挑战
分子动力学模拟是化工、材料领域常用的计算手段,在催化剂设计、药物分子研究、电池材料、聚合物加工等方向均有应用。将MD关键步骤迁移至昇腾NPU,通常需要解决以下问题:
- 力场项多、链路长:经典力场涉及键、角、二面角、非键等多项势能,Host/Kernel分工与数据传递需逐一实现
- 长程静电处理复杂:周期性边界下库仑相互作用需借助Ewald/PME等算法,实现与调试成本较高
- 约束与积分需协同:SHAKE约束与Velocity Verlet积分往往交替执行,异构环境下需额外编排
- 精度需可验证:科学计算场景通常要求与CPU参考实现对比,确认误差在可接受范围内
针对上述挑战,MCE SIG本次发布的4款算子提供完整NPU解决方案:
- GAFF2力场算子:统一实现键、角、二面角、LJ、库仑等五项势能计算,解决力场项多、链路长的问题
- PME长程静电算子:实现Ewald/PME算法,高效处理周期性边界下的长程静电相互作用
- SHAKE键长约束算子+ Velocity Verlet时间积分算子:协同完成约束与积分,支持大时间步长稳定模拟
4款算子均基于Ascend C开发,提供与CPU参考实现对比的测试用例,确保科学计算精度可验证。
3. 典型应用场景与算子组合
一个典型MD时间步可大致表示为:
Velocity Verlet(时间积分)
→ GAFF2(力场:键/角/二面角/LJ/库仑)
→ PME(长程静电:Ewald实空间 + PME倒易空间)
→ SHAKE(键长约束)
→ 回到 Velocity Verlet
上述流程为算法层面的常见组织方式;实际调用顺序、是否与PME/SHAKE配合,需结合具体模拟体系在Host侧配置。4款算子均可独立编译和测试,便于分模块验证与集成。
3.1 GAFF2力场算子
GAFF2(General AMBER Force Field 2)是面向有机小分子的常用力场。本算子实现了GAFF2的五项势能——键伸缩、键角弯曲、二面角扭转、Lennard-Jones 12-6、库仑静电,力约定与GROMACS一致F=-∇E,并支持PBC最小镜像及1-4对缩放规则。
实现与验证情况:
- Ascend C核函数中,sqrt、sin/cos等函数通过位运算与Newton-Raphson迭代实现,文档记录的单项数学误差约 10-7 量级
- 在Atlas A2/A3系列环境下,64原子体系单步力场计算约200 μs(见算子examples说明,仅为参考值)
- 单元测试中,Bond/LJ/Coulomb分项与CPU参考对比均通过;力误差容差分别为10-4(Bond)、10-3(LJ、Coulomb)
3.2 PME长程静电算子
周期性体系中,库仑相互作用不宜直接截断。本算子实现了Ewald求和及PME(Particle Mesh Ewald)相关组件:
- 实空间:erfc修正的短程Coulomb力
- 倒易空间:B样条电荷分配、3D FFT、倒易空间能量/力计算及插值回原子
PME倒易空间部分的理论复杂度为ON log N量级,适用于需要处理长程静电的带电体系。当前仓库以算法实现与模块划分为主,建议开发者结合本算子测试用例自行评估。
3.3 SHAKE键长约束算子
SHAKE用于在积分后对键长施加约束,常用于含氢键体系,以避免时间步长过小。本算子采用NPU Kernel完成单步修正,Host侧循环调用直至收敛(默认容差 10-4 nm)。
单元测试与算法文档中的对比结果显示:与CPU double参考相比,单键长度偏差约 10-5 nm量级;与Velocity Verlet联跑10步后,最大坐标差约10-5 nm。上述数据来自仓库内测试配置,不代表所有体系与规模下的表现。
3.4 Velocity Verlet时间积分算子
Velocity Verlet是MD中常用的时间积分方法。受Ascend C核函数间数据隔离限制,本算子拆为多个Kernel,由Host编排:
Kernel | 功能 |
vv_integrate | 半步速度更新 + 全步位置更新(含PBC折叠) |
vv_finish | 力更新后的全步速度完成 |
thermo_scale | NPT恒温/恒压相关缩放 |
单元测试中,100步NVE模拟的能量漂移指标要求Δ EE< 0.1%,NPU与CPU参考路径的位置偏差要求10-3 nm。算法文档中更小步数下的坐标差可达10-6 nm量级。测试用例通过即表明在当前配置下精度满足预设容差。
5. 技术实现特点
4款算子均基于Ascend C开发,工程上采用相近的组织方式:
特性 | 说明 |
独立编译 | 各算子目录含CMake配置,可单独构建 |
Host/Kernel分层 | 迭代、多Kernel调度在Host侧,计算密集部分在NPU Kernel |
FP32计算 | 采用单精度浮点,并通过文档/测试给出与CPU参考的误差对比 |
多核并行 | 按原子或约束分配到AI Core(具体策略因算子而异) |
测试与示例 | 提供C++单元测试、examples及算法说明文档 |
目标硬件为Atlas A2/A3系列,开发依赖CANN ≥ 7.0(以各算子README为准)。
4. 四款算子在MD主循环中的位置
各算子在MD主循环中的典型调用顺序如下(具体配合方式需结合模拟体系在Host侧配置):
方向 | 相关算子 |
有机分子体系力场计算 | GAFF2 |
带电周期性体系的长程静电 | PME |
含氢键体系的键长约束 | SHAKE |
常规NVE/NPT时间推进 | Velocity Verlet |
6. 参与与反馈
欢迎化工、材料、能源领域的开发者下载代码、运行测试用例(详情见https://gitcode.com/cann/mat-chem-sim-pred/tree/master/simulation/AI4MD),并通过Issue或社区渠道反馈问题与改进建议。算子仍在持续迭代中,欢迎加入MCE SIG共同完善面向流程工业的开源算子库。
联系公邮:mce-owner@cann.osinfra.cn
