CANN ops-math AcosGrad 算子解析：反余弦反向梯度计算原理、参数约束与 NPU 源码实现-平芜编程栈

算子库
人工智能
CANN

【免费下载链接】ops-math

本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。

AcosGrad 是 CANN ops-math 算子库中为反余弦（Acos）前向算子配套的反向梯度算子，用于在 NPU 上高效计算acos(x)关于输入的梯度。本文以 math/acos_grad/README.md 为主线，完整梳理该算子的功能公式、产品支持矩阵、输入输出参数、约束规则与图模式调用方法，并结合仓库内算子定义、shape/dtype 推导、tiling 分片和昇腾 Kernel 源码，逐层还原其在 Ascend 硬件上的真实实现细节，帮助读者既能直接上手调用，又能深入理解一个标准单算子从「算子原型 → 宿主侧推导 → 核函数实现」的完整落地链路。

产品支持情况

AcosGrad 算子已在 CANN ops-math 当前仓库中登记了对以下产品形态的支持（详见 README）：

产品	是否支持
Ascend 950PR / Ascend 950DT	√
Atlas A3 训练系列产品 / Atlas A3 推理系列产品	√
Atlas A2 训练系列产品 / Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	√
Atlas 推理系列产品	√
Atlas 训练系列产品	√

从算子构建配置 math/acos_grad/CMakeLists.txt 可以看到，本仓库中 AcosGrad 的宿主侧（host）tiling 与算子定义当前以ascend950为支持的计算单元，对应arch35架构目录（SUPPORT_TILING_DIR "arch35"），即上述支持列表在代码层面的具体落地形态。需要说明的是，该清单以当前仓库登记为准，实际可用性请以随版本发布的官方支持矩阵为最终依据。

功能说明

算子功能与公式

AcosGrad 计算的是Acos（反余弦）前向算子的反向梯度，即给定前向输入y与上游传入的梯度dy，求出对原始输入张量的梯度z。其逐元素计算公式为：

$$ z_i = -1 \cdot dy_i \cdot \dfrac{1}{\sqrt{1 - y_i^2}} $$

其中：

$y_i$：前向 Acos 算子的输入张量，值域期望落在 $[-1, 1]$；
$dy_i$：上游（loss 侧）传入的梯度；
$z_i$：对原始输入张量的梯度，等于上游梯度乘以 $-1/\sqrt{1 - y_i^2}$。

超出定义域的行为

由于反余弦的定义域为 $[-1, 1]$，当输入超出该范围（$|y_i| > 1$）时，$1 - y_i^2 < 0$，对负数开平方根的结果为 NaN，因此最终输出梯度同样为 NaN。这一行为在 Kernel 实现的注释中亦有明确说明，属于预期的数值语义，而非错误分支。

参数说明

AcosGrad 共包含 2 个输入和 1 个输出，全部要求 ND 格式。详细参数定义如下：

参数名	输入/输出/属性	描述	数据类型	数据格式
y	输入	前向 Acos 算子的输入张量。值域期望落在 [-1, 1]。	FLOAT16, FLOAT32, BFLOAT16	ND
dy	输入	上游传入的梯度张量，shape 与 dtype 与 y 一致。	FLOAT16, FLOAT32, BFLOAT16	ND
z	输出	对原始输入张量的梯度，shape 与 dtype 与 y 一致。	FLOAT16, FLOAT32, BFLOAT16	ND

上述接口在代码层有两处权威定义，二者保持一致：

图 IR 原型定义math/acos_grad/op_graph/acos_grad_proto.h 中通过REG_OP(AcosGrad)声明了INPUT(y)、INPUT(dy)、OUTPUT(z)三个端口，TensorType限定为{DT_BF16, DT_FLOAT16, DT_FLOAT}；
算子定义（OpDef）math/acos_grad/op_host/acos_grad_def.cpp 中为每个端口声明了相同的数据类型列表ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16与格式列表ge::FORMAT_ND，同时将三个端口均标记为REQUIRED（必选），并开启AutoContiguous自动连续化处理。

值得注意的是，acos_grad_proto.h 的注释中还说明了该算子与TensorFlow 的 AcosGrad 算子兼容，这为框架侧迁移与语义对齐提供了依据。

约束说明

根据 README，使用 AcosGrad 算子必须满足以下约束：

y与dy的shape 必须完全一致；
y与dy的dtype 必须完全一致；
仅支持 ND 格式；
支持非连续 Tensor，且非连续 Tensor 的维度不大于 8。

这些约束不仅停留在文档层面，在代码中也有强校验：

