CANN ops-math aclnnAtan2 算子接口解析两段式调用流程与逐元素反正切计算实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 仓库中 experimental/math/atan2/docs/aclnnAtan2.md 为核心系统讲解aclnnAtan2算子的功能定义、两段式接口原型、参数约束与完整调用示例并结合 math/atan2 目录下的源码实现深入剖析其内部的计算链路与类型推导逻辑。读完本文你将掌握在 NPU 上通过 CANN aclnn 接口正确完成atan2逐元素运算的完整流程并能独立编译、运行与验证该算子。功能说明逐元素反正切 atan2aclnnAtan2用于计算两个输入张量x1分子即 y 分量与x2分母即 x 分量的逐元素反正切值计算公式为$$ \text{out}_i \text{atan2}(x1_i,\ x2_i) $$与单参数atan不同atan2(y, x)通过同时接收 y 与 x 两个分量可以正确处理所有象限包括x 0的边界情况。其结果值域为(−π, π]实际计算中相当于以坐标(x2, x1)为参数求极角angle因此在工程上常用于极坐标变换、向量方向角计算、信号相位求解等场景。从公式语义上看x1对应数学上的 y 分量分子x2对应 x 分量分母这一点在接口参数说明中有明确标注使用时切勿颠倒否则结果会偏移。产品支持情况根据原文档aclnnAtan2支持的产品如下产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品 / A200I A2 Box 异构组件√仓库 math/atan2/docs/aclnnAtan2aclnnInplaceAtan2.md 中给出的产品支持矩阵更广覆盖 Atlas A2/A3 训练与推理系列、Atlas 训练系列产品、Ascend 950PR/Ascend 950DT 等同时也明确了 Atlas 200I/500 A2 推理产品不支持。具体以你所使用环境实际安装的 CANN 版本对应产品为准。两段式接口与函数原型aclnnAtan2遵循 CANN 算子库的两段式接口设计必须先调用第一段接口aclnnAtan2GetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器executor再调用第二段接口aclnnAtan2真正执行计算。aclnnStatus aclnnAtan2GetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, aclTensor *y, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnAtan2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)两个阶段的职责划分清晰第一段GetWorkspaceSize完成入参校验、构建算子执行图并据此推算出运行所需 workspace 大小同时将执行器句柄返回给调用方。该阶段不执行实际计算可视为规划阶段。第二段aclnnAtan2接收第一段返回的 executor、workspace 与用户指定的 Stream将计算任务提交到 NPU 上异步执行。aclnnAtan2GetWorkspaceSize 参数详解第一段接口参数说明如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1输入分子张量公式中的 y 分量支持空 Tensorx1 与 x2 的 shape 必须一致x1 与 x2 的数据类型必须一致FLOAT16、FLOAT、BFLOAT16ND0-8√x2输入分母张量公式中的 x 分量支持空 Tensorx1 与 x2 的 shape 必须一致x1 与 x2 的数据类型必须一致FLOAT16、FLOAT、BFLOAT16ND0-8√y输出输出张量逐元素 atan2 结果值域 (−π, π]输出 shape 与 x1 一致输出数据类型与 x1 一致FLOAT16、FLOAT、BFLOAT16ND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----值得注意的几个约束点shape 一致性本接口要求x1与x2的 shape 严格一致输出y的 shape 与x1一致不支持广播。这与仓库稳定版 math/atan2 目录中支持 broadcast 的变体不同使用前请以当前文档experimental 版本为准。数据类型一致性x1、x2、y三者数据类型必须一致仅支持 FLOAT16、FLOAT、BFLOAT16。非连续 Tensor三个张量均支持非连续内存布局对应 ACL_FORMAT_ND 格式下的 strided 视图框架会在内部通过 ViewCopy 等手段处理见下文源码解析。空 Tensorx1、x2支持空 Tensor此时第一段接口直接返回workspaceSize 0无需申请 workspace。返回值与错误码第一段接口会完成入参校验出现异常时返回对应的aclnnStatus错误码完整错误码定义参见 aclnn返回码返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 x1、x2 或 y 是空指针ACLNN_ERR_PARAM_INVALID161002x1 或 x2 的数据类型不在支持的范围之内仅支持 FLOAT16、FLOAT、BFLOAT16ACLNN_ERR_PARAM_INVALID161002x1 与 x2 的数据类型不同ACLNN_ERR_PARAM_INVALID161002x1 与 x2 的 shape 不同在 math/atan2/op_api/aclnn_atan2.cpp 中可以看到第一段接口的校验顺序与文档描述的报错场景一一对应CheckNotNull(self, other, out)任一参数为空指针即返回ACLNN_ERR_PARAM_NULLPTRCheckDtypeValid(self, other)通过OP_CHECK_DTYPE_NOT_SUPPORT检查输入类型是否在支持列表内CheckShape(self, other, out)校验维度不超过 8并通过OP_CHECK_BROADCAST_AND_INFER_SHAPE与OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE完成 shape 推导与一致性比对。值得注意的是上述源码还处理了空 Tensor 的快速路径if (self-IsEmpty() || other-IsEmpty())时直接置*workspaceSize 0并返回成功这与文档中支持空 Tensor的说明一致。aclnnAtan2 参数详解第二段接口参数说明如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnAtan2GetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream其中workspace需要调用方根据第一段接口返回的workspaceSize通过aclrtMalloc在 Device 侧申请若workspaceSize为 0 则无需申请。stream决定算子在哪个异步流上排队执行调用后一般需通过aclrtSynchronizeStream同步等待任务完成再读取结果。源码级实现原理第一段接口内部的计算图构建从 math/atan2/op_api/aclnn_atan2.cpp 的ExecAtan2GetWorkspaceSize可以看到算子执行流由多个底层 l0op 原子算子组合而成Self ──► l0op::Contiguous ──► l0op::Cast ──► l0op::Atan2 ──► l0op::Cast ──► l0op::ViewCopy ──► Out Other ─► l0op::Contiguous ──► l0op::Cast ─┘各环节作用如下Contiguous将输入self/other规整为连续内存视图这是对非连续 Tensor 支持的底层实现保障Cast输入侧通过InferDtype对两个输入做类型提升PromoteType若提升后的类型不在算子内核支持列表中则回落为 FLOAT。从源码看算子内核含 aicore/aicpu实际支持的类型为 FLOAT、FLOAT16、DOUBLE910B 及以上追加 BF16而 API 输入层额外允许 INT8/INT16/INT32/INT64/UINT8/BOOL 等类型由 Cast 统一转换为内核可计算的浮点类型l0op::Atan2 内核执行真正的逐元素反正切计算Cast输出侧将内核结果转换回out声明时的数据类型ViewCopy将计算结果写入可能非连续的输出张量out。最后通过uniqueExecutor-GetWorkspaceSize()汇总整条链路的 workspace 需求并返回给调用方。内核分发AICORE 与 AICPU在 math/atan2/op_api/atan2.cpp 中l0op::Atan2会根据输入数据类型决定走哪条内核路径AICORE 路径当 dtype 属于 FLOAT、FLOAT16、BF16AICORE_DTYPE_SUPPORT_LIST时走ADD_TO_LAUNCHER_LIST_AICORE使用 AI Core 加速计算AICPU 路径其他类型回退到 AICPU 内核ADD_TO_LAUNCHER_LIST_AICPU。对应地算子定义文件 math/atan2/op_host/atan2_def.cpp 注册了输入x1/x2与输出y数据类型为 BF16、FLOAT16、FLOAT并针对 ascend950 配置了支持动态 rank/shape 的 AICore 配置项InferShape 实现 math/atan2/op_host/atan2_infershape.cpp 校验两输入 dtype 相同输出 shape 取二者广播结果、输出 dtype 继承自x1。AICORE 内核模板AICORE 内核实现在 math/atan2/op_kernel/atan2_apt.cpp采用模板化 kernel 入口atan2_opschMode, T内部基于Ops::Base的BroadcastSch调度框架配合 arch35/atan2_dag.h 中定义的算子 DAG 完成对x1、x2的广播与逐元素计算最终写入y。从源码结构可以看出该实现复用了社区统一的 broadcast 调度方案保证了多形态输入下的计算正确性与性能。约束说明本接口experimental 版本约束如下无额外约束原文档注明无。需在实际使用中自行把握的约束均来自参数表即shape 必须一致、数据类型必须一致且限定在 FLOAT16/FLOAT/BFLOAT16、维度不超过 8。仓库稳定版变体math/atan2还声明了确定性计算特性aclnnAtan2默认确定性实现即相同输入下多次运行结果一致可参考 确定性计算。调用示例以下完整示例来自 experimental/math/atan2/docs/aclnnAtan2.md展示了从 ACL 初始化、构造 Tensor、两段式调用到结果回拷的完整流程。编译与运行的具体步骤请参考编译与运行样例。#include iostream #include vector #include cmath #include acl/acl.h #include aclnnop/aclnn_atan2.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); 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); 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); 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]; } *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 初始化 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t shape {4, 4}; // x1: y 分量x2: x 分量 std::vectorfloat x1HostData {-3, -2, -1, 0, 1, 2, 3, -3, -2, -1, 0, 1, 2, 3, -3, -2}; std::vectorfloat x2HostData { 0, 1, 2, 3,-3,-2,-1, 0, 1, 2, 3,-3,-2,-1, 0, 1}; std::vectorfloat yHostData(16, 0.0f); void* x1DeviceAddr nullptr; void* x2DeviceAddr nullptr; void* yDeviceAddr nullptr; aclTensor* x1 nullptr; aclTensor* x2 nullptr; aclTensor* y nullptr; ret CreateAclTensor(x1HostData, shape, x1DeviceAddr, aclDataType::ACL_FLOAT, x1); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(x2HostData, shape, x2DeviceAddr, aclDataType::ACL_FLOAT, x2); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(yHostData, shape, yDeviceAddr, aclDataType::ACL_FLOAT, y); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用 aclnnAtan2 两段式接口 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnAtan2GetWorkspaceSize(x1, x2, y, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAtan2GetWorkspaceSize 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 aclnnAtan2(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAtan2 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. 拷贝结果并打印 auto size GetShapeSize(shape); std::vectorfloat outData(size, 0); ret aclrtMemcpy(outData.data(), outData.size() * sizeof(float), yDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] %f (ref: %f)\n, i, outData[i], std::atan2(x1HostData[i], x2HostData[i])); } // 6. 释放资源 aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(y); aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(yDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例代码的关键要点入参与出参的内存生命周期x1、x2、y的 Device 内存通过aclrtMalloc申请、数据通过aclrtMemcpy搬运aclTensor视图通过aclCreateTensor创建、aclDestroyTensor释放示例中采用ACL_FORMAT_ND与连续 strides实际工程中可构造非连续视图以验证框架的 ViewCopy 支持。结果校验示例在打印时使用std::atan2(x1HostData[i], x2HostData[i])作为参考值做逐元素对比方便快速验证 NPU 计算结果正确性。workspace 按需申请仅当workspaceSize 0时才申请并最终释放 workspace 内存这是所有 aclnn 两段式接口的通用写法。测试与验证仓库为aclnnAtan2提供了完整的测试覆盖可作为验证与二次开发的参考单元测试 math/atan2/tests/ut/op_api/test_aclnn_atan2.cpp通过OP_API_UT宏批量覆盖多种 shape 与数据类型组合并包含空指针入参的异常用例算子 host 侧测试 math/atan2/tests/ut/op_host/test_atan2_infershape.cpp验证 InferShape 广播推导与 dtype 继承逻辑可独立运行的完整样例 math/atan2/examples/test_aclnn_atan2.cpp以 RAII 方式管理 Stream、Device 内存与 Tensor 生命周期演示了工程化的调用写法。从测试与实现可以确认aclnnAtan2在 ops-math 中属于API 接口层 算子定义层 AICORE/AICPU 内核层三层结构完备的数学算子本文聚焦的 experimental 版本是其在特定产品Atlas A2 系列上的接口文档要求输入输出 shape 一致、dtype 一致若需要广播语义与更宽的数据类型支持可参考稳定版 math/atan2 目录及其文档。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考