文档HOT
下载
中文
注册

aclnnBaddbmm/aclnnInplaceBaddbmm

接口原型

  • aclnnBaddbmm和aclnnInplaceBaddbmm实现相同的功能，其使用区别如下，请根据自身实际场景选择合适的算子。
    • aclnnBaddbmm：需新建一个输出张量对象存储计算结果。
    • aclnnInplaceBaddbmm：无需新建输出张量对象，直接在输入张量的内存中存储计算结果。
  • 每个算子分为两段接口，必须先调用“aclnnXxxGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnXxx”接口执行计算。
  • aclnnBaddbmm两段式接口如下：
    • 第一段接口：aclnnStatus aclnnBaddbmmGetWorkSpaceSize(const aclTensor *self, const aclTensor *batch1, const aclTensor *batch2, const aclScalar *beta, const aclScalar *alpha, aclTensor *out, int8_t cubeMathType, uint64_t *workspaceSize, aclOpExecutor **executor)
    • 第二段接口：aclnnStatus aclnnBaddbmm(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)
  • aclnnInplaceBaddbmm两段式接口如下：
    • 第一段接口：aclnnStatus aclnnInplaceBaddbmmGetWorkSpaceSize(const aclTensor *selfRef, const aclTensor *batch1, const aclTensor *batch2, const aclScalar *beta, const aclScalar *alpha, int8_t cubeMathType, uint64_t *workspaceSize, aclOpExecutor **executor)
    • 第二段接口：aclnnStatus aclnnInplaceBaddbmm(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述

  • 算子功能：计算α与batch1、batch2矩阵乘结果的乘积，再与β和self的乘积求和。其中batch1、batch2必须是3维张量，两个shape不支持做broadcast，而self需支持和batch1@batch2的结果做broadcast。
  • 计算公式：

  • 示例：

    假设self的shape是(A, M, K），batch1@batch2的shape是(1, M, 1)，那么输出out的shape为(A, M, K)。每一个维度的数字需要相等或其中一个为1。

    注意，若此时batch1@batch2的shape是(2, M, 2)，那么算子会因不能完成broadcast出现报错。

aclnnBaddbmmGetWorkspaceSize

  • 接口定义：

    aclnnStatus aclnnBaddbmmGetWorkSpaceSize(const aclTensor *self, const aclTensor *batch1, const aclTensor *batch2, const aclScalar *beta, const aclScalar *alpha, aclTensor *out, int8_t cubeMathType, uint64_t *workspaceSize, aclOpExecutor **executor)

  • 参数说明：
    • self：Device侧的aclTensor，数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品支持），且数据类型与batch1@batch2的数据类型需满足数据类型推导规则，shape需要与batch1@batch2满足broadcast关系。数据格式支持ND，支持非连续Tensor，支持空Tensor。
      • 对于数据类型，只能batch1@batch2的dtype向self推导，不能反过来，因此若self为FLOAT16，batch1@batch2为FLOAT，会报错。
      • 对于shape，只能是self通过broadcast变成和batch1@batch2一样的shape，不可以是batch1@batch2的shape通过broadcast变成和self一样。例如self的shape是(1, 3, 5)， batch1@batch2的shape是(1, 1, 1)，执行会报错。
    • batch1：Device侧的aclTensor，数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品支持），且数据类型与batch2、self的数据类型需满足数据类型推导规则。shape需要与batch2满足bmm输入约束条件。数据格式支持ND，支持非连续Tensor，支持空Tensor。
    • batch2：Device侧的aclTensor，数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品支持），且数据类型与batch1、self的数据类型需满足数据类型推导规则。shape需要与batch1满足bmm输入约束条件。数据格式支持ND，支持非连续Tensor，支持空Tensor。
    • beta：Host侧的aclScalar，数据类型支持FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL，且数据类型需要可转换成self与batch1@batch2推导后的数据类型。
    • alpha：Host侧的aclScalar，数据类型支持FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL，且数据类型需要可转换成self与batch1@batch2推导后的数据类型。
    • out：Device侧的aclTensor，数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品支持），且数据类型需要与self保持一致，shape要求与batch1@batch2的shape保持一致。数据格式支持ND，支持非连续Tensor，支持空Tensor。
    • cubeMathType：Host侧的整型，判断Cube单元应该使用哪种计算逻辑进行运算，支持INT8类型的枚举值，枚举值如下：
      • 0：KEPP_DTYPE，保持输入的数据类型进行计算。
      • 1：ALLOW_FP32_DOWN_PRECISION，允许转换输入数据类型降低精度计算。
    • workspaceSize：返回用户需要在Device侧申请的workspace大小。
    • executor：返回op执行器，包含了算子计算流程。
  • 返回值：

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

    第一段接口完成入参校验，出现以下场景时报错：

    • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：传入的self、batch1或batch2是空指针。
    • 返回161002（ACLNN_ERR_PARAM_INVALID）：
      • self、batch1、batch2或out的数据类型或数据格式不在支持的范围内。
      • self、batch1、batch2或out的数据类型不一致。
      • self、batch1、batch2或out的数据格式不一致。
      • self不能与batch1@batch2做broadcast操作。

aclnnBaddbmm

  • 接口定义：

    aclnnStatus aclnnBaddbmm(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

  • 参数说明：
    • workspace：在Device侧申请的workspace内存起址。
    • workspaceSize：在Device侧申请的workspace大小，由第一段接口aclnnBaddbmmGetWorkspaceSize获取。
    • executor：op执行器，包含了算子计算流程。
    • stream：指定执行任务的AscendCL stream流。
  • 返回值：

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

aclnnInplaceBaddbmmGetWorkspaceSize

  • 接口定义：

    aclnnStatus aclnnInplaceBaddbmmGetWorkSpaceSize(const aclTensor *selfRef, const aclTensor *batch1, const aclTensor *batch2, const aclScalar *beta, const aclScalar *alpha, int8_t cubeMathType, uint64_t *workspaceSize, aclOpExecutor **executor)

  • 参数说明：
    • selfRef：Device侧的aclTensor，输入/输出张量，数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品支持），且数据类型与batch1@batch2的数据类型需满足数据类型推导规则，shape需要与batch1@batch2满足broadcast关系。数据格式支持ND，支持非连续Tensor，支持空Tensor。
      • 对于数据类型，只能batch1@batch2的dtype向selfRef推导，不能反过来，因此若selfRef为FLOAT16，batch1@batch2为FLOAT，会报错。
      • 对于shape，只能是selfRef通过broadcast变成和batch1@batch2一样的shape，不可以是batch1@batch2的shape通过broadcast变成和selfRef一样。例如selfRef的shape是(1, 3, 5)， batch1@batch2的shape是(1, 1, 1)，执行会报错。
    • batch1：Device侧的aclTensor，数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品支持），且数据类型与batch2、selfRef的数据类型需满足数据类型推导规则。shape需要与batch2满足bmm输入约束条件。数据格式支持ND，支持非连续Tensor，支持空Tensor。
    • batch2：Device侧的aclTensor，数据类型支持FLOAT16、FLOAT、BFLOAT16（仅Atlas A2训练系列产品支持），且数据类型与batch1、selfRef的数据类型需满足数据类型推导规则。shape需要与batch1满足bmm输入约束条件。数据格式支持ND，支持非连续Tensor，支持空Tensor。
    • beta：Host侧的aclScalar，数据类型支持FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL，且数据类型需要可转换成selfRef与batch1@batch2推导后的数据类型。
    • alpha：Host侧的aclScalar，数据类型支持FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL，且数据类型需要可转换成selfRef与batch1@batch2推导后的数据类型。
    • cubeMathType：Host侧的整型，判断Cube单元应该使用哪种计算逻辑进行运算，支持INT8类型的枚举值，枚举值如下：
      • 0：KEPP_DTYPE，保持输入的数据类型进行计算。
      • 1：ALLOW_FP32_DOWN_PRECISION，允许转换输入数据类型降低精度计算。
    • workspaceSize：返回用户需要在Device侧申请的workspace大小。
    • executor：返回op执行器，包含了算子计算流程。
  • 返回值：

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

    第一段接口完成入参校验，出现以下场景时报错：

    • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：传入的selfRef、batch1或batch2是空指针。
    • 返回161002（ACLNN_ERR_PARAM_INVALID）：
      • selfRef、batch1、batch2或out的数据类型或数据格式不在支持的范围内。
      • selfRef、batch1、batch2或out的数据类型不一致。
      • selfRef、batch1、batch2或out的数据格式不一致。
      • selfRef不能与batch1@batch2做broadcast操作。
      • batch1和batch2的第一维度不相等。
      • batch1和batch2的维度不是三维。
      • batch1的最后一维和batch2的倒数第二维不相等。

aclnnInplaceBaddbmm

  • 接口定义：

    aclnnStatus aclnnInplaceBaddbmm(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

  • 参数说明：
    • workspace：在Device侧申请的workspace内存起址。
    • workspaceSize：在Device侧申请的workspace大小，由第一段接口aclnnInplaceBaddbmmGetWorkspaceSize获取。
    • executor：op执行器，包含了算子计算流程。
    • stream：指定执行任务的AscendCL stream流。
  • 返回值：

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

调用示例

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_baddbmm.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, aclrtContext* context, aclrtStream* stream) {
  // 固定写法，AscendCL初始化
  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 = aclrtCreateContext(context, deviceId);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
  ret = aclrtSetCurrentContext(*context);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext 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);
  // 调用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::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];
  }

  // 调用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/context/stream初始化，参考AscendCL对外接口列表
  // 根据自己的实际device填写deviceId
  int32_t deviceId = 0;
  aclrtContext context;
  aclrtStream stream;
  auto ret = Init(deviceId, &context, &stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  // 2. 构造输入与输出，需要根据API的接口自定义构造
  std::vector<int64_t> selfShape = {1, 2, 4};
  std::vector<int64_t> batch1Shape = {1, 2, 3};
  std::vector<int64_t> batch2Shape = {1, 3, 4};
  std::vector<int64_t> outShape = {1, 2, 4};
  void* selfDeviceAddr = nullptr;
  void* batch1DeviceAddr = nullptr;
  void* batch2DeviceAddr = nullptr;
  void* outDeviceAddr = nullptr;
  aclTensor* self = nullptr;
  aclTensor* batch1 = nullptr;
  aclTensor* batch2 = nullptr;
  aclScalar* alpha = nullptr;
  aclScalar* beta = nullptr;
  aclTensor* out = nullptr;
  std::vector<float> selfHostData = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<float> batch1HostData = {1, 1, 1, 2, 2, 2};
  std::vector<float> batch2HostData = {1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4};
  std::vector<float> outHostData(8, 0);
  int8_t cubeMathType = 1;
  float alphaValue = 1.2f;
  float betaValue = 1.0f;
  // 创建self aclTensor
  ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建batch1 aclTensor
  ret = CreateAclTensor(batch1HostData, batch1Shape, &batch1DeviceAddr, aclDataType::ACL_FLOAT, &batch1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建batch2 aclTensor
  ret = CreateAclTensor(batch2HostData, batch2Shape, &batch2DeviceAddr, aclDataType::ACL_FLOAT, &batch2);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // 创建alpha aclScalar
  alpha = aclCreateScalar(&alphaValue,aclDataType::ACL_FLOAT);
  CHECK_RET(alpha != nullptr, return ret);
  // 创建beta aclScalar
  beta = aclCreateScalar(&betaValue,aclDataType::ACL_FLOAT);
  CHECK_RET(beta != nullptr, return ret);
  // 创建out aclTensor
  ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 3. 调用CANN算子库API，需要修改为具体的API名称
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // 调用aclnnBaddbmm第一段接口
  ret = aclnnBaddbmmGetWorkspaceSize(self, batch1, batch2, alpha, beta, out, cubeMathType, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBaddbmmGetWorkspaceSize 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);
  }
  // 调用aclnnBaddbmm第二段接口
  ret = aclnnBaddbmm(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBaddbmm 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::vector<float> 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: %f\n", i, resultData[i]);
  }

  // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
  aclDestroyTensor(self);
  aclDestroyTensor(batch1);
  aclDestroyTensor(batch2);
  aclDestroyScalar(alpha);
  aclDestroyScalar(beta);
  aclDestroyTensor(out);
  return 0;
}
父主题： NN类算子接口