CANN/mat-chem-sim-pred Lennard-Jones力场融合算子

LJForceFused【免费下载链接】mat-chem-sim-pred面向工业领域聚焦计算仿真、预测两大核心场景构建面向流程工业机理数据双轮驱动的领域计算层推动AI for Science在材料化学领域的深度应用。项目地址: https://gitcode.com/cann/mat-chem-sim-pred作者刘非(Magic_LF)学术指导黄剑兴huangjianxing张玉橙Splendid2025产品支持情况产品是否支持Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×功能说明算子功能实现 Lennard-Jones 力场的融合计算一次完成距离计算、势能计算和力向量计算适用于分子动力学模拟场景。计算公式Lennard-Jones 势能$$ V_{LJ}(r) 4\varepsilon \left[ \left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^{6} \right] $$Lennard-Jones 力$$ F_{LJ}(r) \frac{24\varepsilon}{r} \left[ 2\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^{6} \right] \cdot \frac{\vec{r}}{|r|} $$其中$r$ 为两原子间距离$\varepsilon$ 为势阱深度$\sigma$ 为零势能距离$\vec{r}$ 为距离向量应用场景应用领域典型场景说明催化剂设计分子吸附模拟模拟气体分子在催化剂表面的吸附行为材料研发聚合物性能预测计算高分子链间相互作用力制药工程药物分子对接预测药物与靶点蛋白的结合能石油化工原油组分分析模拟烃类分子间范德华力膜分离气体渗透模拟计算气体分子穿透膜材料的能垒电池材料电解液设计模拟溶剂分子与锂离子的相互作用价值说明为什么需要分子动力学 LJ 专用融合算子1. 内存搬运优化传统分步计算 LJ 力场需要多次内存搬运读取坐标计算距离 → 写回距离矩阵2次搬运读取距离计算势能 → 写回势能2次搬运读取距离计算力分量 → 写回力向量6次搬运xyz三个方向各2次读取力分量做归约 → 写回最终力6次搬运总计16次 HBM 搬运融合算子优化一次读取坐标到片上 UB在 UB 内完成距离、势能、力向量全部计算一次写回最终结果总计2次 HBM 搬运优化点传统方案融合算子HBM 搬运次数16次2次中间结果存储需要不需要内存带宽占用高低Kernel 启动开销多次1次2. 计算优化优化技术说明截断距离优化只计算 r cutoff 的原子对跳过远距离原子减少 O(N²) 计算量牛顿第三定律F_ij -F_ji每对原子只计算一次计算量减半多核并行原子均匀分配到多个 AI Core充分利用 NPU 算力向量化计算利用 Ascend C 向量指令批量处理原子坐标和力向量3. 精度保证特性说明强制 FP32科学计算需要数值稳定性不使用 FP16能量守恒力场计算精度直接影响分子动力学模拟的能量守恒性可复现性相同输入保证相同输出满足科研可复现要求参数说明参数名输入/输出描述数据类型数据格式positions输入原子坐标shape为[N, 3]N为原子数FLOAT32NDepsilon属性势阱深度单位eVFLOAT32-sigma属性零势能距离单位AngstromFLOAT32-cutoff属性截断距离超过此距离的原子对不计算相互作用FLOAT32-forces输出每个原子受到的力shape为[N, 3]FLOAT32NDenergy输出系统总势能FLOAT32-约束说明输入坐标必须为 FLOAT32 类型原子数 N 不超过 65535cutoff 必须大于 0调用说明调用方式样例代码说明aclnn接口test_aclnn_lj_force通过 aclnnLJForce 接口调用 LJForceFused 算子算子特性特性说明融合计算距离 势能 力向量一次完成减少内存访问多核并行原子均匀分配到多个 AI Core 并行计算截断优化只计算 r cutoff 的原子对精度保证强制 FP32 精度满足科学计算需求精度测试与 PyTorch CPU 双精度参考实现对比原子数力向量最大相对误差能量相对误差是否通过200.12%0.08%✓500.15%0.11%✓1000.18%0.14%✓2560.21%0.16%✓精度验证标准相对误差 1%性能数据原子数原子对数PyTorchNPU融合算子加速比642,0160.54 ms0.57 ms0.96x1288,12825.57 ms0.75 ms34.21x25632,640174.96 ms0.85 ms206.23x512130,816183.00 ms1.45 ms126.36x【免费下载链接】mat-chem-sim-pred面向工业领域聚焦计算仿真、预测两大核心场景构建面向流程工业机理数据双轮驱动的领域计算层推动AI for Science在材料化学领域的深度应用。项目地址: https://gitcode.com/cann/mat-chem-sim-pred创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考