已关闭
[Requirement|需求建议]: 增加 ProdForceSeA 支持 Ascend950 实现 SIMT #2451
ideal_创建于  7月31日关闭于  8月3日
ideal_
ideal_成员
7月31日 创建

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 与数据生成/对比脚本,形成完整验证闭环。

💡 备注(选填)

likedislike
ideal_ideal_成员
7月31日 添加了label:requirement
陈思
陈思成员
7月31日 评论:
likedislike
CANN-robotCANN-robot成员
7月31日 将 wang-shilong32 设为负责人
CANN-robotCANN-robot成员
8月3日 关闭了 issue
CANN-robotCANN-robot成员
8月3日 添加了label:resolved
ideal_ideal_成员
25 天前 issue状态由 进行中 改变为 已完成