aclnnEmbeddingDenseBackward

Applicable Products

Function

Operator function: Implements the backward computation of aclnnEmbedding, accumulating the row of grad corresponding to the same index indices onto out.

Function Prototype

Each operator is divided into a two-phase API. First call the "aclnnEmbeddingDenseBackwardGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the "aclnnEmbeddingDenseBackward" API to perform the computation.

• aclnnStatus aclnnEmbeddingDenseBackwardGetWorkspaceSize(const aclTensor *grad, const aclTensor *indices, uint64_t numWeights, uint64_t paddingIdx, bool scaleGradByFreq, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)

• aclnnStatus aclnnEmbeddingDenseBackward(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)

aclnnEmbeddingDenseBackwardGetWorkspaceSize

• Parameters:

  • grad(aclTensor*, computation input): Original gradient of the data, an aclTensor on the Device side. Supports 2 to 8 dimensions. After combining all axes except the last axis, the shape is the same as the combined shape of indices. Supports non-contiguous tensor. The data format supports ND.

    • Atlas training products: The data type supports FLOAT16 and FLOAT.

    • Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products, Ascend 950PR/Ascend 950DT: The data type supports BFLOAT16, FLOAT16, and FLOAT.

  • indices(aclTensor*, computation input): Index values corresponding to the grad input, an aclTensor on the Device side. The value range is [0, numWeights). Supports 1 to 8 dimensions. Supports non-contiguous tensor. The data format supports ND. The data type supports FLOAT, FLOAT16, DOUBLE, INT32, INT64, INT16, INT8, UINT8, and BOOL.

  • numWeights(uint64_t, computation input): Specifies the size of the first axis of the output tensor.

  • paddingIdx(uint64_t, computation input): Fills the paddingIdx-th row of the output tensor with zeros. If paddingIdx is negative, no processing is performed.

  • scaleGradByFreq(bool, computation input): Whether to scale the gradient based on word frequency. If set to true, scales the result by word frequency; if set to false, no processing is performed.

  • out(aclTensor*, computation output): Output of the gradient summation result, an aclTensor on the device side, with two dimensions. The size of the first axis is numWeights, and the size of the last axis is the same as that of grad. The data type is the same as that of grad. The data format supports only ND.

    • Atlas training products: The data type supports FLOAT16 and FLOAT.

    • Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products, Ascend 950PR/Ascend 950DT: The data type supports BFLOAT16, FLOAT16, and FLOAT.

  • workspaceSize(uint64_t *, output parameter): Returns the workspace size to be allocated on the Device side.

  • executor(aclOpExecutor **, output parameter): Returns the op executor, which contains the operator computation flow.

• Return value

  aclnnStatus: return code. For details, see aclnn Return Code.

  The first-phase API completes input validation. The following scenarios trigger an error:
  Returns 161001 (ACLNN_ERR_PARAM_NULLPTR): 1. The passed grad, indices, or out is a null pointer.
  Returns 161002 (ACLNN_ERR_PARAM_INVALID): 1. The data type or data format of grad, indices, or out is not within the supported range.
                                      2. The dimensions of grad or indices exceed 8.
                                      3. The shape of grad and indices does not meet the constraints.
                                      4. The shape of out does not match the deduced result.

aclnnEmbeddingDenseBackward

• Parameters:

  • workspace(void *, input parameter): Address of the workspace memory applied on the Device side.

  • workspaceSize(uint64_t, input parameter): Size of the workspace applied on the Device side, obtained by the first-phase API aclnnEmbeddingDenseBackwardGetWorkspaceSize.

  • executor(aclOpExecutor *, input parameter): Op executor, which contains the operator computation flow.

  • stream(aclrtStream, input parameter): Specifies the stream for executing the task.

• Return value

  aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

