十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

CANN ops-math AssignSub 算子深度解析:从算子定义到 aclnn 调用实践

CANN ops-math AssignSub 算子深度解析:从算子定义到 aclnn 调用实践 CANN ops-math AssignSub 算子深度解析从算子定义到 aclnn 调用实践【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读AssignSub 是 CANN ops-math 数学算子库中用于计算var - value并将结果写入输出的基础算子在权重更新如 SGD 类优化器等场景中扮演减法赋值的关键角色。本文以 experimental/math/assign_sub/README.md 为骨架结合算子定义、Shape 推导、Tiling、AscendC Kernel 与单测用例系统讲解其功能规格、参数约束、两段式 aclnn 调用方式与底层实现原理帮助你快速上手在 NPU 上调用该算子并理解其性能优化设计。一、产品支持情况AssignSub 算子当前支持以下产品形态详见 README 及 docs/aclnnAssignSub.md产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas A3 训练系列产品/Atlas A3 推理系列产品√产品形态详细说明请参见昇腾产品形态说明可查阅 CANN 官方文档的昇腾产品形态说明章节。从源码实现看该算子通过 op_host/assign_sub_def.cpp 中的AICore().AddConfig(ascend910b).AddConfig(ascend910_93)注册了对应的 AI Core 编译配置分别对应 Atlas A2Ascend 910B与 Atlas A3Ascend 910_93系列芯片与文档中的产品支持情况一一对应。二、功能说明AssignSub 的算子功能为计算var - value并将结果写入输出张量var_out。计算公式如下$$ var_out var - value $$从语义上看该算子等价于原地减法赋值类似var - value但输入输出通过独立张量传递适合在反向传播后的参数更新阶段以显式算子形式插入计算图。它的计算模式是**逐元素element-wise**运算两个输入张量在对应位置逐个做减法。三、参数说明3.1 算子参数参数名输入/输出说明var输入被减数张量数据类型支持 FLOAT16、INT8、FLOAT、INT32、UINT8、BF16、INT64数据格式支持 ND。value输入减数张量数据类型与 var 一致shape 与 var 一致数据格式支持 ND。var_out输出输出张量shape 与 var 一致数据类型与 var 一致数据格式支持 ND。上述数据类型与格式约束在 op_host/assign_sub_def.cpp 的算子定义中有完整体现var、value、var_out三个端口均声明为REQUIRED参数依次支持DT_FLOAT16、DT_INT8、DT_FLOAT、DT_INT32、DT_UINT8、DT_BF16、DT_INT64七种数据类型Format 与 UnknownShapeFormat 均限定为FORMAT_ND并统一调用.AutoContiguous()保证输入在计算前被处理为连续内存布局。3.2 aclnn 两段式接口参数在 aclnnAscendCL NN 算子调用方式下每个算子分为两段式接口必须先调用aclnnAssignSubGetWorkspaceSize获取工作空间大小并完成执行器创建再调用aclnnAssignSub执行计算aclnnStatus aclnnAssignSubGetWorkspaceSize( const aclTensor* var, const aclTensor* value, const aclTensor* varOut, uint64_t* workspaceSize, aclOpExecutor** executor);aclnnStatus aclnnAssignSub( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream);aclnnAssignSubGetWorkspaceSize的参数说明如下参数名输入/输出说明var输入被减数张量数据类型支持 FLOAT16、INT8、FLOAT、INT32、UINT8、BF16、INT64。value输入减数张量数据类型与 var 一致shape 与 var 一致。varOut输出输出张量shape 与 var 一致。workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。executor输出返回 op 执行器。需要说明的是从 op_host/assign_sub_tiling.cpp 的GetWorkspaceSize实现看AssignSub 当前将 workspace 大小置为 0WS_SYS_SIZE即正常调用时无需额外申请 workspace 空间示例代码中workspaceSize 0才分配内存的写法是 aclnn 调用的通用防御性范式可覆盖算子内部未来可能引入 workspace 需求的情况。四、约束与限制shape 与数据类型必须完全一致var 和 value 的 shape 及数据类型必须完全一致不支持 broadcast。这一点在 op_host/assign_sub_infershape.cpp 的 InferShape 中有强校验varShape与valueShape的维度数GetDimNum必须相等逐维GetDim(i)也必须相等任一不一致即通过OP_LOGE记录错误日志并返回GRAPH_FAILED。数据格式仅支持 ND。int8/uint8 的减法溢出按模 256 环绕处理即结果按 8 位无符号环绕语义计算这与 tests/ut/op_kernel/assign_sub_data/gen_data.py 中 golden 数据的生成逻辑一致——int8 结果先以int16做减法再按((result 128) % 256 - 128)回绕到 int8 范围uint8 则按result % 256回绕。int64 类型输入值范围限制在 int32 可表示范围内[-2^311, 2^31-1]这是因为 Kernel 内部使用 int32 作为 int64 的中间计算类型详见下文 Kernel 实现分析超出范围会产生截断导致的错误结果。五、调用说明与完整示例README 中给出了测试命令的调用方式可参考 docs/zh/invocation/quick_op_invocation.md 中描述的 build.sh 流程进行算子调用测试。目录描述examples/test_aclnn_assign_sub.cpp通过 aclnn 调用的方式调用 AssignSub 算子。5.1 完整可运行的 aclnn 调用示例仓库提供的 test_aclnn_assign_sub.cpp 是一个完整的最小可运行示例其主流程如下初始化环境aclInit(nullptr)初始化 ACLaclrtSetDevice(deviceId)设置设备示例中使用设备 0aclrtCreateStream(stream)创建流。构造张量示例构造{4, 5}形状的 FLOAT 张量var全部填充 10.0fvalue全部填充 3.0f预期输出为全 7.0f。CreateAclTensor模板函数完成 host 数据到 device 内存的搬运aclrtMallocaclrtMemcpy并按 ND 格式计算连续 strides 后通过aclCreateTensor创建aclTensor。两段式调用uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnAssignSubGetWorkspaceSize(var, value, varOut, workspaceSize, executor); CHECK_RET(ret ACLNN_SUCCESS, LOG_PRINT(aclnnAssignSubGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } ret aclnnAssignSub(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACLNN_SUCCESS, LOG_PRINT(aclnnAssignSub failed. ERROR: %d\n, ret); return ret); ret aclrtSynchronizeStream(stream);资源释放依次aclDestroyTensor销毁张量、aclrtFree释放 device 内存、aclDestroyOpExecutor销毁执行器、aclrtDestroyStream销毁流、aclrtResetDevice复位设备、aclFinalize结束 ACL。需要包含的头文件为#include acl/acl.h #include aclnnop/aclnn_assign_sub.h5.2 编译与测试算子的 CMake 构建入口位于 experimental/math/assign_sub/CMakeLists.txt并挂接在 experimental/math/CMakeLists.txt 的算子集合中。Tiling 单测位于 tests/ut/op_host/test_assign_sub_tiling.cpp测试用例assign_sub_0使用{3, 5, 7, 9}的 INT8 ND 张量、64 核、262144 字节 UB 的模拟环境断言 Tiling 结果为945 32 32依次对应 totalNum945、blockFactor32、ubFactor32且 workspace 为 0。Kernel 侧数据生成脚本 tests/ut/op_kernel/assign_sub_data/gen_data.py 使用{3, 5, 7, 9}形状的全 1 张量按上述环绕规则生成 golden 数据用于与 NPU 计算结果对比验证。六、源码级实现原理6.1 算子注册OpDefop_host/assign_sub_def.cpp 通过继承OpDef并在构造函数中声明输入输出端口完成算子注册最后由OP_ADD(AssignSub)注册到算子库。三个端口均声明为REQUIRED说明 var、value、var_out 在调用时缺一不可。6.2 Shape 推导InferShape / InferDataTypeop_host/assign_sub_infershape.cpp 中InferShapeAssignSub校验 var 与 value 的维度数和各维大小完全一致然后把varShape整体赋给输出即*outShape *varShape。这也从图编译层面保证了输出 shape 与 var 一致的规格。InferDataTypeAssignSubcontext-SetOutputDataType(0, context-GetInputDataType(0))将输出的数据类型直接继承自输入 var。6.3 Tiling 策略op_host/assign_sub_tiling.cpp 实现了 Tiling 函数AssignSubTilingFunc核心思路如下获取平台信息通过GetPlatformInfo读取 AIV 核数GetCoreNumAiv与 UB 内存大小GetCoreMemSize用于后续并行度与片上空间的计算。按数据类型选择调度模式GetDtypeInfo为七种数据类型分别生成 TilingKeyASSIGNSUB_TPL_SCH_MODE_0~MODE_6并给出对应的dtypeSize与perElemBytes。其中 INT8/UINT8 因需在计算中使用 half 中间类型并做模 256 处理额外计入2 * sizeof(uint16_t)的临时缓冲区开销BF16 计入2 * sizeof(float)INT64 计入2 * sizeof(int32_t)。计算三级因子写入 op_kernel/assign_sub_tiling_data.h 定义的AssignSubTilingData结构totalNum输入总元素数由存储 shape 的元素个数得到blockFactor每个核承担的元素块大小由CeilDiv(totalNum, coreNum)向上对齐到alignNum由 UB 块大小GetUbBlockSize除以 dtypeSize 得到算出并做 512 字节的粗粒度对齐尝试COARSE_ALIGN_BYTES以提升 GM 访问效率ubFactor单次流水一次 CopyIn/Compute/CopyOut 循环处理的元素数由ubSize / perElemBytes向下对齐到alignNum得到且不超过 blockFactor。设置执行参数SetBlockDim(usedCoreNum)设置实际使用的核数SetTilingKey(tilingKey)让 Kernel 侧按数据类型模板展开。从单测 tests/ut/op_host/test_assign_sub_tiling.cpp 的期望值945 32 32可以看到945 个元素被划分为每个核 32 个元素的块每个核内再按 32 个元素为一片ubFactor进行多轮流水处理。6.4 AscendC Kernel 实现op_kernel/assign_sub.cpp 使用if constexpr按 TilingKey即数据类型展开模板实例化NsAssignSub::AssignSubhalf/int8_t/float/int32_t/uint8_t/bfloat16_t/int64_t并统一调用op.Init(...)与op.Process()。op_kernel/assign_sub.h 中的核心设计流水结构TPipe配合三个TQue输入 var、输入 value、输出各 2 个 buffer构成双缓冲流水Process()按ubFactor分片循环每轮依次执行CopyInGM→UB非对齐场景使用DataCopyPad、Compute、CopyOutUB→GM通过EnQue/DeQue实现生产者-消费者同步。DtypeTrait 特化每种类型定义中间计算类型ComputeT与策略标志FLOAT16/FLOAT/INT32直接Sub无需类型转换INT8/UINT8先Cast到 half 做减法再经ShiftLeft/ShiftRight的 8 位算术移位实现模 256 环绕对应 README 中的溢出约束最后转回原类型BF16提升到 float 做减法结果用CAST_RINT舍入回 BF16INT64使用 int32 作为中间计算类型源码注释明确说明该设计因此输入值必须限制在 int32 可表示范围内超出会产生截断错误——这与 README 的 int64 范围约束完全对应。6.5 调度模式的模板参数化op_kernel/assign_sub_tiling_key.h 通过ASCENDC_TPL_ARGS_DECL与ASCENDC_TPL_SEL宏将schMode0~6对应七种数据类型声明为编译期模板参数使 Kernel 在编译时即可针对每种类型展开出最合适的指令序列避免了运行时的类型分支开销。七、贡献说明贡献者贡献方贡献算子贡献时间贡献内容Xzz西工大智能感知交互实验室AssignSub2026/07/12新增AssignSub算子八、小结AssignSub 是 ops-math 中结构清晰、规格精简的逐元素减法算子功能上实现var - value并输出到独立张量实现上由 assign_sub_def.cpp 完成算子注册、assign_sub_infershape.cpp 完成 shape/数据类型推导、assign_sub_tiling.cpp 完成多核与 UB 双缓冲调度、assign_sub.h 完成按类型的流水计算使用上通过aclnnAssignSubGetWorkspaceSizeaclnnAssignSub两段式接口即可在 Atlas A2/A3 系列产品上完成调用。需要注意三点规格限制不支持 broadcast、仅支持 ND 格式、int8/uint8 按模 256 环绕以及 int64 值域受 int32 中间类型限制。掌握该算子的调用范式后可以类推到 ops-math 中其他结构相近的 element-wise 算子。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表