acos_grad_def.cpp 中将y、dy的UnknownShapeFormat均限制为FORMAT_ND，并在 AICore 配置中开启DynamicShapeSupportFlag(true)与DynamicRankSupportFlag(true)，说明该算子支持动态 shape 与动态 rank（这正好呼应了"支持非连续 Tensor、维度不大于 8"的能力）；
Tiling 阶段 acos_grad_tiling_arch35.cpp 会校验yShape.GetShapeSize() == dyShape.GetShapeSize() == zShape.GetShapeSize()，三者元素总数不一致时直接报错返回GRAPH_FAILED，同时校验输入 dtype 必须落在{DT_FLOAT16, DT_FLOAT, DT_BF16}集合内。

从算子原型到 NPU 执行的完整链路

AcosGrad 的落地链路与 CANN 单算子标准流程一致，可分为宿主侧（host）与设备侧（device）两部分，仓库目录结构即对应这一分层：

math/acos_grad/ ├── op_graph/ # 算子 IR 原型、图推导（InferShape / InferDataType） ├── op_host/ # 算子定义、shape 推导、tiling 分片（含 arch35） ├── op_kernel/ # 昇腾核函数实现（含 arch35 模板与 tiling 数据结构） ├── framework/ # 框架侧插件适配 ├── examples/ # 图模式调用样例 └── tests/ # 单元测试

图推导：shape 与 dtype 的确定

Shape 推导：acos_grad_infershape.cpp 中的InferShape4AcosGrad直接从输入端口 0 取出y的 shape，并将输出z的 shape 直接赋值为*yShape，即"输出 shape 等于输入 y 的 shape"，与文档中"z 的 shape 与 y 一致"完全对应；
数据类型推导：acos_grad_graph_infer.cpp 中的InferDataType4AcosGrad将输出z的 dtype 设置为输入y的 dtype，保证输出与输入类型一致。

算子定义中的 AICore 配置

acos_grad_def.cpp 中还通过OpAICoreConfig配置了关键运行参数：

OpAICoreConfig aiCoreConfig; aiCoreConfig.DynamicCompileStaticFlag(true) .DynamicFormatFlag(false) .DynamicShapeSupportFlag(true) .PrecisionReduceFlag(true) .NeedCheckSupportFlag(false) .DynamicRankSupportFlag(true) .ExtendCfgInfo("opFile.value", "acos_grad_apt"); this->AICore().AddConfig("ascend950", aiCoreConfig);

其中ExtendCfgInfo("opFile.value", "acos_grad_apt")指明了设备侧核函数实现文件为 acos_grad_apt.cpp，PrecisionReduceFlag(true)表示允许精度优化策略。

Kernel 入口与实现

核函数入口 acos_grad_apt.cpp 以__global__ __aicore__形式暴露acos_grad(y, dy, z, workspace, tiling)，通过GET_TILING_DATA_WITH_STRUCT(AcosGradTilingData, tilingData, tiling)解析宿主侧传入的 tiling 数据，并以编译期注入的DTYPE_Y宏实例化模板NsAcosGrad::KernelAcosGrad<DTYPE_Y>。

实际计算逻辑位于 math/acos_grad/op_kernel/arch35/acos_grad.h，其核心 Compute 流程严格按公式展开为向量指令序列：

Mul(tmpF32A, yF32, yF32, currentNum); // y^2 Muls(tmpF32A, tmpF32A, -1.0f, currentNum); // -y^2 Adds(tmpF32B, tmpF32A, 1.0f, currentNum); // 1 - y^2 Sqrt(tmpF32B, tmpF32B, currentNum); // sqrt(1 - y^2) Muls(dyF32, dyF32, -1.0f, currentNum); // -dy Div(dyF32, dyF32, tmpF32B, currentNum); // -dy / sqrt(1 - y^2)

实现中还体现了两个值得关注的工程细节：

精度路径分治：对于 FLOAT32 直接计算；对于 FLOAT16 / BFLOAT16，则先Cast到 float32 完成全部中间运算，最后再以CAST_RINT舍入模式Cast回原类型，从而在低精度输入下尽量保证中间结果的数值精度；
流水与缓冲：Kernel 使用TPipe与双缓冲队列（BUFFER_NUM = 2）组织CopyIn → Compute → CopyOut流水，输入y、dy与输出z各自独立队列，同时在 VECCALC 区申请 float32 临时缓冲，支撑向量级数据搬运与计算的并行。

Tiling 分片策略

acos_grad_tiling_arch35.cpp 实现了宿主侧 tiling 计算，其分片策略的核心逻辑如下：

