ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

CANN ops-math 算子 aclnnTan 完整指南:NPU 逐元素正切计算 API 用法与实现原理

2026/9/20 23:43:31 拓冰建站 浏览量
CANN ops-math 算子 aclnnTan 完整指南:NPU 逐元素正切计算 API 用法与实现原理 CANN ops-math 算子 aclnnTan 完整指南NPU 逐元素正切计算 API 用法与实现原理【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathaclnnTan是 CANN ops-math 数学算子库中用于计算输入张量逐元素正切值的 aclnnAscend CANN Lite NativeAPI通过tan(x) sin(x) / cos(x)的恒等式在 Atlas A2 训练/推理系列产品上实现 NPU 加速计算。本文以 experimental/math/tan/docs/aclnnTan.md 为骨架结合仓库中算子 Host 侧定义、Tiling 与 Device 侧 Kernel 源码完整讲解aclnnTanGetWorkspaceSize/aclnnTan双阶段调用的参数语义、错误码、约束边界并深入剖析其 InferShape、多核 Tiling、float16 升精度计算等底层实现。读完本文你将能够独立编写、编译、运行并验证基于 aclnnTan 的 NPU 正切计算程序。功能描述Tan 算子计算输入张量x的逐元素正切值$$out \tan(x) \frac{\sin(x)}{\cos(x)}$$其功能对标 PyTorch 的torch.tan见 experimental/math/tan/README.md 功能说明。该算子具备以下核心特征不支持广播out的 shape 与x完全相同属于单输入单输出逐元素算子支持数据类型float32、float16且out的数据类型必须与x一致支持标量与空 tensor标量输入内部按 shape{1}处理空 tensor元素数为 0时无需 workspace 直接返回成功支持动态 shape / 动态 rank算子定义中显式开启了动态 shape、动态 rank 支持详见下文源码分析。支持的产品型号产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√从源码角度进一步验证算子在 tan_def.cpp 中通过OpAICoreConfig注册了ascend910barch22 架构的 AICore 配置并开启DynamicCompileStaticFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag说明该算子面向的是 arch22 架构的昇腾 AI Core。函数原型aclnn 算子采用“查询 workspace → 执行”的双阶段异步调用模型Tan 算子对外暴露两个 APIaclnnStatus aclnnTanGetWorkspaceSize( const aclTensor *x, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor); aclnnStatus aclnnTan( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);调用约定如下先调用aclnnTanGetWorkspaceSize完成参数校验、shape 推导InferShape与执行器aclOpExecutor创建并返回算子执行所需的 workspace 大小调用方依据返回的workspaceSize自行分配设备内存再调用aclnnTan将 workspace 与执行器提交到指定 ACL stream 上异步执行通过aclrtSynchronizeStream同步等待执行完成后再释放资源。aclnnTanGetWorkspaceSize参数说明参数名输入/输出描述x输入数据类型float32、float16。数据格式ND。支持非连续 tensor。out输出数据类型与 x 相同。数据格式ND。shape 须与 x 相同。支持非连续 tensor。workspaceSize输出算子执行所需 workspace 大小单位为 Byte。由本函数返回调用方须据此分配 workspace 内存。executor输出算子执行器包含算子计算流信息由本函数返回后传入 aclnnTan 执行。从实现看Tan 算子实际不依赖 workspace在 tan_tiling.cpp 中GetWorkspaceSize将 workspace 设置为WS_SYS_SIZE 0。因此对于常规输入aclnnTanGetWorkspaceSize返回的workspaceSize为 0调用aclnnTan时传入nullptr即可仅当输入为空 tensor 等边界场景时才会走“直接返回成功”的短路路径。示例代码中的“若 workspaceSize 0 才分配”写法正是为了兼容这一约定。返回值说明返回aclnnStatus错误码详见下文错误码章节。aclnnTan参数说明参数名输入/输出描述workspace输入workspace 内存地址。若 workspaceSize 为 0可传入 nullptr。workspaceSize输入workspace 大小由 aclnnTanGetWorkspaceSize 返回。executor输入算子执行器由 aclnnTanGetWorkspaceSize 返回。stream输入ACL stream用于异步调度算子执行。返回值说明返回aclnnStatus错误码详见下文错误码章节。错误码错误码描述ACLNN_SUCCESS0执行成功。ACLNN_ERR_PARAM_NULLPTR输入/输出 tensor 指针为空。ACLNN_ERR_PARAM_INVALID参数非法包括数据类型不支持、out shape 与 x 不一致等。ACLNN_ERR_INNER_CREATE_EXECUTOR内部创建算子执行器失败。ACLNN_ERR_INNER_NULLPTR内部 tensor 分配失败。ACLNN_ERR_INNER_INFERSHAPE_ERROR内部 InferShape 失败。其中ACLNN_ERR_PARAM_INVALID与源码中的两道校验环节对应其一是算子注册层面的数据类型约束tan_def.cpp 中输入输出均声明为{ge::DT_FLOAT, ge::DT_FLOAT16}且格式限定FORMAT_ND其二是 InferShape 阶段 shape 一致性校验见下节。约束说明输入x支持 float32 和 float16 数据类型out的数据类型须与x相同out的 shape 须与x相同不支持广播支持标量输入内部转为 shape{1}处理支持空 tensor元素数为 0此时 workspaceSize 为 0直接返回成功workspace 须在调用aclnnTan之前分配在 stream 中算子执行完成后方可释放当输入值接近 $\frac{\pi}{2} k\pi$$k$ 为整数时正切函数结果趋向无穷大可能出现精度下降或溢出。最后一条约束是正切函数本身的数学特性在 $\pi/2$ 的奇数倍附近 $\cos(x) \to 0$sin/cos除法结果趋于无穷应用侧应结合数值范围评估精度需求。调用示例以下示例展示了 Tan 算子的完整调用流程与仓库 examples/test_aclnn_tan.cpp 中的官方示例同源#include cstdio #include vector #include acl/acl.h #include aclnn_tan.h int main() { // 1. 初始化 ACL 及设备 aclInit(nullptr); aclrtSetDevice(0); aclrtStream stream; aclrtCreateStream(stream); // 2. 准备输入数据fp32shape[2,4] // x [0.0, 0.5, 1.0, -1.0, 0.25, -0.5, 2.0, -2.0] // out tan(x) int64_t shape[] {2, 4}; int64_t strides[] {4, 1}; float x_host[] {0.0f, 0.5f, 1.0f, -1.0f, 0.25f, -0.5f, 2.0f, -2.0f}; float out_host[8] {0}; void *x_dev nullptr, *out_dev nullptr; size_t nbytes 8 * sizeof(float); aclrtMalloc(x_dev, nbytes, ACL_MEM_MALLOC_NORMAL_ONLY); aclrtMalloc(out_dev, nbytes, ACL_MEM_MALLOC_NORMAL_ONLY); aclrtMemcpy(x_dev, nbytes, x_host, nbytes, ACL_MEMCPY_HOST_TO_DEVICE); // 3. 创建 aclTensor aclTensor *x aclCreateTensor(shape, 2, ACL_FLOAT, strides, 0, ACL_FORMAT_ND, shape, 2, x_dev); aclTensor *out aclCreateTensor(shape, 2, ACL_FLOAT, strides, 0, ACL_FORMAT_ND, shape, 2, out_dev); // 4. 查询 workspace 大小并分配 uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; aclnnTanGetWorkspaceSize(x, out, workspaceSize, executor); void *workspace nullptr; if (workspaceSize 0) aclrtMalloc(workspace, workspaceSize, ACL_MEM_MALLOC_NORMAL_ONLY); // 5. 执行算子 aclnnTan(workspace, workspaceSize, executor, stream); aclrtSynchronizeStream(stream); // 6. 取回结果 aclrtMemcpy(out_host, nbytes, out_dev, nbytes, ACL_MEMCPY_DEVICE_TO_HOST); printf(out [%.4f, %.4f, %.4f, %.4f, %.4f, %.4f, %.4f, %.4f]\n, out_host[0], out_host[1], out_host[2], out_host[3], out_host[4], out_host[5], out_host[6], out_host[7]); // 期望: [0.0000, 0.5463, 1.5574, -1.5574, 0.2553, -0.5463, -2.1850, 2.1850] // 7. 释放资源 if (workspace) aclrtFree(workspace); aclrtFree(x_dev); aclrtFree(out_dev); aclDestroyTensor(x); aclDestroyTensor(out); aclrtDestroyStream(stream); aclrtResetDevice(0); aclFinalize(); return 0; }调用流程要点拆解上述示例可归纳为七个标准步骤这也是所有 aclnn 单算子调用的通用范式初始化环境aclInit→aclrtSetDevice→aclrtCreateStream数据准备Host 端构造输入数据aclrtMalloc分配设备内存aclrtMemcpyACL_MEMCPY_HOST_TO_DEVICE上传创建 aclTensor通过aclCreateTensor描述 shape、rank、数据类型ACL_FLOAT、strides、格式ACL_FORMAT_ND与设备地址示例中的strides {4, 1}表示行主序连续存储查询并分配 workspaceaclnnTanGetWorkspaceSize返回workspaceSize与executor按返回值条件分配 workspaceTan 常规输入下为 0可传nullptr异步执行aclnnTan提交任务后调用aclrtSynchronizeStream等待完成结果回传aclrtMemcpyACL_MEMCPY_DEVICE_TO_HOST取回结果并与预期值比对资源释放按 workspace → 设备内存 → tensor → stream → 设备 →aclFinalize的顺序逆序释放。仓库示例 test_aclnn_tan.cpp 进一步演示了工程化写法封装Init/CreateAclTensor辅助函数、通过CHECK_RET宏检查每个 ACL 调用的返回值、并以atol 1e-4、rtol 1e-4的相对/绝对误差阈值与std::tan的 golden 值逐元素比对见该文件第 135-148 行。若需要创建 float16 输入只需将数据类型改为ACL_HALF并在 Host 侧使用half类型数组即可。源码级实现原理Shape 推导输出 shape 恒等于输入 shapeTan 的 InferShape 实现位于 tan_infershape.cpp。InferShape4Tan从InferShapeContext中取出输入 shape直接赋值给输出 shapestatic ge::graphStatus InferShape4Tan(gert::InferShapeContext* context) { const gert::Shape* input_shape context-GetInputShape(0); ... *output_shape *input_shape; return ge::GRAPH_SUCCESS; }这正是“不支持广播、outshape 与x相同”这一约束在框架层的落地无论输入是标量、1D 还是 8D输出 shape 均被推导为与输入完全一致。算子注册与动态 shape 支持tan_def.cpp 通过OP_ADD(Tan)将算子注册到算子定义注册表输入x与输出y均声明DataType({ge::DT_FLOAT, ge::DT_FLOAT16})、Format({ge::FORMAT_ND, ge::FORMAT_ND})并调用.AutoContiguous()支持非连续输入aicoreConfig910B开启DynamicRankSupportFlag(true)与DynamicShapeSupportFlag(true)说明算子支持动态 rank 与动态 shapePrecisionReduceFlag(true)允许精度降低优化与 float16 升精度计算路径配合使用。Tiling多核切分 UB 切分Tan 的 Tiling 逻辑位于 tan_tiling.cpp核心任务是计算出TanTilingData三个字段定义于 tan_tiling_data.hstruct TanTilingData { int64_t totalNum 0; // 总元素数 int64_t blockFactor 0; // 每个 AI Core 处理的元素数 int64_t ubFactor 0; // 每次 UB 循环迭代处理的元素数 };多核切分blockFactor CeilDiv(totalNum, coreNum)usedCoreNum CeilDiv(totalNum, blockFactor)将总元素均匀分摊到各 AI Core当totalNum coreNum时部分 Core 因blockLength_被钳制为 0 而处于空闲状态见 tan.h。UB 切分与 buffer 规划每个 Core 内按 UB 容量分块且 float32 与 float16 使用不同的 buffer 配比数据类型buffer 数分配依据代码注释float326BUFFER_NUM_FP32输入队列 ×2 输出队列 ×2 中间结果 tmpBuf1/tmpBuf2 ×2均按 float 大小计float164BUFFER_NUM_FP16输入/输出队列按 half共 4 个 half 2 个 float 大小 两个 float 中间缓冲44合计等效 4 个 float对应代码为tiling-ubFactor FloorAlign(FloorDiv((ubCanUse / typeSize), BUFFER_NUM_XXX), ubBlockSize)其中typeSize 4因为float16 路径内部也会提升到 float32 计算见下节。TilingKey 选择根据输入 dtype 设置不同的 TilingKey——TAN_TPL_SCH_MODE_0float32或TAN_TPL_SCH_MODE_1float16定义见 tan_tiling_key.h。空 tensor 短路当totalNum 0时TilingData 三个字段全部置 0SetBlockDim(1)直接返回成功与文档“空 tensor 时 workspaceSize 为 0”的描述对应。Kernelfloat32 直算与 float16 升精度计算Device 侧 Kernel 入口位于 tan.cpp根据 TilingKey 模板分发到NsTan::Tanfloat或NsTan::Tanhalf见 tan.h。两条计算路径如下float32 路径直接计算sinVal Sin(x) // 计算 sin(x) cosVal Cos(x) // 计算 cos(x) y Div(sinVal, cosVal) // tan(x) sin(x) / cos(x)对应 tan.h 中Tanfloat::Compute特化实现使用tmpBuf1、tmpBuf2两块 VECCALC 缓冲暂存 sin、cos 中间结果。float16 路径升精度计算x_fp32 Cast(x, CAST_NONE) // half - float32 cosVal Cos(x_fp32) // 先算 cos避免输入被覆盖 sinVal Sin(x_fp32) // 再算 sin result Div(sinVal, cosVal) // tan sin / cos y Cast(result, CAST_ROUND) // float32 - half对应 tan.h 中Tanhalf::Compute特化实现。这里的执行顺序有讲究先把 half 输入 Cast 到 float32 存入sinVal缓冲随后先算 Cos 再算 Sin因为 Sin 会原地覆盖sinVal必须先保存好 cos 结果最后 Div 复用sinVal缓冲再以CAST_ROUND舍入模式 Cast 回 half。整个流程将 float16 的运算提升到 float32 精度域完成这正是 README 中“float16 全量 20 条用例 100% 通过rtol1e-3, atol1e-3”的精度保障。流水线调度Kernel 使用双缓冲BUFFER_NUM 2与TPipe队列机制通过inputQueue/outputQueue实现 CopyInGM→UB→ ComputeUB 向量计算→ CopyOutUB→GM三级流水重叠CopyIn/CopyOut使用DataCopyPad完成带 pad 的数据搬运Process主循环按ubLength_分块迭代末块用余数currentNum处理见 tan.h。构建与测试前提条件CANN Toolkit 已安装如/home/developer/Ascend/cann-9.0.0已设置环境变量source /home/developer/Ascend/ascend-toolkit/set_env.sh。编译自定义算子包cd ops/tan bash build.sh --socascend910b --pkg编译成功后算子包位于build/custom_opp_ubuntu_aarch64.run执行安装bash build/custom_opp_ubuntu_aarch64.run算子将安装到$ASCEND_HOME_PATH/opp/vendors/tan_custom/安装完成后即可在业务代码中通过#include aclnn_tan.h调用 aclnnTan。测试方法bash build.sh --socascend910b --pkg -a # 编译 UT ST bash build.sh --socascend910b --pkg -u # 编译 仅 UT bash build.sh --socascend910b --pkg -s # 编译 仅 ST全量 56 条测试用例36 条 float32 20 条 float16。ST 测试支持两种模式Mock 模式CPU Golden无需 NPU用于开发验证与真实 NPU 模式需要 NPU 设备验证实际精度。官方精度表现float32 在rtol1e-4, atol1e-6标准下全量通过float16 在rtol1e-3, atol1e-3标准下全量通过见 README.md 精度说明。总结aclnnTan 是 ops-math 中结构清晰的单输入逐元素数学算子范例对外以aclnnTanGetWorkspaceSizeaclnnTan双阶段异步接口提供服务支持 float32/float16、动态 shape、标量与空 tensor对内由 tan_def.cpp 注册算子、tan_infershape.cpp 完成 shape 推导、tan_tiling.cpp 完成多核与 UB 两级切分、tan.h 实现 float32 直算与 float16 升精度两条 Kernel 路径。理解 aclnnTan 的调用约定与实现细节不仅能直接用于 NPU 上的正切计算场景也为阅读 ops-math 中其他 aclnn 算子的 API 文档与源码提供了一套可复用的方法论。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考