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

资讯详情

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

CANN ops-math 中的 AtanGrad 算子:反正切梯度计算的原理、实现与 aclnn 调用实战

CANN ops-math 中的 AtanGrad 算子:反正切梯度计算的原理、实现与 aclnn 调用实战 CANN ops-math 中的 AtanGrad 算子反正切梯度计算的原理、实现与 aclnn 调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 仓库中 experimental/math/atan_grad 算子为核心系统讲解 AtanGrad反正切函数 atan 的输入梯度算子的数学原理、产品支持情况、参数与约束并结合仓库中的 op_api、op_host、op_kernel 源码剖析其 aclnn 两段式调用链路、Tiling 多核切分策略与 kernel 精度设计。读完本文你将掌握 AtanGrad 算子在 Ascend NPU 上的完整实现机制并能够参照仓库示例完成自定义算子包的编译、安装与 aclnn 接口调用。算子概述与数学原理AtanGrad 是用于神经网络反向传播的逐元素梯度算子计算反正切函数atan(x)对输入x的梯度。其核心公式为$$ dx_i dy_i \times \frac{1}{1 x_i^2} $$其中x为前向计算中 atan 函数的输入张量dy为上游loss 侧反向传入的梯度张量dx为算子输出的输入梯度张量与dy逐元素相乘后回传。该公式来源于 atan 的解析导数d(atan(x))/dx 1/(1x²)即dx dy · atan(x)。从实现角度看计算可以拆解为四步等效分步t_i x_i × x_i—— 计算x²g_i t_i 1.0—— 计算分母1 x²值恒不小于 1r_i 1 / g_i—— 取倒数即 atan 的导数dx_i dy_i × r_i—— 乘以上游梯度。这一分步在仓库的 kernel 实现 op_kernel/atan_grad.h 中被直接体现Mul求x*x、Adds加 1.0、再以Div完成dy/(1x²)。产品支持情况产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品×Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I / 500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×需要说明的是算子原型注册文件 op_host/atan_grad_def.cpp 中通过this-AICore().AddConfig(ascend910b, aiCoreConfig)与AddConfig(ascend950, aiCoreConfig)将算子绑定到 ascend910b对应 Atlas A2 系列与 ascend950对应 Ascend 950 系列两条 AICore 配置kernel 入口文件 op_kernel/atan_grad.cpp 也标注了 arch35 架构、支持 ascend910b/ascend950与 README 中的产品支持矩阵一致。参数说明AtanGrad 共 3 个参数两个输入、一个输出均为 ND 格式的稠密张量。参数名输入/输出描述数据类型数据格式x输入前向计算输入张量对应公式中 x为 atan 函数自变量。FLOAT16、FLOAT、BFLOAT16NDdy输入上游传入的梯度张量对应公式中 dy。数据类型须与 x 完全一致。FLOAT16、FLOAT、BFLOAT16NDdx输出输出的输入梯度张量对应公式中 dx。数据类型须与 x 完全一致。FLOAT16、FLOAT、BFLOAT16ND从源码看参数校验的落地aclnn 接口实现 op_api/aclnn_atan_grad.cpp 中的CheckParams将上述约束转化为四个显式的校验步骤CheckNotNull对 x、dy、dx 三个指针执行OP_CHECK_NULL任一为空即返回ACLNN_ERR_PARAM_NULLPTRCheckDtypeValid通过OP_CHECK_DTYPE_NOT_MATCH强制dy、dx的 dtype 与x完全一致再通过ATAN_GRAD_DTYPE_SUPPORT_LISTDT_FLOAT16、DT_FLOAT、DT_BF16白名单过滤不支持的类型CheckFormat通过IsPrivateFormat拒绝私有格式仅允许 ND 等公开格式CheckShape先用OP_CHECK_MAX_DIM将维度限制在ACLNN_MAX_SHAPE_RANK 8以内再强制x-GetViewShape() dy-GetViewShape()注释明确指出 shape 不一致会带来 GM 越界访问风险。在 L0 层 op_api/atan_grad.cpp 中AtanGradInferShape直接将输出 shape 置为输入 x 的 shape逐元素算子天然如此并由IsAiCoreSupport再次核对 dtype 白名单。约束说明aclnnAtanGrad 为默认确定性实现算子行为确定同一输入多次运行结果一致。x、dy、dx 三者数据类型必须完全一致不支持隐式类型转换。x、dy、dx 三者 shape 必须完全相同不支持广播broadcast。支持空 Tensor元素个数为 0。支持 0-8 维 Tensor0 维表示标量scalar此时 dy 和 dx 也必须为 0 维。当 x 取值极大如 fp16 最大值时x²可能溢出为 inf此时1/inf 0dx 0属于正常数值行为。空 Tensor 与 0 维标量的特殊处理空 TensoraclnnAtanGradGetWorkspaceSize中通过HasEmptyTensor(x, dy)提前检测若任一输入为空直接将*workspaceSize 0并返回成功跳过后续建图与调度0 维标量Tiling 侧 op_host/atan_grad_tiling.cpp 用EnsureNotScalar将 0 维 shape 归一为{1}从而与常规向量路径复用同一套多核切分逻辑总元素为 0Tiling 中totalNum 0时设置blockDim 0并提前返回避免空启动。aclnn 两段式接口调用调用方式调用方式调用样例说明aclnn 调用test_aclnn_atan_grad调用前需完成自定义算子包的编译与安装bash build.sh --socascend910b。README 当前标注接口文档和设计文档待补充因此示例代码是本算子最权威的调用参考。两段式接口的执行链路从 op_api/aclnn_atan_grad.cpp 的文件头注释可以看到 AtanGrad 采用标准的 ACLNN L2 两段式接口第一段aclnnAtanGradGetWorkspaceSizeCREATE_EXECUTOR→CheckParams→HasEmptyTensor→Contiguous→l0op::AtanGrad→ViewCopy→ 返回workspaceSize与executor第二段aclnnAtanGrad通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)在指定 stream 上真正执行。值得注意的两个实现细节输入会先经过l0op::Contiguous归一为连续内存算子计算完成后通过l0op::ViewCopy将结果写入用户传入的 dx 张量视图L0 层AtanGrad内部依次完成AtanGradInferShape输出 shape 输入 shape、IsAiCoreSupportdtype 与平台支持确认、executor-AllocTensor分配输出张量、AtanGradAiCore经ADD_TO_LAUNCHER_LIST_AICORE调度 kernel四个步骤。示例编译与运行仓库示例 test_aclnn_atan_grad.cpp 以 fp32、shape[4, 8]的测试数据演示了完整流程编译算子包cd ops/atan_grad bash build.sh --socascend910b安装指令参照 build.sh 输出编译示例g -stdc17 -o test_aclnn_atan_grad test_aclnn_atan_grad.cpp \ -I${ASCEND_TOOLKIT_HOME}/include \ -L${ASCEND_TOOLKIT_HOME}/lib64 \ -lacl_op_compiler -lascendcl \ -Wl,-rpath,${ASCEND_TOOLKIT_HOME}/lib64运行./test_aclnn_atan_grad。示例程序的关键步骤依次为aclInit与aclrtSetDevice初始化 ACL →aclrtMalloc在 Device 侧分配内存并通过aclrtMemcpy上传 x、dy →aclCreateTensor构造 ND 格式 aclTensor → 调用aclnnAtanGradGetWorkspaceSize获取 workspace 大小逐元素算子通常为 0→ 按需分配 workspace → 调用aclnnAtanGrad执行 →aclrtSynchronizeStream等待完成 → 结果拷回 Host。精度验证方式示例在 Host 侧实现了 CPU Golden 参考实现ComputeGolden用 double 精度计算dy * (1.0 / (1.0 x*x))再与 NPU 输出逐元素比较最大相对误差MARE并以 fp32 阈值10 * 2^-13 ≈ 0.00122判定 PASS/FAIL。这为在自定义场景中验证算子精度提供了可复制的模板。kernel 实现与精度设计数据流与缓冲kernel 类实现在 op_kernel/atan_grad.h采用经典的CopyIn → Compute → CopyOut流水结构CopyIn通过DataCopyPad配合DataCopyExtParamsblockLen为 uint32_t规避DataCopyParams.blockLen仅 uint16_t、最大 65535 字节的限制将 GM 上的 x、dy 搬入 UBCompute在 UB 中完成核心计算CopyOut将结果 dx 从 UB 写回 GM。缓冲策略由模板参数BUFFER_MODE决定BUFFER_NUM BUFFER_MODE ? 2 : 1即 0 为单缓冲、1 为双缓冲通过TQue与TBuf实现数据搬运与计算的异步流水。按 dtype 分叉的精度方案kernel 注释明确记录了穿刺验证得到的精度结论三种数据类型采用不同计算路径fp32 路径直接计算。ReciprocalINTRINSIC 模式精度不足MERE≈2.8e-3因此改用Div(dx, dy, tmp)一步完成dy/(1x²)将 MERE 降至 1.19e-7fp16 / bf16 路径升精度到 fp32 计算ComputeUpcast。先用Cast(CAST_NONE)将 x、dy 无损升到 fp32在 fp32 下完成Mul → Adds → Div四步计算后再用Cast(CAST_RINT)银行家舍入降回原精度。fp16 直接使用 Reciprocal 时 MERE≈1.07e-3 超阈值必须升精度bf16 采用Cast(CAST_NONE) fp32 计算 Cast(CAST_RINT)后 MERE1.091e-3 达标。对应地Process中用if constexpr (std::is_same_vT, float)在编译期分派fp32 走Compute其余走ComputeUpcast。模板组合与 kernel 入口模板参数在 op_kernel/atan_grad_tiling_key.h 中声明共 6 个组合fp16×SB、fp16×DB、fp32×SB、fp32×DB、bf16×SB、bf16×DB。kernel 入口 op_kernel/atan_grad.cpp 通过REGISTER_TILING_DEFAULT与GET_TILING_DATA_WITH_STRUCT读取 tiling 参数实例化NsAtanGrad::AtanGradD_T_X, BUFFER_MODE后执行Init与Process。Tiling 多核切分策略Tiling 逻辑集中在 op_host/atan_grad_tiling.cpp核心步骤可概括为获取平台信息通过PlatformAscendC取得 AIV Core 数量coreNum与 UB 大小ubSize多核切分blockFactor CeilAlign(CeilDiv(totalNum, coreNum), ubBlockSize)即把总元素数按核心数均分并对齐到 DMA 最小粒度32B / sizeof(T)实际使用核数为CeilDiv(totalNum, blockFactor)双缓冲决策useDoubleBuffer (totalNum MIN_SPLIT_THRESHOLD) ? 1 : 0阈值常量MIN_SPLIT_THRESHOLD 1024元素超过 1024 启用双缓冲以隐藏搬运延迟UB 切分ubFactor FloorAlign(FloorDiv(ubSize / typeSize, bufferNum), ubBlockSize)。注意 fp32 与 fp16/bf16 的bufferNum不同——fp32 直接计算只需 4单缓冲/7双缓冲个 buffer而 fp16/bf16 升精度路径需要 xFp32、dyFp32、tmp、dxFp32 共 4 个 fp32 临时 buffer因此为 7单缓冲/10双缓冲ubFactor下限被钳制为ubBlockSize以保证 CopyIn 正常workspace 置 0逐元素算子无需额外 workspaceSetWorkspace将currentWorkspace[0]置 0这也是示例中 workspaceSize 通常为 0 的根因模板选择ASCENDC_TPL_SEL_PARAM(context, dTypeX, useDoubleBuffer)根据 dtype × 缓冲模式实例化对应 kernel 模板。TilingData 结构体定义在 op_kernel/atan_grad_tiling_data.h包含三个字段totalNum总元素数量、blockFactor每核负责的元素数量、ubFactor每次 UB 循环处理的元素数量。算子原型注册op_host/atan_grad_def.cpp 通过OP_ADD(AtanGrad)注册算子原型三个张量x、dy、dx均为REQUIREDdtype 限定DT_FLOAT16 / DT_FLOAT / DT_BF16Format 限定 ND且均声明AutoContiguousOpAICoreConfig开启DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)并关闭NeedCheckSupportFlag声明PrecisionReduceFlag(true)通过AddConfig将算子绑定到 ascend910b 与 ascend950 两个平台构建入口 op_host/CMakeLists.txt 中的add_modules_sources(OPTYPE atan_grad ACLNNTYPE aclnn)将该算子的 op_type 与 aclnn 接口一并纳入模块编译。总结AtanGrad 是一个典型的逐元素反向梯度算子数学上以dx dy/(1x²)完成 atan 的梯度回传工程上则以算子原型注册 → Tiling 多核/UB 切分 → kernel 模板实例化 → aclnn 两段式接口的完整链路落地到 Ascend NPU。其实现中的两个亮点值得在开发其他梯度算子时复用一是按 dtype 分叉的精度策略fp32 用 Div 替代 Reciprocal、fp16/bf16 升精度计算 银行家舍入回降二是确定性、零 workspace 的逐元素设计配合[4, 8]形状的 CPU Golden 精度比对示例可作为快速验证与二次开发的直接参考。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表