aclnnEmbeddingBag
Applicable Products
| Product | Supported |
|---|---|
| √ | |
| √ | |
| √ | |
| × | |
| √ | |
| √ |
Function
Description: Obtains a group of aggregated values from weight based on indices, and then performs max, sum, or mean aggregation on the obtained values according to the offset of offsets and the aggregation mode specified by mode. The remaining parameters further refine the control of the computation process.
Shape deduction formula:
Assume that the shape of the input weight is (numWeight, embeddingDim), the shape of indices is (bagIndices), and the shape of offsets is (bagOffsets).
Atlas A3 training products/Atlas A3 inference products ,Atlas A2 training products/Atlas A2 inference products ,Atlas inference products ,Atlas training products :- When mode is set to sum mode:
The shape of output is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false. The shape of offset2bag is (bagIndices,). The shape of bagSize is (bagOffsets - 1) if includeLastOffset is true, or (bagOffsets, ) if includeLastOffset is false. The shape of maxIndices is (bagOffsets - 1) if includeLastOffset is true, or (bagOffsets, ) if includeLastOffset is false.- When mode is set to mean mode:
The shape of output is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false. The shape of offset2bag is (bagIndices,). The shape of bagSize is (bagOffsets - 1) if includeLastOffset is true, or (bagOffsets, ) if includeLastOffset is false. The shape of maxIndices is (bagOffsets - 1) if includeLastOffset is true, or (bagOffsets, ) if includeLastOffset is false.- When mode is set to max mode:
The shape of output is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false. The shape of offset2bag is (bagIndices,). The shape of bagSize is (bagOffsets - 1) if includeLastOffset is true, or (bagOffsets, ) if includeLastOffset is false. The shape of maxIndices is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false.Ascend 950PR/Ascend 950DT :- When mode is set to sum:
The shape of output is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false. The shape of offset2bag is (bagIndices,). The shape of bagSize is (bagOffsets,). The shape of maxIndices is (0,).- When mode is set to mean:
The shape of output is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false. The shape of offset2bag is (bagIndices,). The shape of bagSize is bagOffsets. The shape of maxIndices is (0,).- When mode is set to max:
The shape of output is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false. The shape of offset2bag is (bagIndices,). The shape of bagSize is (bagOffsets,). The shape of maxIndices is (bagOffsets - 1, embeddingDim) if includeLastOffset is true, or (bagOffsets, embeddingDim) if includeLastOffset is false.
Function Prototype
Each operator is divided into a two-phase API. The "aclnnEmbeddingBagGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnEmbeddingBag" API is called to perform the computation.
aclnnStatus aclnnEmbeddingBagGetWorkspaceSize(
const aclTensor* weight,
const aclTensor* indices,
const aclTensor* offsets,
bool scaleGradByFreq,
int64_t mode,
bool sparse,
const aclTensor* perSampleWeights,
bool includeLastOffset,
int64_t paddingIdx,
aclTensor* output,
aclTensor* offset2bag,
aclTensor* bagSize,
aclTensor* maxIndices,
uint64_t* workspaceSize,
aclOpExecutor** executor)aclnnStatus aclnnEmbeddingBag(
void* workspace,
uint64_t workspaceSize,
aclOpExecutor* executor,
aclrtStream stream)`aclnnEmbeddingBagGetWorkspaceSize
Parameters
Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-contiguous Tensor weight Input Word embedding matrix that contains embedding vectors of all words. - FLOAT, FLOAT16, BFLOAT16 ND 2 √ indices Input Tensor containing indices, which specifies which word embedding vectors to extract from weight. - INT32, INT64 ND 0-2 √ offsets Input Offset tensor used to split indices into multiple bags. - INT32, INT64 - 0-1 - scaleGradByFreq Input Controls whether to scale the gradient by word frequency. When scaleGradByFreq is set to true, the gradient is scaled by word frequency; when scaleGradByFreq is set to false, it is not. - - - - mode Input Used to control the aggregation mode. 0 indicates sum aggregation mode, 1 indicates mean aggregation mode, and other values indicate max aggregation mode. - - - - sparse Input Used to control the sparse mode. - When set to false, indicates that weight is a non-sparse matrix;
- When set to true, indicates that weight is a sparse matrix.
- - - - perSampleWeights Input Specifies the sample weights. Can be non-nullptr only in sum mode, and must be nullptr in other modes. FLOAT, FLOAT16, BFLOAT16 - 0-2 - includeLastOffset Input Controls whether to include the last offset. - When set to false, indicates that the last offset is not included.
- When set to true, indicates that the last offset is included.
- When indices is 2-dimensional, only includeLastOffset set to false is supported.
- - - - paddingIdx Input Controls the indices that do not participate in computation. - - - - - Output Output Result after aggregation of the word embedding matrix. - Consistent with weight. ND 2 √ offset2bag Output Starting offset of the bag. The data type of offset2bag is consistent with the higher precision of indices and offsets. INT32, INT64 ND 0-1 √ bagSize Output Size of each bag. The data type of bagSize is consistent with the higher precision of indices and offsets. INT32, INT64 ND 0-1 √ maxIndices Output When mode is set to max, the row where the maximum value of the embedding vector resides. The data type of maxIndices is consistent with the higher precision of indices and offsets. INT32, INT64 ND 1 √ workspaceSize Output Returns the workspace size to be applied on the Device side. - - - - - executor Output Returns the op executor, which includes the computation process of the operator. - - - - - Atlas A2 training products/Atlas A2 inference products indices only support 0-1 dimensions, and perSampleWeights only supports 1 dimension.
Return Value
aclnnStatus: return code. For details, see aclnn Return Code.
The first-phase API performs input parameter validation and returns an error in the following scenarios:
Return Value Error Code Description ACLNN_ERR_PARAM_NULLPTR 161001 The passed weight, indices, offsets, output, offset2bag, bagSize, and maxIndices are null pointers. ACLNN_ERR_PARAM_INVALID 161002 The data type of weight is not supported, or the weight dimension is not 2. The data type of indices is not supported, or the indices dimension is not 1. The offsets data type is not supported, or the offsets dimension is not 1D. Neither the indices nor the offsets data type is INT32 or INT64. When perSampleWeights is not nullptr, its data type is inconsistent with weight, perSampleWeights is not 1D, the number of elements in perSampleWeights is not equal to that of indices, or perSampleWeights is not nullptr in non-sum mode. The output data type is inconsistent with weight, or its shape does not match the definition. The data type and shape of offset2bag, bagSize, and maxIndices do not match the deduced data type and shape.
aclnnEmbeddingBag
Parameters
Parameter Input/Output Description workspace Input Workspace memory address applied on the Device side. workspaceSize Input Workspace size applied on the Device side, obtained from the first-phase API aclnnEmbeddingBagGetWorkspaceSize. executor Input Operator executor, which includes the computation flow of the operator. stream Input Specifies the stream for task execution. Return Value
aclnnStatus: return code. For details, see aclnn Return Code.
Constraints
Deterministic computation:
- aclnnEmbeddingBag defaults to a deterministic implementation.
sparse and scaleGradByFreq only support input False.
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_bag.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 for 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 apply for 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 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 failed. ERROR: %d\n", ret); return ret);
// 2. Construct the input and output. Custom construction is required based on the API interface.
std::vector<int64_t> weightShape = {3, 3};
std::vector<int64_t> indicesShape = {6};
std::vector<int64_t> offsetsShape = {4};
std::vector<int64_t> perSampleWeightsShape = {6};
std::vector<int64_t> outputShape = {4, 3};
std::vector<int64_t> offset2bagShape = {6};
std::vector<int64_t> bagSizeShape = {4};
std::vector<int64_t> maxIndicesShape = {4};
void* weightDeviceAddr = nullptr;
void* indicesDeviceAddr = nullptr;
void* offsetsDeviceAddr = nullptr;
void* perSampleWeightsDeviceAddr = nullptr;
void* outputDeviceAddr = nullptr;
void* offset2bagDeviceAddr = nullptr;
void* bagSizeDeviceAddr = nullptr;
void* maxIndicesDeviceAddr = nullptr;
aclTensor* weight = nullptr;
aclTensor* indices = nullptr;
aclTensor* offsets = nullptr;
aclTensor* perSampleWeights = nullptr;
aclTensor* output = nullptr;
aclTensor* offset2bag = nullptr;
aclTensor* bagSize = nullptr;
aclTensor* maxIndices = nullptr;
std::vector<float> weightHostData = {1, 1, 1, 1, 1, 1, 1, 1, 1};
std::vector<int64_t> indicesHostData = {1, 2, 0, 2, 2, 1};
std::vector<int64_t> offsetsHostData = {0, 2, 4, 5};
std::vector<float> perSampleWeightsHostData = {1, 1, 1, 1, 1, 1};
std::vector<float> outputHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
std::vector<int64_t> offset2bagHostData = {0, 0, 0, 0, 0, 0};
std::vector<int64_t> bagSizeHostData = {0, 0, 0, 0};
std::vector<int64_t> maxIndicesHostData = {0, 0, 0, 0};
//Create the weight aclTensor.
ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_FLOAT, &weight);
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 offsets aclTensor.
ret = CreateAclTensor(offsetsHostData, offsetsShape, &offsetsDeviceAddr, aclDataType::ACL_INT64, &offsets);
CHECK_RET(ret == ACL_SUCCESS, return ret);
//Create perSampleWeights aclTensor.
ret = CreateAclTensor(perSampleWeightsHostData, perSampleWeightsShape, &perSampleWeightsDeviceAddr, aclDataType::ACL_FLOAT, &perSampleWeights);
CHECK_RET(ret == ACL_SUCCESS, return ret);
//Create output aclTensor.
ret = CreateAclTensor(outputHostData, outputShape, &outputDeviceAddr, aclDataType::ACL_FLOAT, &output);
CHECK_RET(ret == ACL_SUCCESS, return ret);
//Create offset2bag aclTensor.
ret = CreateAclTensor(offset2bagHostData, offset2bagShape, &offset2bagDeviceAddr, aclDataType::ACL_INT64, &offset2bag);
CHECK_RET(ret == ACL_SUCCESS, return ret);
//Create bagSize aclTensor.
ret = CreateAclTensor(bagSizeHostData, bagSizeShape, &bagSizeDeviceAddr, aclDataType::ACL_INT64, &bagSize);
CHECK_RET(ret == ACL_SUCCESS, return ret);
//Create maxIndices aclTensor.
ret = CreateAclTensor(maxIndicesHostData, maxIndicesShape, &maxIndicesDeviceAddr, aclDataType::ACL_INT64, &maxIndices);
CHECK_RET(ret == ACL_SUCCESS, return ret);
//Non-tensor parameter.
bool scaleGradByFreq = false;
int64_t mode = 0;
bool sparse = false;
bool includeLastOffset = false;
int64_t paddingIdx = 1;
// 3. Call the CANN operator library API. Replace with the specific API name.
uint64_t workspaceSize = 0;
aclOpExecutor* executor;
// Call the first-phase API of aclnnEmbeddingBag.
ret = aclnnEmbeddingBagGetWorkspaceSize(weight, indices, offsets, scaleGradByFreq, mode, sparse, perSampleWeights,
includeLastOffset, paddingIdx, output, offset2bag, bagSize, maxIndices, &workspaceSize, &executor);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnEmbeddingBagGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
// Apply for 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 aclnnEmbeddingBag.
ret = aclnnEmbeddingBag(workspaceAddr, workspaceSize, executor, stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnEmbeddingBag failed. ERROR: %d\n", ret); return ret);
// 4. (Boilerplate) Synchronously wait for the task execution to complete.
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 side. Modify based on the specific API definition.
auto outputSize = GetShapeSize(outputShape);
std::vector<float> outputResultData(outputSize, 0);
ret = aclrtMemcpy(outputResultData.data(), outputResultData.size() * sizeof(outputResultData[0]), outputDeviceAddr,
outputSize * sizeof(float), 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 < outputSize; i++) {
LOG_PRINT("outputResult[%ld] is: %f\n", i, outputResultData[i]);
}
auto offset2bagSize = GetShapeSize(offset2bagShape);
std::vector<int64_t> offset2bagResultData(offset2bagSize, 0);
ret = aclrtMemcpy(offset2bagResultData.data(), offset2bagResultData.size() * sizeof(offset2bagResultData[0]), offset2bagDeviceAddr,
offset2bagSize * sizeof(int64_t), 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 < offset2bagSize; i++) {
LOG_PRINT("offset2bagResult[%ld] is: %ld\n", i, offset2bagResultData[i]);
}
auto bagSizeSize = GetShapeSize(bagSizeShape);
std::vector<int64_t> bagSizeResultData(bagSizeSize, 0);
ret = aclrtMemcpy(bagSizeResultData.data(), bagSizeResultData.size() * sizeof(bagSizeResultData[0]), bagSizeDeviceAddr,
bagSizeSize * sizeof(int64_t), 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 < bagSizeSize; i++) {
LOG_PRINT("bagSizeResult[%ld] is: %ld\n", i, bagSizeResultData[i]);
}
auto maxIndicesSize = GetShapeSize(maxIndicesShape);
std::vector<int64_t> maxIndicesResultData(maxIndicesSize, 0);
ret = aclrtMemcpy(maxIndicesResultData.data(), maxIndicesResultData.size() * sizeof(maxIndicesResultData[0]), maxIndicesDeviceAddr,
maxIndicesSize * sizeof(int64_t), 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 < maxIndicesSize; i++) {
LOG_PRINT("maxIndicesResult[%ld] is: %ld\n", i, maxIndicesResultData[i]);
}
// 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
aclDestroyTensor(weight);
aclDestroyTensor(indices);
aclDestroyTensor(offsets);
aclDestroyTensor(perSampleWeights);
aclDestroyTensor(output);
aclDestroyTensor(offset2bag);
aclDestroyTensor(bagSize);
aclDestroyTensor(maxIndices);
// 7. Release device resources. Modify based on the specific API definition.
aclrtFree(weightDeviceAddr);
aclrtFree(indicesDeviceAddr);
aclrtFree(offsetsDeviceAddr);
aclrtFree(perSampleWeightsDeviceAddr);
aclrtFree(outputDeviceAddr);
aclrtFree(offset2bagDeviceAddr);
aclrtFree(bagSizeDeviceAddr);
aclrtFree(maxIndicesDeviceAddr);
if (workspaceSize > 0) {
aclrtFree(workspaceAddr);
}
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
return 0;
}