• Atlas training products:

  • For scenarios where scale is true, if the last dimension of grad is embeddingDim and its size exceeds the specified range, an error is reported. The valid range is as follows:

    • When indices is int32, the following must be satisfied:
    $$embeddingDim < \frac{180192 - countsSize * 4}{36}$$
    • When indices is int64, the following must be satisfied:
    $$embeddingDim < \frac{180192 - countsSize * 8}{20}$$
    • The formula for countsSize is as follows, where coreNum represents the number of AI processor cores:
    $$countsSize = numWeights / coreNum + numWeights \% coreNum$$

• Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:

  • When the parameter shape exceeds the following limits, the output cannot guarantee high precision. If deterministic computation is enabled, high performance also cannot be guaranteed.

    • After grad is coalesced into a two-dimensional shape, the first dimension exceeds INT32_MAX (2147483647).

    • numWeights exceeds INT32_MAX (2147483647).

• When the dimension of indices after concatenation exceeds INT32_INF (2139095040), high performance cannot be guaranteed.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_embedding_dense_backward.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) {
  // Boilerplate: resource initialization.
  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);
  // Call aclrtMalloc to allocate device-side memory.
  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);
  // Call aclrtMemcpy to copy host-side data to the device-side memory.
  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);

  // Compute the strides of a contiguous tensor.
  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];
  }

  // Call the aclCreateTensor API to create an 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. (Boilerplate) Initialize the device/stream. See the acl API manual.
  // Fill in the deviceId based on your actual device.
  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. Construct the input and output. Customize the construction based on the API.
  uint64_t numWeights = 4;
  uint64_t paddingIdx = 0;
  bool scaleGradByFreq = false;
  std::vector<int64_t> gradOutputShape = {2, 3};
  std::vector<int64_t> indicesShape = {2};
  std::vector<int64_t> outShape = {4, 3};
  void* gradOutputDeviceAddr = nullptr;
  void* indicesDeviceAddr = nullptr;
  void* outDeviceAddr = nullptr;
  aclTensor* gradOutput = nullptr;
  aclTensor* indices = nullptr;
  aclTensor* out = nullptr;

  std::vector<float> gradOutputHostData = {1, 2, 3, 4, 5, 6};
  std::vector<int64_t> indicesHostData = {1, 2};
  std::vector<float> outHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

  // Create the gradOutput aclTensor.
  ret = CreateAclTensor(gradOutputHostData, gradOutputShape, &gradOutputDeviceAddr, aclDataType::ACL_FLOAT, &gradOutput);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the indices aclTensor.
  ret = CreateAclTensor(indicesHostData, indicesShape, &indicesDeviceAddr, aclDataType::ACL_INT64, &indices);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the out aclTensor.
  ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 3. Call the CANN operator library API. Replace it with the specific API name.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // Call the first-phase API of aclnnEmbeddingDenseBackward.
  ret = aclnnEmbeddingDenseBackwardGetWorkspaceSize(gradOutput, indices, numWeights, paddingIdx, scaleGradByFreq, out,
                                                    &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnEmbeddingDenseBackwardGetWorkspaceSize failed. ERROR: %d\n", ret);
            return ret);
  // Allocate device memory based on the workspaceSize calculated by the first-phase API.
  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);
  }
  // Call the second-phase API of aclnnEmbeddingDenseBackward.
  ret = aclnnEmbeddingDenseBackward(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnEmbeddingDenseBackward failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Wait for the task execution to complete synchronously.
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // 5. Obtain the output value and copy the result from the device memory to the host memory. Modify based on the specific API definition.
  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 resultData from device to host failed. ERROR: %d\n", ret);
            return ret);
  for (int64_t i = 0; i < size; i++) {
    LOG_PRINT("resultData[%ld] is: %f\n", i, resultData[i]);
  }

  // 6. Release aclTensor. Modify based on the specific API definition.
  aclDestroyTensor(gradOutput);
  aclDestroyTensor(indices);
  aclDestroyTensor(out);

  // 7. Release device resources. Modify based on the specific API definition.
  aclrtFree(gradOutputDeviceAddr);
  aclrtFree(indicesDeviceAddr);
  aclrtFree(outDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}