以ELEM_ALIGN = 512为对齐粒度，将总元素数按核数（AIV 核数）切分为blockFormer（每个 block 的元素数）与blockNum（block 数量），并将blockNum设为SetBlockDim的核维度，实现多核并行；
依据 Unity 缓冲（UB）大小（UB_SIZE_BYTES = 184 * 1024）与各 dtype 的每元素占用（FLOAT 为 32 字节/元素、FLOAT16/BF16 为 28 字节/元素）计算单次进 UB 的处理量ubFormer，进一步把每个 block 拆成整块循环ubLoop与尾块ubTail；
当totalLength == 0时走空 tensor 分支，直接清零 tiling 数据并SetBlockDim(1)；
通过useDoubleBuffer = (totalLength > 1024) ? 1 : 0决定是否启用双缓冲，并将该参数通过ASCENDC_TPL_SEL_PARAM下发给核函数模板。

这套"多核均分 + UB 内循环分块 + 尾块处理 + 双缓冲开关"的组合，保证了算子对不同规模（从 7 元素的小 shape 到数十万元素的大 shape）都能获得合理的内核资源利用率。

调用说明：图模式调用

AcosGrad 当前在仓库中提供了图模式（GEIR）调用方式，调用样例如下：

调用方式	调用样例	说明
图模式调用	examples/test_geir_acos_grad.cpp	通过算子 IR 构图方式调用 AcosGrad 算子。

以 examples/test_geir_acos_grad.cpp 为例，图模式调用的核心流程可分为四步：

构造算子节点并构图：包含acos_grad_proto.h后，通过auto acosGrad = op::AcosGrad("acos_grad")创建算子节点，使用宏ADD_INPUT(1, y, inDtype, yShape)、ADD_INPUT(2, dy, inDtype, yShape)添加y、dy两个输入（样例中 shape 取{4, 2}，dtype 取DT_FLOAT），并用ADD_OUTPUT(1, z, inDtype, yShape)声明输出z，最后graph.AddOp(...)加入计算图；
初始化 GE 运行环境：调用ge::GEInitialize(global_options)，其中全局选项包含{"ge.exec.deviceId", "0"}与{"ge.graphRunMode", "1"}，随后通过graph.SetInputs(inputs).SetOutputs(outputs)将算子挂接到图的输入输出上；
会话执行：创建ge::Session，通过session->AddGraph(graph_id, graph, graph_options)添加计算图，再以session->RunGraph(graph_id, input, output)驱动整图在 NPU 上执行；执行前还会调用aclgrphDumpGraph将构图结果 dump 到./dump目录便于调试；
结果导出：ProcessIOData将输入输出张量按二进制形式写出到./tc_ge_irrun_test_0009_npu_input_*.bin/tc_ge_irrun_test_0009_npu_output_*.bin文件，并在终端逐元素打印输出，最后调用ge::GEFinalize()收尾释放资源。

该样例同时演示了 GE 图模式下单算子验证的通用骨架，可推广到 ops-math 其他单算子的调试。

单元测试验证

仓库为 AcosGrad 提供了宿主侧单元测试，覆盖 shape 推导的正确性：

tests/ut/op_host/test_acos_grad_infershape.cpp 通过InfershapeContextPara构造算子上下文并执行 shape 推导用例，覆盖了1D fp32（16 元素）、1D 非对齐尾块（7 元素）、2D fp16（4×8）、4D bf16（2×3×4×5）、大 shape 多核场景（1×64×2×64）、超大 1D fp16（416910 元素）以及空 tensor（0 元素）等多种形态，均断言输出 shape 与输入y完全一致；
此外 tests/ut/op_host/arch35/test_acos_grad_tiling_arch35.cpp 针对 arch35 架构的 tiling 计算逻辑提供了专门的 UT 覆盖。

这些测试既验证了文档中"输出 shape 与 dtype 与 y 一致"的约束，也覆盖了空 tensor、非对齐长度等边界情况，可作为自行验证算子行为的参考样例。

小结

AcosGrad 是一个典型的逐元素（element-wise）反向梯度算子，其接口极简（y、dy两输入 +z一输出）、约束清晰（shape/dtype 一致、ND 格式、支持非连续 Tensor），却在 NPU 实现上体现了完整的分层设计：op_graph负责原型与图推导，op_host负责 shape/dtype 推导与多核分片 tiling，op_kernel负责以 float32 中间精度完成-dy / sqrt(1 - y^2)的向量化计算。对于希望在 CANN ops-math 上扩展或移植类似反向梯度算子的开发者而言，acos_grad_def.cpp、acos_grad_tiling_arch35.cpp 与 acos_grad.h 三者构成了一个值得对照研读的最小完整示例。