CANN ops-math AssignSub 算子深度解析：原址减法实现、图模式调用与源码级原理

CANN ops-math AssignSub 算子深度解析原址减法实现、图模式调用与源码级原理【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 数学算子库中的 AssignSub 算子为讲解主体围绕其功能定义、参数约束、产品适配与图模式调用展开并结合仓库内 op_graph / op_host / op_kernel / framework 四个层面的源码实现帮助读者既能在实际工程中正确调用 AssignSub又能理解该算子从算子原型、Tiling 计算到 NPU Kernel 执行的完整链路。一、算子概述与核心公式AssignSub 是 CANN ops-math 提供的原址in-place减法算子在已有张量var的基础上直接完成减法计算结果写回var所在的内存地址不需要额外分配输出张量。这一点与普通的 elementwise 减法算子如sub会生成新的输出张量有本质区别因此它天然适合用作训练/推理图中的变量Variable更新节点。其数学定义非常简单文档给出的计算公式为out_i var_i - value_i其中var_i是被更新的原张量元素value_i是要减去的值out_i与var_i共享同一块地址。这一行为在算子原型注释中也有明确描述Updates var by subtracting value from it. This operation outputs var after the update is done且明确指出该算子与 TensorFlow 的 AssignSub 算子兼容见 math/assign_sub/op_graph/assign_sub_proto.h。二、产品支持情况当前仓库中的 AssignSub 实现适配了 Ascend 主流训练/推理产品线支持情况如下表产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品√从源码角度可以印证上述适配范围算子的 AICore 配置在 math/assign_sub/op_host/assign_sub_def.cpp 中通过this-AICore().AddConfig(ascend950, aicoreConfig)和AddConfig(ascend350, aicoreConfig)注册对应支持 asceng 950 与 350 系列平台Tiling 与 Kernel 实现则位于arch35目录math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp、math/assign_sub/op_kernel/arch35/assign_sub_dag.h。需要说明的是产品支持矩阵以当前仓库实现为准具体到目标环境时仍建议以对应版本的《CANN 算子支持矩阵》做最终确认。三、参数说明AssignSub 算子共包含两个输入、一个输出和一个可选属性完整参数定义如下表数据范围与算子原型REG_OP(AssignSub)的声明严格一致见 math/assign_sub/op_graph/assign_sub_proto.h参数名输入/输出/属性描述数据类型数据格式var输入输入张量公式中的var_iBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvalue输入输入张量表示要从var_i上减去的值公式中的value_iBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvar输出输出张量公式中的out_i与var_i同地址BFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDuse_locking属性可选是否使用锁来保护更新操作bool—在使用时需要注意以下几点数据类型一致性var、value与输出var必须为同一数据类型。这一约束不仅在算子原型中体现还会在 Tiling 阶段被显式校验AssignSubTiling::CheckDtype()会逐一取出三个张量的 dtype 并比较若不一致会记录The dtypes of var, value and output must be the same错误并返回失败见 math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp。Shape 一致性var、value、输出三者的 shape 也必须完全一致CheckShape()中通过varShape ! valueShape || varShape ! outputShape进行校验同一文件 L74-L96。这与该算子使用逐元素elewise推断规则是一致的。use_locking 属性可选 bool默认值为False。原型注释说明当其为True时减法更新操作将由锁保护避免多线程并发下的竞争为False时行为未定义但可以减少锁竞争带来的开销。当前仓库的 Tiling/Kernel 实现中未对锁语义做额外处理属性主要随图下发给框架侧。约束说明原文档明确约束说明无即在已支持的 dtype / ND 格式范围内没有额外使用约束。四、图模式调用单算子构图完整样例AssignSub 支持图模式调用即通过算子 IR 构图GE Graph 接口的方式在计算图中插入 AssignSub 节点后运行。仓库在 math/assign_sub/examples/test_geir_assign_sub.cpp 中提供了完整的可运行样例下面拆解其关键流程。4.1 定义算子 IR 与构图首先通过ge::op::AssignSub创建算子实例并为两个输入分别创建Data占位节点、设置输入/输出 TensorDescauto assignSub1 op::AssignSub(assignSub1); std::vectorint64_t xShape {1, 16}; std::vectorint64_t yShape {1, 16}; ADD_INPUT(1, var, inDtype, yShape); // 输入 varshape 为 {1, 16} ADD_INPUT(2, value, inDtype, xShape); // 输入 valueshape 为 {1, 16} ADD_INPUT_ATTR(use_locking, false); // 设置属性 use_locking false ADD_OUTPUT(1, var, inDtype, yShape); // 输出 var与输入 var 同 shape样例中的宏展开后会为每个输入调用assignSub1.set_input_var(...)/assignSub1.set_input_value(...)建立数据流边并通过assignSub1.update_output_desc_var(...)声明输出描述。这里示例使用DT_FLOAT16FP16作为输入数据类型shape 取{1, 16}。4.2 初始化 GE 与会话运行构图完成后样例依次完成 GE 初始化、建图、执行三步// 1. 初始化 GE指定 deviceId 与图运行模式 std::mapAscendString, AscendString global_options {{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}}; Status ret ge::GEInitialize(global_options); // 2. 创建 Session 并添加计算图 ge::Session* session new Session(build_options); uint32_t graph_id 0; ret session-AddGraph(graph_id, graph, graph_options); // 3. 运行图输入/输出均为 ge::Tensor std::vectorge::Tensor output; ret session-RunGraph(graph_id, input, output);其中ge.graphRunMode1表示按图运行模式执行运行结束后样例还会通过aclgrphDumpGraph(graph, ./dump, ...)将图结构 dump 到本地便于调试查看构图是否正确。4.3 数据生成与结果落盘样例内置了GenOnesData辅助函数用于构造输入数据根据 shape 计算元素个数结合GetDataTypeSize换算字节数将数据全部初始化为固定值后封装为ge::Tensor。运行完成后输入与输出张量会分别写入tc_ge_irrun_test_0008_npu_input_{i}.bin与tc_ge_irrun_test_0008_npu_output_{i}.bin并在终端逐元素打印输出结果方便与预期var - value结果比对。从该样例可以看到图模式调用的标准姿势op::AssignSub构图 →GEInitialize→Session::AddGraph→RunGraph→GEFinalize这套流程同样适用于 ops-math 中其它算子。五、源码级原理从原型注册到 NPU Kernel 的完整链路了解了调用方式后再深入一层看 AssignSub 在仓库内部是如何被定义、分派与执行的。整条链路可以概括为算子原型 → Host 侧定义/推断 → Tiling 计算 → Kernel 执行四个阶段。5.1 算子原型op_graphmath/assign_sub/op_graph/assign_sub_proto.h 使用REG_OP宏声明算子原型REG_OP(AssignSub) .INPUT(var, TensorType::NumberType()) .INPUT(value, TensorType::NumberType()) .OUTPUT(var, TensorType::NumberType()) .ATTR(use_locking, Bool, false) .OP_END_FACTORY_REG(AssignSub)注意输入与输出同名var这正是原址更新语义在原型层面的体现输出与输入var指向同一块数据。5.2 Host 侧算子定义op_hostmath/assign_sub/op_host/assign_sub_def.cpp 通过OpDef注册了算子的运行时定义包括输入/输出的数据类型白名单DT_BF16, DT_FLOAT16, DT_FLOAT, DT_INT8, DT_INT32, DT_INT64, DT_UINT8数据格式均为FORMAT_ND属性use_lockingOPTIONAL类型默认falseAICore 配置开启动态编译DynamicCompileStaticFlag(true)、动态 rank 与动态 shape 支持DynamicRankSupportFlag(true)/DynamicShapeSupportFlag(true)并将 Kernel 文件指向assign_sub_apt。Shape 推断方面math/assign_sub/op_host/assign_sub_infershape.cpp 直接复用逐元素通用推断逻辑Ops::Base::InferShape4Elewise——这再次印证了 AssignSub 在 shape 语义上与 elementwise 算子一致输入输出 shape 相同。对应的单测 math/assign_sub/tests/ut/op_host/test_assign_sub_infershape.cpp 验证了 FP16、ND 格式下输入输出 shape 推断的结果。5.3 Tiling 计算op_host/arch35在 NPU 上执行前Host 侧需要先完成 Tiling数据分片计算。math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp 实现了AssignSubTiling::RunTiling()其执行顺序为dtype 校验三个张量 dtype 必须一致CheckDtypeshape 校验三个张量 shape 必须一致CheckShape按 dtype 选择分片策略通过Ops::Base::ElewiseBaseTiling::DoTilingAssignSubOpT::OpDag()生成基础 Tiling 数据baseTiling设置 Tiling KeySetTilingData()依据输出 dtype 写入对应的 Tiling KeyFP16/BF16/FP32/INT8/INT32/INT64/UINT8 分别映射为101/102/103/104/105/106/107申请 Workspace设置系统 Workspace 大小SYS_WORKSPACE_SIZE 16MB供 Kernel 侧使用。Tiling 阶段还通过TilingPrepareForAssignSub从平台信息中读取 AIV 核数与 UB 内存大小GetCoreNumAiv()/GetCoreMemSize(CoreMemType::UB)用于决定分块粒度与核数分配。5.4 Kernel 执行op_kernelNPU 侧入口在 math/assign_sub/op_kernel/assign_sub_apt.cpp其assign_sub内核函数以GM_ADDR x1/x2/y分别接收var、value与输出地址核心逻辑是一个按 Tiling Key 分派的 elementwise 调度器if (TILING_KEY_IS(101UL)) { ElementwiseSch0UL, AssignSubOphalf::OpDag sch((tilingData.baseTiling), pipe); sch.Init(x1, x2, y); sch.Process(); } else if (TILING_KEY_IS(103UL)) { ElementwiseSch0UL, AssignSubOpfloat::OpDag sch(...); ... }真正执行减法运算的计算图在 math/assign_sub/op_kernel/arch35/assign_sub_dag.h 中定义通过CopyIn0 / CopyIn1将两个输入从 GM 搬入 UBOpCopyIn自定义向量减法节点SubCustom使用寄存器指令Reg::Sub(vregOutput, vregInput1, vregInput2, mask)完成a - b其中向量宽度按数据类型VECTOR_REG_WIDTH / sizeof(T)计算并通过CeilDivision(count, vl)分循环处理剩余元素CopyOut将结果写回 GM 输出地址最终由DAGSchOutputs组成可调度的计算 DAG。Tiling 数据本身在 math/assign_sub/op_kernel/arch35/assign_sub_tiling_struct.h 中定义为AssignSubTilingData { EleBaseTilingData baseTiling; }即直接复用 elementwise 基础 Tiling 结构这也是 AssignSub 与仓库中其他 elementwise 算子共享同一套调度框架atvoss/elewise的体现。六、框架适配TensorFlow 算子映射除了图模式直接构图AssignSub 还通过框架插件完成了与第三方框架的对接。math/assign_sub/framework/assign_sub_tf_plugin.cpp 中使用REGISTER_CUSTOM_OP将自定义算子与 TensorFlow 侧算子建立映射REGISTER_CUSTOM_OP(AssignSub) .FrameworkType(TENSORFLOW) .OriginOpType(std::vectorge::AscendString{ge::AscendString(AssignSub), ge::AscendString(AssignSubVariableOp)}) .ParseParamsByOperatorFn(AutoMappingByOpFn) .ImplyType(ImplyType::TVM);这里同时映射了 TensorFlow 的AssignSub与AssignSubVariableOp两个原始算子类型后者对应tf.Variable的assign_sub方法参数通过AutoMappingByOpFn自动映射无需手写逐字段转换逻辑。这意味着在 PyTorch/TensorFlow 生态中以var.assign_sub(value)这类语义发起的减法更新可以无缝落到 NPU 上的 AssignSub 内核。七、测试与验证仓库为 AssignSub 提供了 Host 侧单测与 ST 用例可作为功能正确性的验证参考Infershape 单测math/assign_sub/tests/ut/op_host/test_assign_sub_infershape.cpp 构造{ -1 }动态 shape 的 FP16 输入断言推断输出 shape 与预期一致且返回GRAPH_SUCCESSTiling 单测math/assign_sub/tests/ut/op_host/arch35/test_assign_sub_tiling.cpp 覆盖 arch35 平台上的 Tiling 计算ST 用例math/assign_sub/tests/st/arch35/ttk_kernel_assign_sub_st.csv 定义了 Kernel 级 ST 用例矩阵配合 math/assign_sub/tests/assets/golden.py 的 golden 数据生成脚本可批量校验内核输出。八、小结AssignSub 虽然只是一个单输入对的原址减法算子但其仓库实现完整覆盖了从算子原型op_graph、Host 定义与推断op_host、Tiling 计算arch35、NPU Kernelop_kernel到框架适配framework的全栈链路是理解 ops-math 中 elementwise 类算子实现范式的极佳样本。实际使用时只需记住三个要点输入输出 dtype 与 shape 必须一致、数据格式为 ND、通过图模式GE IR或 TensorFlow 插件路径发起调用即可在支持的 Ascend 产品上正确完成原址减法更新。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考