0
代码
代码
Issues
Pull Requests
讨论
Wiki
分析
0
cann-robotcann-robot【社区任务】AscendC实现Atan2算子贡献
4a7bc4cc创建于 22 天前历史提交

aclnnAtan2

产品支持情况

产品 是否支持
Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件 √

功能说明

  • 接口功能：计算两个输入张量 x1（分子）和 x2（分母）的逐元素反正切值，结果范围 (−π, π]，可正确处理所有象限（包括 x=0 的边界情况）。

  • 计算公式：

outi=atan2(x1i, x2i)\text{out}_i = \text{atan2}(x1_i,\ x2_i)

函数原型

每个算子分为两段式接口，必须先调用"aclnnAtan2GetWorkspaceSize"接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器，再调用"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)

aclnnAtan2GetWorkspaceSize

  • 参数说明：

    参数名 输入/输出 描述 使用说明 数据类型 数据格式 维度(shape) 非连续Tensor
    x1 输入 分子张量，公式中的 y 分量。
    • 支持空Tensor。
    • x1 与 x2 的 shape 必须一致。
    • x1 与 x2 的数据类型必须一致。
    		 FLOAT16、FLOAT、BFLOAT16 ND 0-8 √
    x2 输入 分母张量，公式中的 x 分量。
    • 支持空Tensor。
    • x1 与 x2 的 shape 必须一致。
    • x1 与 x2 的数据类型必须一致。
    		 FLOAT16、FLOAT、BFLOAT16 ND 0-8 √
    y 输出 输出张量，逐元素 atan2 结果，值域 (−π, π]。
    • 输出 shape 与 x1 一致。
    • 输出数据类型与 x1 一致。
    		 FLOAT16、FLOAT、BFLOAT16 ND 0-8 √
    workspaceSize 输出 返回需要在 Device 侧申请的 workspace 大小。 - - - - -
    executor 输出 返回 op 执行器，包含了算子计算流程。 - - - - -

  • 返回值：

    aclnnStatus：返回状态码，具体参见aclnn返回码。 第一段接口会完成入参校验，出现以下场景时报错：

    返回码 错误码 描述
    ACLNN_ERR_PARAM_NULLPTR 161001 传入的 x1、x2 或 y 是空指针。
    ACLNN_ERR_PARAM_INVALID 161002 x1 或 x2 的数据类型不在支持的范围之内（仅支持 FLOAT16、FLOAT、BFLOAT16）。
    x1 与 x2 的数据类型不同。
    x1 与 x2 的 shape 不同。

aclnnAtan2

  • 参数说明：

    参数名 输入/输出 描述
    workspace 输入 在 Device 侧申请的 workspace 内存地址。
    workspaceSize 输入 在 Device 侧申请的 workspace 大小，由第一段接口 aclnnAtan2GetWorkspaceSize 获取。
    executor 输入 op 执行器，包含了算子计算流程。
    stream 输入 指定执行任务的 Stream。

  • 返回值：

    aclnnStatus：返回状态码，具体参见aclnn返回码。

约束说明

无。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。

#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::vector<int64_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::vector<T>& hostData, const std::vector<int64_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::vector<int64_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::vector<int64_t> shape = {4, 4};
  // x1: y 分量；x2: x 分量
  std::vector<float> x1HostData = {-3, -2, -1, 0, 1, 2, 3, -3, -2, -1, 0, 1, 2, 3, -3, -2};
  std::vector<float> x2HostData = { 0,  1,  2, 3,-3,-2,-1,  0,  1,  2, 3,-3,-2,-1,  0,  1};
  std::vector<float> 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::vector<float> 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;
}