ARTICLE DETAIL

资讯详情

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

CANN ops-math 中 aclnnAtan 算子两段式接口调用实战指南

CANN ops-math 中 aclnnAtan 算子两段式接口调用实战指南 CANN ops-math 中 aclnnAtan 算子两段式接口调用实战指南【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读aclnnAtan 是 CANN ops-math 数学算子库中用于对输入张量逐元素执行反正切arctangent运算的 aclnn 接口面向 Atlas A2 训练系列产品与 Atlas A2 推理系列产品提供 NPU 加速能力。本文以 experimental/math/atan/docs/aclnnAtan.md 为主干结合 experimental/math/atan 目录下的 op_api、op_host、op_kernel 源码与示例代码系统讲解 aclnnAtan 的函数原型、参数约束、错误码、完整调用示例以及算子从 aclnn 接口到 L0 算子、再到 AiCore/AiCpu Kernel 的底层实现链路。读完本文你将掌握如何在设备侧正确申请 workspace 并完成一次 Atan 算子计算也能理解该算子在 NPU 上如何通过泰勒展开与 tiling 切分实现高性能计算。产品支持情况产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√aclnnAtan 算子当前仅在 Atlas A2 训练系列产品与 Atlas A2 推理系列产品上得到支持使用前请确认运行环境属于上述产品形态。功能说明算子功能aclnnAtan 对输入矩阵张量的每个元素执行反正切运算后输出属于典型的逐元素element-wise数学算子公式如下$$ out_{i}tan^{-1}(input_{i}) $$即输出张量out的每个元素为输入张量input对应元素x的反正切值arctan(x)。在 experimental/math/atan/README.md 中同样给出了该算子的逐元素定义并补充了其数值计算所依据的级数展开式$$ y_i \sum_{n0}^{\infty} \frac{(-1)^{n} x_i^{2n1}}{2n1} $$该级数即反正切函数的泰勒展开形式仓库中 Kernel 层的实际实现正是围绕这一数学基础进行的分段逼近优化详见后文Kernel 层算法实现。输入输出参数参数名输入/输出/属性描述数据类型数据格式x输入公式中的输入张量 xFLOAT、FLOAT16、BFLOAT16NDout输出公式中的输出张量 yFLOAT、FLOAT16、BFLOAT16ND函数原型与两段式调用每个 aclnn 算子都采用两段式接口设计必须先调用第一段接口aclnnAtanGetWorkspaceSize完成入参校验并获取执行所需的 workspace 大小再调用第二段接口aclnnAtan在指定 Stream 上执行计算。aclnnStatus aclnnAtanGetWorkspaceSize(const aclTensor *input, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnAtan(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)两段式设计的核心价值在于第一段接口在 Host 侧完成参数校验、算子图构建与 workspace 大小推算不实际触发计算第二段接口基于第一段返回的 executor 与 workspace 内存真正下发计算任务。这样可以让调用方根据实际需要精确申请 Device 侧内存同时复用已构建好的执行器。aclnnAtanGetWorkspaceSize 参数说明第一段接口用于校验入参并计算所需 workspace 大小参数如下参数类型方向说明inputaclTensor*计算输入输入 TensorDevice 侧的 aclTensor。当类型为 FLOAT16、BFLOAT16 时在内部会转化为 FLOAT32 进行运算输出 FLOAT32 类型Atlas A2 产品上数据类型支持 FLOAT、FLOAT16、BFLOAT16。支持非连续的 Tensor数据格式支持 ND非连续 Tensor 的维度不大于 8且 shape 需要与 out 一致outaclTensor*计算输出输出 TensorDevice 侧的 aclTensor。支持非连续的 Tensor数据格式支持 ND非连续 Tensor 的维度不大于 8且 shape 需要与 input 一致Atlas A2 产品上数据类型支持 FLOAT、FLOAT16、BFLOAT16workspaceSizeuint64_t*出参返回需要在 Device 侧申请的 workspace 大小供第二段接口使用executoraclOpExecutor**出参返回 op 执行器封装了算子计算流程返回值与入参校验错误码第一段接口返回aclnnStatus状态码具体定义参见 aclnn返回码。第一段接口完成入参校验出现以下场景时报错161001 (ACLNN_ERR_PARAM_NULLPTR)1. 传入的input或out是空指针。 161002 (ACLNN_ERR_PARAM_INVALID)1. input或out的数据类型不在支持的范围之内。 2. input和out的shape不一致。 3. input或out的维数大于8。从源码 experimental/math/atan/op_api/aclnn_atan.cpp 可以看到这些校验的具体实现CheckParamsAtan通过CheckDtypeValid1In1Out校验输入输出数据类型是否落在支持列表内通过CheckSameShape1In1Out校验 input 与 out 的 shape 是否一致任一项不满足均返回ACLNN_ERR_PARAM_INVALID而空指针检查由CHECK_NOT_NULL(input, out)完成对应ACLNN_ERR_PARAM_NULLPTR。此外当输入为空张量input-IsEmpty()时接口直接返回成功并将workspaceSize置为 0避免无效计算。aclnnAtan 参数说明第二段接口在指定 Stream 上真正执行计算参数如下参数类型方向说明workspacevoid*入参在 Device 侧申请的 workspace 内存地址由调用方根据第一段接口计算出的 workspaceSize 申请workspaceSizeuint64_t入参在 Device 侧申请的 workspace 大小由第一段接口 aclnnAtanGetWorkspaceSize 获取executoraclOpExecutor*入参op 执行器包含了算子计算流程由第一段接口返回streamaclrtStream入参指定执行任务的 Stream第二段接口同样返回aclnnStatus状态码参见 aclnn返回码。在源码实现中第二段接口通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)调用框架统一能力完成计算其中L2_DFX_PHASE_2用于打点记录算子执行信息便于后续定位与性能分析。调用示例下面给出完整的调用示例可在 Atlas A2 环境上编译运行。示例中构造了 shape 为{4, 2}的 FLOAT 类型输入经过aclnnAtanGetWorkspaceSize→ 申请 workspace →aclnnAtan→ 同步 Stream → 拷贝结果回 Host 的完整流程最后打印每个元素的反正切结果。具体编译与执行过程请参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_atan.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // check根据自己的需要处理 CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t selfShape {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f,0.8f}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. aclnnAtan接口调用示例 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnAtan第一段接口 ret aclnnAtanGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAtanGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 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); } // 调用aclnnAtan第二段接口 ret aclnnAtan(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAtan failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧需要根据具体API的接口定义修改 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %lf\n, i, resultData[i]); } // 6. 释放aclTensor需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例代码与 experimental/math/atan/examples/test_aclnn_atan.cpp 保持一致的调用骨架可在构建算子库后参考运行。对照样例可知完整调用流程可归纳为七个步骤初始化aclInit、aclrtSetDevice、aclrtCreateStream固定写法构造输入输出在 Host 侧准备数据通过aclrtMalloc申请 Device 内存、aclrtMemcpy拷贝数据再用aclCreateTensor创建 aclTensorND 格式需按 shape 计算连续 strides两段式调用先调用aclnnAtanGetWorkspaceSize获取 workspaceSize 与 executor再按需aclrtMalloc申请 workspace最后调用aclnnAtan执行同步等待aclrtSynchronizeStream等待任务执行结束结果回拷aclrtMemcpy将 Device 侧结果拷贝回 Host 并打印释放 aclTensoraclDestroyTensor释放资源aclrtFree释放 Device 内存与 workspaceaclrtDestroyStream、aclrtResetDevice、aclFinalize。其中第 3 步的两段式调用是本算子使用上的关键约束workspace 内存地址与大小必须分别由第一段接口的出参提供不能自行估算或省略。源码级实现原理aclnn 接口层校验、Cast 与 ViewCopy 组装在 experimental/math/atan/op_api/aclnn_atan.cpp 中ExecAtanGetWorkspaceSize展示了第一段接口的完整计算图组装逻辑创建 OpExecutor进行空指针与参数校验对输入执行l0op::Contiguous将非连续输入统一为连续 Tensor判断输入数据类型是否在输出支持列表FLOAT、FLOAT16、DOUBLE、BF16内不在则用l0op::Cast提升为 FLOAT32 再参与运算——这正对应文档中FLOAT16、BFLOAT16 转化为 FLOAT32 进行运算的说明调用l0op::Atan完成核心反正切计算用l0op::Cast将中间结果转换回 out 声明的数据类型通过l0op::ViewCopy将结果写入 outout 可能是非连续 Tensor通过uniqueExecutor-GetWorkspaceSize()汇总整个计算链所需的 workspace 大小并返回。此外该文件还提供了 in-place 版本的aclnnInplaceAtanGetWorkspaceSize/aclnnInplaceAtan其中输入输出共用同一个 TensorinputRef并通过CheckInplaceDtypeValid限定 in-place 场景下支持的数据类型FLOAT、FLOAT16、BF16可用于需要原地计算的场景。L0 算子层AiCore / AiCpu 动态分发experimental/math/atan/op_api/atan.cpp 实现了 L0 算子l0op::Atan的分发逻辑根据当前芯片类型与输入数据类型通过IsAiCoreSupport判断是否走 AiCore 路径——在 ASCEND910B、ASCEND910_93 以及寄存器级RegBase平台上 FLOAT、FLOAT16、BF16 均可走 AiCore其余场景仅 FLOAT、FLOAT16 走 AiCore不满足条件时降级为 AiCpu 路径AtanAiCpu通过ADD_TO_LAUNCHER_LIST_AICPU加入任务队列。这样保证了同一 aclnn 接口在不同硬件形态下的可用性。算子定义与形状推导experimental/math/atan/op_host/atan_def.cpp 中通过OpDef注册了 Atan 算子输入x与输出y均声明为 REQUIRED 参数支持数据类型为DT_FLOAT、DT_FLOAT16、DT_BF16数据格式为 ND并配置了 AiCore 的DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)等动态特性。形状推导逻辑在 experimental/math/atan/op_host/atan_infershape.cpp 中InferShapeAtan直接将输出形状赋值为输入形状*yShape *xShape体现了逐元素算子输出与输入同 shape 的特性。Tiling 计算数据切分与核间负载均衡experimental/math/atan/op_host/atan_tiling.cpp 负责将输入数据按多核并行方式切分。其核心思路是从平台信息中获取 AiCore 核数coreNum与 workspace 大小从 UBUnified Buffer内存大小出发按 256 字节对齐划分空间块根据数据类型确定单次处理的数据块数dataNumberFLOAT 为 10FLOAT16/BF16 为 19推算每次 tile 处理的数据量tileDataNum将输入总长度按 256 字节对齐后均分到各个核余数部分由大核多承担一个块bigCoreDataNum与smallCoreDataNum之分并计算每核的 tile 循环次数finalBigTileNum/finalSmallTileNum与尾块处理量tailDataNum最终将结果写入AtanTilingData结构体见 experimental/math/atan/op_kernel/atan_tiling_data.h并通过context-SetBlockDim(coreNum)设置核数、SetTilingKey设置编译模板键。该结构体包含smallCoreDataNum、bigCoreDataNum、finalBigTileNum、finalSmallTileNum、smallTailDataNum、bigTailDataNum、tileDataNum、tailBlockNum八个字段Kernel 侧据此定位各自的数据区间并确定循环次数。Kernel 层算法分段逼近的反正切实现Kernel 实现位于 experimental/math/atan/op_kernel/atan.cpp 与 experimental/math/atan/op_kernel/atan.h。NsAtan::AtanT是模板类通过Init/Process/CopyIn/Compute/CopyOut组成典型的数据流水Init根据 tiling 数据确定当前核的全局数据起始位置、tile 数与尾块大小并初始化输入输出队列及多个 VECCALC 临时缓冲区CopyIn将 Global Memory 数据按 tile 拷贝进 UB 队列CopyOut将计算结果拷回 Global Memory两者通过TQue双缓冲BUFFER_NUM 2实现搬运与计算重叠Compute中实现核心算法先将输入夹在[-10000, 10000]内以避免后续泰勒计算溢出当epsilon0.0001时tan(pi/2 - 0.0001) ≈ 10000注释中明确说明了该边界选取依据随后提取符号位sign x / |x|对|x|利用恒等式将大数值区间折叠到小角度区间——计算(x-1)/(x1)构造第二路待逼近值再调用两次DoTaylor分别逼近两个区间的反正切值最后加π/4并取两路结果的最小值、乘回符号位得到最终结果DoTaylor使用常数TAN_PI_BY_EIGHTtan(π/8)、CONST_PI_BY_FOURπ/4、CONST_PI_BY_EIGHTπ/8与系数表TAYLOR[] {1, -1/3, 1/5, -1/7, 1/9, -1/11, 1/13}通过霍纳法Horner 形式逐层Mul/Adds计算反正切的泰勒多项式最终叠加π/8完成区间校正。对于 half 与 bfloat16_t 输入Compute 会先将数据Cast到 FLOAT32 运算对应文档中转化为 FLOAT32 进行运算的描述计算完成后再转回原类型输出以保证中间计算精度。tiling 模板键由 experimental/math/atan/op_kernel/atan_tiling_key.h 中的ASCENDC_TPL_ARGS_DECL声明预留了schMode模板参数以支持不同调度模式。测试与验证仓库在 experimental/math/atan/tests 下提供了三层测试op_api 层tests/ut/op_api/test_aclnn_atan.cpp 直接调用 aclnnAtan 接口验证端到端功能op_host 层tests/ut/op_host/test_atan_tiling.cpp 校验 tiling 计算结果的正确性op_kernel 层tests/ut/op_kernel/test_atan.cpp 配合 gen_data.py 与 compare_data.py 生成输入数据并对比 Kernel 输出与标准结果。如需在本地复现可参照 docs/zh/context/compile_and_run_sample.md 与 README.md 中的构建说明完成算子库编译与单测执行。小结aclnnAtan 是 CANN ops-math 提供的标准逐元素反正切算子接口其使用要点可归纳为严格遵循两段式调用先 GetWorkspaceSize 后执行、保证 input 与 out shape 一致且维度不大于 8、FLOAT16/BFLOAT16 输入在内部以 FLOAT32 精度计算。从实现层面看该算子依次经过 aclnn 接口层的参数校验与 Cast/ViewCopy 组装、L0 算子的 AiCore/AiCpu 动态分发、Host 侧 tiling 的多核切分最终在 Kernel 层通过双缓冲流水与基于泰勒级数的分段逼近算法完成 NPU 上的高性能反正切计算。上述调用骨架与实现原理同样适用于 ops-math 中其他逐元素数学算子可作为在 Atlas A2 产品上开发与优化数学算子的参考范本。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表