aclnnTransposeQuantBatchMatMul
Applicable Products
| Product | Supported |
|---|---|
| √ | |
| × | |
| × | |
| × | |
| × | |
| × |
Function
Description: Performs quantized matrix multiplication of tensor x1 and tensor x2, supporting the K-C quantization mode. Only three-dimensional tensors are supported. Tensors support transposition, and the transpose sequence changes based on the input sequence. permX1 represents the transpose sequence of tensor x1, and permX2 represents the transpose sequence of tensor x2. The dimension with a sequence value of 0 is the batch dimension, and the remaining two dimensions are used for matrix multiplication.
Example:
Assume that the shape of x1 is (M, B, K), the shape of x2 is (B, K, N), x1Scale and x2Scale are not None, and batchSplitFactor equals 1. The shape of the computed output out is (M, B, N).
Function Prototype
Each operator is divided into a two-phase API. You must first call the "aclnnTransposeQuantBatchMatMulGetWorkspaceSize" API to obtain the input parameters and calculate the required workspace size based on the process, and then call the "aclnnTransposeQuantBatchMatMul" API to perform the computation.
aclnnStatus aclnnTransposeQuantBatchMatMulGetWorkspaceSize(
const aclTensor* x1,
const aclTensor* x2,
const aclTensor* bias,
const aclTensor* x1Scale,
const aclTensor* x2Scale,
const int32_t dtype,
const int32_t groupSize,
const aclIntArray* permX1,
const aclIntArray* permX2,
const aclIntArray* permY,
const int32_t batchSplitFactor,
aclTensor* out,
uint64_t* workspaceSize,
aclOpExecutor** executor)aclnnStatus aclnnTransposeQuantBatchMatMul(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
const aclrtStream stream)aclnnTransposeQuantBatchMatMulGetWorkSpaceSize
Parameters:
Parameter Input/Output Description Instruction Data Type Data Format Shape Non-contiguous Tensor x1 (aclTensor*) Input Represents the first matrix of the matrix multiplication. - The data type must satisfy the data type deduction rules with x2 (see Deduction Relationship and Constraints).
- Currently, only FLOAT8_E5M2 and FLOAT8_E4M3FN are supported.
FLOAT8_E5M2, FLOAT8_E4M3FN ND 3 √ x2 (aclTensor*) Input Represents the second matrix of the matrix multiplication. - The data type must comply with the data type deduction rules with x1 (see Deduction Relationship and Constraints).
- The k dimension of x2 must be equal in size to the k dimension of x1.
- Currently, only FLOAT8_E5M2 and FLOAT8_E4M3FN data types are supported.
FLOAT8_E5M2, FLOAT8_E4M3FN ND 3 √ bias (aclTensor*) Input Represents the bias matrix for matrix multiplication. Reserved parameter, currently not supported. BFLOAT16, FLOAT16, FLOAT32 - - - x1Scale (aclTensor*) Input Indicates the quantization coefficient of the left matrix. The shape supports only one dimension and must be equal to [m]. FLOAT32 ND 1 √ x2Scale (aclTensor*) Input Represents the quantization scale of the right matrix. The shape supports only one dimension and must satisfy being equal to [n]. FLOAT32 ND 1 √ dtype (int32_t) Input Specifies the data type of the output matrix. Supported values: 1, 27. - A value of 1 indicates that the output matrix type is FLOAT16.
- A value of 27 indicates that the output matrix type is BFLOAT16.
INT32 - - - groupSize (int32_t) Input Specifies the quantization group size. This is a reserved parameter and currently only supports a value of 0. Other values do not take effect. Non-zero values do not take effect in the current configuration. INT32 - - - permX1 (aclIntArray*) Input Indicates the transpose sequence of the first matrix in the matrix multiplication. It is an aclIntArray on the host side. Supports [1, 0, 2]. INT64 - 1 - permX2 (aclIntArray*) Input Specifies the transpose sequence of the second matrix in the matrix multiplication. It is an aclIntArray on the host side. Supports [0, 1, 2]. INT64 - 1 - permY (aclIntArray*) Input Transpose sequence of the matrix multiplication output matrix, which is a host-side aclIntArray. Supports [1, 0, 2]. INT64 - 1 - batchSplitFactor (int32_t) Input Specifies the split size of the B dimension in the matrix multiplication output matrix. It is a host-side integer and currently only supports a value of 1. Currently, only a value of 1 is supported. INT32 - - - out (aclTensor*) Output Represents the output matrix of the matrix multiplication, i.e., out in the formula. - The data type must be consistent with the data type derived from x1 and x2 (see Deduction Relationship and Constraints).
- Currently, only non-empty x1Scale and x2Scale, groupSize of 0, batchSplitFactor of 1, and output shape of (M, B, N) are supported.
BFLOAT16, FLOAT16 ND 3 - 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 x1, x2, out, x1Scale, x2Scale, permX1, permX2, or permY is a null pointer. ACLNN_ERR_PARAM_INVALID 161002 The data type of x1, x2, or out is not within the supported range. The second dimension of x1 is not equal to the first dimension of x2. The dimension sizes of x1, x2, permX1, permX2, and permY are not equal to 3. The dimension sizes of x1Scale and x2Scale are not equal to 1. batchSplitFactor is not within the supported range. The data types of x1Scale and x2Scale are not within the supported range. The values of permX1, permX2, and permY are not within the supported range. The shapes of x1Scale and x2Scale do not meet the requirements.
aclnnTransposeQuantBatchMatMul
Parameters:
Parameter Input/Output Description workspace Input Memory address of the workspace applied for on the Device side. workspaceSize Input Workspace size applied on the device side, obtained through the first-phase API aclnnTransposeQuantBatchMatMulGetWorkspaceSize. executor Input Operator executor, which contains the operator computation flow. stream Input Specifies the stream for executing the task. Return value
aclnnStatus: return code. For details, see aclnn Return Code.
Constraints
Deterministic computation description: aclnnTransposeQuantBatchMatMul defaults to a deterministic implementation.
Ascend 950PR/Ascend 950DT :permX1 and permY support [1, 0, 2], and permX2 supports input [0, 1, 2].
x1Scale and x2Scale are 1-dimensional, with x1Scale being (M,) and x2Scale being (N,).
out and dtype support float16 and bfloat16.
Example
The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.
#include <iostream>
#include <memory>
#include <vector>
#include <limits>
#include "acl/acl.h"
#include "aclnnop/aclnn_transpose_quant_batch_mat_mul.h"
#define CHECK_RET(cond, return_expr) \
do { \
if (!(cond)) { \
return_expr; \
} \
} while (0)
#define CHECK_FREE_RET(cond, return_expr) \
do { \
if (!(cond)) { \
Finalize(deviceId, stream); \
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;
}
// Conversion function from BF16 to float.
float bf16_to_float(uint16_t bf16)
{
uint16_t sign = (bf16 >> 15) & 0x1;
uint16_t exp = (bf16 >> 7) & 0xFF; // 8-bit exponent
uint16_t mant = bf16 & 0x7F;
// Handle special values.
if (exp == 0) {
if (mant == 0) {
return sign ? -0.0f : 0.0f;
} else {
// Denormalized BF16 to float
return (sign ? -1.0f : 1.0f) * (float)mant * (1.0f / (1 << 7)) * (1.0f / (1, 127));
}
} else if (exp == 255) {
// Infinity or NaN
if (mant == 0) {
return sign ? -std::numeric_limits<float>::infinity() : std::numeric_limits<float>::infinity();
} else {
return std::numeric_limits<float>::quiet_NaN();
}
} else {
// Normalized number
float f_exp = (float)(exp - 127); // Offset by 127.
float f_mant = (float)mant / (1 << 7); // 7-bit mantissa.
float f = (sign ? -1.0f : 1.0f) * (1.0f + f_mant) * (1 << (int)f_exp);
return f;
}
}
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 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 data from host to device 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);
// Calculate 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;
}
void Finalize(int32_t deviceId, aclrtStream stream)
{
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
}
int AclnnTransposeQuantBatchMatmulTest(int32_t deviceId, aclrtStream& stream)
{
auto ret = Init(deviceId, &stream);
// Perform check as needed.
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
// 2. Construct the inputs and outputs. Customize the construction based on the API.
int32_t M = 32;
int32_t K = 512;
int32_t N = 128;
int32_t Batch = 16;
std::vector<int64_t> x1Shape = {M, Batch, K};
std::vector<int64_t> x2Shape = {Batch, K, N};
std::vector<int64_t> x1ScaleShape = {M};
std::vector<int64_t> x2ScaleShape = {N};
std::vector<int64_t> outShape = {M, Batch, N};
std::vector<int64_t> permX1Series = {1, 0, 2};
std::vector<int64_t> permX2Series = {0, 1, 2};
std::vector<int64_t> permYSeries = {1, 0, 2};
void* x1DeviceAddr = nullptr;
void* x2DeviceAddr = nullptr;
void* x1ScaleDeviceAddr = nullptr;
void* x2ScaleDeviceAddr = nullptr;
void* outDeviceAddr = nullptr;
aclTensor* x1 = nullptr;
aclTensor* x2 = nullptr;
aclTensor* x1Scale = nullptr;
aclTensor* x2Scale = nullptr;
aclTensor* out = nullptr;
std::vector<int8_t> x1HostData(GetShapeSize(x1Shape), 0x38);
std::vector<int8_t> x2HostData(GetShapeSize(x2Shape), 0x38);
std::vector<float> x1ScaleHostData(GetShapeSize(x1ScaleShape), 1);
std::vector<float> x2ScaleHostData(GetShapeSize(x2ScaleShape), 1);
std::vector<uint16_t> outHostData(GetShapeSize(outShape), 0); // bf16
// Create the x1 aclTensor.
ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x1);
std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x1TensorPtr(x1, aclDestroyTensor);
std::unique_ptr<void, aclError (*)(void*)> x1deviceAddrPtr(x1DeviceAddr, aclrtFree);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the x2 aclTensor.
ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x2);
std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2TensorPtr(x2, aclDestroyTensor);
std::unique_ptr<void, aclError (*)(void*)> x2deviceAddrPtr(x2DeviceAddr, aclrtFree);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the x1Scale aclTensor.
ret = CreateAclTensor(x1ScaleHostData, x1ScaleShape, &x1ScaleDeviceAddr, aclDataType::ACL_FLOAT, &x1Scale);
std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x1ScaleTensorPtr(x1Scale, aclDestroyTensor);
std::unique_ptr<void, aclError (*)(void*)> x1ScaledeviceAddrPtr(x1ScaleDeviceAddr, aclrtFree);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the x2Scale aclTensor.
ret = CreateAclTensor(x2ScaleHostData, x2ScaleShape, &x2ScaleDeviceAddr, aclDataType::ACL_FLOAT, &x2Scale);
std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2ScaleTensorPtr(x2Scale, aclDestroyTensor);
std::unique_ptr<void, aclError (*)(void*)> x2ScaledeviceAddrPtr(x2ScaleDeviceAddr, aclrtFree);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the out aclTensor.
ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_BF16, &out);
std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outTensorPtr(out, aclDestroyTensor);
std::unique_ptr<void, aclError (*)(void*)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
CHECK_RET(ret == ACL_SUCCESS, return ret);
aclIntArray* permX1 = aclCreateIntArray(permX1Series.data(), permX1Series.size());
aclIntArray* permX2 = aclCreateIntArray(permX2Series.data(), permX2Series.size());
aclIntArray* permY = aclCreateIntArray(permYSeries.data(), permYSeries.size());
uint64_t workspaceSize = 0;
aclOpExecutor* executor = nullptr;
std::unique_ptr<void, aclError (*)(void*)> executorAddrPtr(nullptr, aclrtFree);
int32_t batchSplitFactor = 1;
int32_t groupSize = 0;
int32_t dtype = 27; // bf16
// Example of the aclnnTransposeQuantBatchMatMul API call.
// 3. Call the CANN operator library API. Replace it with the specific API name.
// Call the first-phase API of aclnnTransposeQuantBatchMatMul.
ret = aclnnTransposeQuantBatchMatMulGetWorkspaceSize(
x1, x2, (const aclTensor*)nullptr, x1Scale, x2Scale, dtype, groupSize, permX1, permX2, permY, batchSplitFactor,
out, &workspaceSize, &executor);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTransposeQuantBatchMatMulGetWorkspaceSize 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);
executorAddrPtr.reset(workspaceAddr);
}
// Call the second-phase API of aclnnTransposeQuantBatchMatMul.
ret = aclnnTransposeQuantBatchMatMul(workspaceAddr, workspaceSize, executor, stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTransposeQuantBatchMatMul failed. ERROR: %d\n", ret); return ret);
// 4. (Boilerplate) Synchronize and wait for 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 memory. Modify based on the specific API definition.
auto size = GetShapeSize(outShape);
std::vector<uint16_t> resultData(size, 0); // bf16
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);
float resultDataBF16 = 0;
for (int64_t i = 0; i < size; i++) {
resultDataBF16 = bf16_to_float(resultData[i]);
LOG_PRINT("result[%ld] is: %f\n", i, resultDataBF16);
}
return ACL_SUCCESS;
}
int main()
{
// 1. (Boilerplate) Initialize the device/stream. See the ACL API manual.
// Set the deviceId based on the actual device.
int32_t deviceId = 0;
aclrtStream stream;
auto ret = AclnnTransposeQuantBatchMatmulTest(deviceId, stream);
CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTransposeQuantBatchMatMulTest failed. ERROR: %d\n", ret);
return ret);
Finalize(deviceId, stream);
return 0;
}