已关闭
[Requirement|需求建议]: 增加 ProdForceSeA 支持 Ascend950 实现 SIMT #2451
ideal_创建于 7月31日关闭于 8月3日
7月31日 添加了label:requirement
陈思
7月31日 评论:
7月31日 评论:
/assign @wang-shilong32


7月31日 将 wang-shilong32 设为负责人
8月3日 关闭了 issue
8月3日 添加了label:resolved
25 天前 issue状态由 进行中 改变为 已完成
Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
一、背景信息 (必填)
新增了 ProdForceSeA 算子,面向 Ascend950 平台实现 SIMT(Grid-Stride)模式下的 DeepMD-kit DeepPot-SE Angular 原子力计算。算子接收 net_deriv、in_deriv、nlist、natoms 四个输入,输出形状为 [nframes, nall, 3] 的 atom_force,整体覆盖算子定义、Shape/类型推导、Tiling、SIMT Kernel 及配套测试与 Golden 等完整算子开发流程。
二、价值/作用 (必填)
增加 ProdForceSeA 支持 Ascend950 实现 SIMT
三、设计方案 (必填)
算子定义与注册: 在 prod_force_se_a_def.cpp 和 prod_force_se_a_proto.h 中定义 ProdForceSeA 算子,声明四个输入、一个输出和 n_a_sel/n_r_sel 两个必选属性,通过 OpAICoreConfig 在 Ascend950 上启用动态形状和动态编译。
Shape/类型推导: 在 prod_force_se_a_infershape.cpp 中实现基于 natoms 值依赖的 Shape 推导(输出 [nframes, nall, 3]),并进行帧数一致性校验;在 prod_force_se_a_graph_infer.cpp 中实现输出 dtype 跟随 net_deriv 的类型推导。
Tiling 策略: 在 prod_force_se_a_tiling.cpp 中实现完整的 Tiling 函数,包含输入维度、dtype、Shape 一致性和属性语义的多层校验,基于工作负载动态计算所需核数(ComputeNeedCoreNum)并设置 BlockDim 与 UB 内存分配。
SIMT Kernel 实现: 在 prod_force_se_a_simt.h 中实现 OpProdForceSeASimt 模板函数,采用 Grid-Stride 分帧、帧内串行完成清零-中心力(减)-邻居力(加)的计算,使用 volatile 局部变量阻止 FMA 融合以保证与 numpy Golden 的两步舍入一致,并通过模板化 uint32_t/uint64_t 索引路径适配不同地址范围。
测试与 Golden: 新增 test_prod_force_se_a.cpp(Kernel UT)、test_prod_force_se_a_tiling.cpp(Tiling UT)、test_prod_force_se_a_infershape.cpp(Infershape UT)及 Python 端 Golden 函数 golden.py 与数据生成/对比脚本,形成完整验证闭环。
💡 备注(选填)