企业官网建设流程全解析

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 主流训练/推理产品线，支持情况如下表：

从源码角度可以印证上述适配范围：算子的 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）：

在使用时需要注意以下几点：

1. 数据类型一致性：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）。
2. Shape 一致性：var、value、输出三者的 shape 也必须完全一致，CheckShape()中通过varShape != valueShape || varShape != outputShape进行校验（同一文件 L74-L96）。这与该算子使用逐元素（elewise）推断规则是一致的。
3. use_locking 属性：可选 bool，默认值为False。原型注释说明：当其为True时，减法更新操作将由锁保护，避免多线程并发下的竞争；为False时行为未定义，但可以减少锁竞争带来的开销。当前仓库的 Tiling/Kernel 实现中未对锁语义做额外处理，属性主要随图下发给框架侧。
4. 约束说明：原文档明确"约束说明：无"，即在已支持的 dtype / ND 格式范围内没有额外使用约束。

四、图模式调用：单算子构图完整样例

AssignSub 支持图模式调用，即通过算子 IR 构图（GE Graph 接口）的方式在计算图中插入 AssignSub 节点后运行。仓库在 math/assign_sub/examples/test_geir_assign_sub.cpp 中提供了完整的可运行样例，下面拆解其关键流程。

4.1 定义算子 IR 与构图

首先通过ge::op::AssignSub创建算子实例，并为两个输入分别创建Data占位节点、设置输入/输出 TensorDesc：

auto assignSub1 = op::AssignSub("assignSub1"); std::vector<int64_t> xShape = {1, 16}; std::vector<int64_t> yShape = {1, 16}; ADD_INPUT(1, var, inDtype, yShape); // 输入 var，shape 为 {1, 16} ADD_INPUT(2, value, inDtype, xShape); // 输入 value，shape 为 {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_FLOAT16（FP16）作为输入数据类型，shape 取{1, 16}。

4.2 初始化 GE 与会话运行

构图完成后，样例依次完成 GE 初始化、建图、执行三步：

// 1. 初始化 GE，指定 deviceId 与图运行模式 std::map<AscendString, 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::vector<ge::Tensor> output; ret = session->RunGraph(graph_id, input, output);

其中ge.graphRunMode=1表示按图运行模式执行；运行结束后样例还会通过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_graph）

math/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_host）

math/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_locking：OPTIONAL类型，默认false；
• AICore 配置：开启动态编译（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()，其执行顺序为：

1. dtype 校验：三个张量 dtype 必须一致（CheckDtype）；
2. shape 校验：三个张量 shape 必须一致（CheckShape）；
3. 按 dtype 选择分片策略：通过Ops::Base::ElewiseBaseTiling::DoTiling<AssignSubOp<T>::OpDag>()生成基础 Tiling 数据baseTiling；
4. 设置 Tiling Key：SetTilingData()依据输出 dtype 写入对应的 Tiling Key，FP16/BF16/FP32/INT8/INT32/INT64/UINT8 分别映射为101/102/103/104/105/106/107；
5. 申请 Workspace：设置系统 Workspace 大小（SYS_WORKSPACE_SIZE = 16MB），供 Kernel 侧使用。

Tiling 阶段还通过TilingPrepareForAssignSub从平台信息中读取 AIV 核数与 UB 内存大小（GetCoreNumAiv()/GetCoreMemSize(CoreMemType::UB)），用于决定分块粒度与核数分配。

5.4 Kernel 执行（op_kernel）

NPU 侧入口在 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)) { ElementwiseSch<0UL, AssignSubOp<half>::OpDag> sch(&(tilingData.baseTiling), &pipe); sch.Init(x1, x2, y); sch.Process(); } else if (TILING_KEY_IS(103UL)) { ElementwiseSch<0UL, AssignSubOp<float>::OpDag> sch(...); ... }

真正执行减法运算的计算图在 math/assign_sub/op_kernel/arch35/assign_sub_dag.h 中定义：

• 通过CopyIn0 / CopyIn1将两个输入从 GM 搬入 UB（OpCopyIn）；
• 自定义向量减法节点SubCustom使用寄存器指令Reg::Sub(vregOutput, vregInput1, vregInput2, mask)完成a - b，其中向量宽度按数据类型VECTOR_REG_WIDTH / sizeof(T)计算，并通过CeilDivision(count, vl)分循环处理剩余元素；
• CopyOut将结果写回 GM 输出地址，最终由DAGSch<Outputs>组成可调度的计算 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::vector<ge::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_SUCCESS；
• Tiling 单测：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 Kernel（op_kernel）到框架适配（framework）的全栈链路，是理解 ops-math 中 elementwise 类算子实现范式的极佳样本。实际使用时只需记住三个要点：输入输出 dtype 与 shape 必须一致、数据格式为 ND、通过图模式（GE IR）或 TensorFlow 插件路径发起调用，即可在支持的 Ascend 产品上正确完成原址减法更新。

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math