aclnnFlatQuant
Applicable Products
| Product | Supported |
|---|---|
| √ | |
| √ | |
| √ | |
| × | |
| × | |
| × |
Function
Description: This fused operator performs two small matrix multiplications on the input matrix x in one pass: it right-multiplies x by the input matrix kroneckerP2 and left-multiplies x by the input matrix kroneckerP1, then quantizes the matrix multiplication result. Currently, pertoken and pergroup Quantization Modes are supported, corresponding to int4 and float4_e2m1 quantization output types, respectively.
Matrix multiplication formula:
- Input x is right-multiplied by kroneckerP2:
- kroneckerP1 left-multiplies x':
Quantization computation mode:
pertoken quantization mode:
- Compute the maximum absolute value along dimension 0 of x'' and divide by (7 / clipRatio) to calculate the quantization scale for INT4 format:
- Compute the output out:
pergroup quantization mode
After matrix multiplication, the shape of x'' is [K, M, N]. When computing mx_quantize in the pergroup quantization mode, reshape it to [K, M*N], denoted as x2.
Group elements along the second dimension of x2 by groupsize, with elements e0, e1, ..., e31. Compute emax.
- Compute reduceMaxValue and sharedExp.
- Compute quantScale.
- Share one quantScale per groupsize and compute out.
Function Prototype
Each operator is divided into a two-phase API. You must first call the aclnnFlatQuantGetWorkspaceSize API to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow, and then call the aclnnFlatQuant API to perform the computation.
aclnnStatus aclnnFlatQuantGetWorkspaceSize(
const aclTensor *x,
const aclTensor *kroneckerP1,
const aclTensor *kroneckerP2,
double clipRatio,
aclTensor *out,
aclTensor *quantScale,
uint64_t *workspaceSize,
aclOpExecutor **executor)aclnnStatus aclnnFlatQuant(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream)aclnnFlatQuantGetWorkspaceSize
Parameters:
Parameter Input/Output Description Instruction Data Type Data Format Shape Non-contiguous Tensor x (aclTensor*) Input The input raw data, corresponding to `x` in the formula. - Empty tensors are not supported.
- The shape is [K, M, N], where K does not exceed 262144, and M and N do not exceed 256.
- If the data type of out is INT32, N must be an integer multiple of 8.
- If the data type of out is INT4, N must be an even number.
- If the data type of out is FLOAT4_E2M1, N must be an even number.
FLOAT16, BFLOAT16 ND 3 √ kroneckerP1 (aclTensor*) Input Specifies the input computation matrix 1, corresponding to `kroneckerP1` in the formula. - Empty tensors are not supported.
- The shape is [M, M], where M is consistent with the M dimension in x.
- The data type is the same as that of the input parameter x.
FLOAT16, BFLOAT16 ND 2 √ kroneckerP2 (aclTensor*) Input Input computation matrix 2, corresponding to `kroneckerP2` in the formula. - Does not support empty tensors.
- Shape is [N, N], where N is consistent with the N dimension in x.
- The data type is consistent with the data type of the input parameter x.
FLOAT16, BFLOAT16 ND 2 √ clipRatio (double) Input Controls the clipping ratio for quantization, corresponding to `clipRatio` in the formula. - The input data range is (0, 1].
- - - - out (aclTensor*) Output Output tensor, corresponding to out in the formula. - Non-empty tensor required.
- When the data type is INT32, the last dimension of the shape is 1/8 of the last dimension of the input parameter x, and the other dimensions are the same as x.
- When the data type is INT4, the shape is the same as the input parameter x.
- When the data type is FLOAT4_E2M1, the shape is [K, M*N].
INT4, INT32, FLOAT4_E2M1 ND 3 or 2 √ quantScale (aclTensor*) Output Output quantization factor, corresponding to quantScale in the formula. - Empty tensors are not supported.
- When the output type is INT4 or INT32, the shape is [K], where K is consistent with the K dimension in x.
- When the output type is FLOAT8_E8M0, the shape is [K, ceilDiv(M*N, 64), 2].
FLOAT32, FLOAT_E8M0 ND 1 or 3 √ workspaceSize (uint64_t *) Output Returns the workspace size that the user needs to apply for on the Device side. - - - - - executor (aclOpExecutor**) Output 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 performs input parameter validation and returns an error in the following scenarios:
Return Code Error Code Description ACLNN_ERR_PARAM_NULLPTR 161001 The required input (x, kroneckerP1, kroneckerP2) or required output (out, quantScale) is a null pointer. ACLNN_ERR_PARAM_INVALID 161002 The data type or data format of x, kroneckerP1, kroneckerP2, out, or quantScale is not within the supported range. The data type of kroneckerP1 or kroneckerP2 is inconsistent with that of x. The dimension of x is not 3. The first dimension of x is outside the range [1, 262144], or the second dimension is outside [1, 256], or the third dimension is outside [1, 256]. The dimension of kroneckerP1 is not 2, or its first dimension and second dimension are inconsistent with the second dimension of x. The dimension of kroneckerP2 is not 2, or its first dimension and second dimension are inconsistent with the third dimension of x. In the int4 or int32 scenario, the dimension of quantScale is not 1, or its first dimension is inconsistent with the first dimension of x. In the float4_e2m1 scenario, the dimension of quantScale is not 3, or its first dimension is inconsistent with the first dimension of x. The value of clipRatio is outside the range (0, 1]. When the data type of out is INT4, the size of the last axis of x's shape is not an even number, or x's shape is inconsistent with out's shape. When the data type of out is INT32, the last axis of x's shape is not 8 times the size of the last axis of out's shape, or the sizes of the non-last axes of x's shape and out's shape are inconsistent.
aclnnFlatQuant
Parameters:
Parameter Input/Output Description workspace Input Address of the workspace memory applied for on the Device side. workspaceSize Input Size of the workspace applied for on the Device side, obtained from the first-phase API aclnnFlatQuantGetWorkspaceSize. executor Input Operator executor that contains the operator computation flow. stream Input Specifies the stream for task execution. Return value:
aclnnStatus: return code. For details, see aclnn Return Code.
Constraints
Deterministic computation:
- aclnnFlatQuant defaults to a deterministic implementation.
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_flat_quant.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 shape_size = 1;
for (auto i : shape) {
shape_size *= i;
}
return shape_size;
}
int Init(int32_t deviceId, aclrtStream* stream)
{
// Boilerplate: initialize resources.
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 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);
// Copy host-side data to device memory by calling aclrtMemcpy.
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];
}
// Create an aclTensor by calling the aclCreateTensor API.
*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.
// Set deviceId based on the actual device.
int32_t deviceId = 0;
aclrtStream stream;
auto ret = Init(deviceId, &stream);
// Handle based on your own needs.
CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
// 2. Construct inputs and outputs. Custom construction is required based on the API.
std::vector<int64_t> xShape = {16, 16, 16};
std::vector<int64_t> kroneckerP1Shape = {16, 16};
std::vector<int64_t> kroneckerP2Shape = {16, 16};
std::vector<int64_t> outShape = {16, 16, 2};
std::vector<int64_t> quantScaleShape = {16};
void* xDeviceAddr = nullptr;
void* kroneckerP1DeviceAddr = nullptr;
void* kroneckerP2DeviceAddr = nullptr;
void* outDeviceAddr = nullptr;
void* quantScaleDeviceAddr = nullptr;
aclTensor* x = nullptr;
aclTensor* kroneckerP1 = nullptr;
aclTensor* kroneckerP2 = nullptr;
aclTensor* out = nullptr;
aclTensor* quantScale = nullptr;
double clipRatio = 1.0;
std::vector<aclFloat16> xHostData(16 * 16 * 16, aclFloatToFloat16(1));
std::vector<aclFloat16> kroneckerP1HostData(16 * 16, aclFloatToFloat16(1));
std::vector<aclFloat16> kroneckerP2HostData(16 * 16, aclFloatToFloat16(1));
std::vector<int32_t> outHostData(16 * 16 * 2, 1);
std::vector<float> quantScaleHostData(16, 0);
// Create the x aclTensor.
ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT16, &x);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the kroneckerP1 aclTensor.
ret = CreateAclTensor(
kroneckerP1HostData, kroneckerP1Shape, &kroneckerP1DeviceAddr, aclDataType::ACL_FLOAT16, &kroneckerP1);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the kroneckerP2 aclTensor.
ret = CreateAclTensor(
kroneckerP2HostData, kroneckerP2Shape, &kroneckerP2DeviceAddr, aclDataType::ACL_FLOAT16, &kroneckerP2);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the out aclTensor.
ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_INT32, &out);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the quantScale aclTensor.
ret = CreateAclTensor(
quantScaleHostData, quantScaleShape, &quantScaleDeviceAddr, aclDataType::ACL_FLOAT, &quantScale);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// 3. Call the CANN operator library API. Replace it with the specific API.
uint64_t workspaceSize = 0;
aclOpExecutor* executor;
// Call the first-phase API of aclnnFlatQuant.
ret = aclnnFlatQuantGetWorkspaceSize(
x, kroneckerP1, kroneckerP2, clipRatio, out, quantScale, &workspaceSize, &executor);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlatQuantGetWorkspaceSize 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 aclnnFlatQuant.
ret = aclnnFlatQuant(workspaceAddr, workspaceSize, executor, stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlatQuant 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<int32_t> resultData(size, 0);
ret = aclrtMemcpy(
resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(int32_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 < size; i++) {
LOG_PRINT("result[%ld] is: %d\n", i, resultData[i]);
}
auto quantScaleSize = GetShapeSize(quantScaleShape);
std::vector<float> quantScaleResultData(quantScaleSize, 0);
ret = aclrtMemcpy(
quantScaleResultData.data(), quantScaleResultData.size() * sizeof(quantScaleResultData[0]),
quantScaleDeviceAddr, quantScaleSize * 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 < quantScaleSize; i++) {
LOG_PRINT("result[%ld] is: %f\n", i, quantScaleResultData[i]);
}
// 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
aclDestroyTensor(x);
aclDestroyTensor(kroneckerP1);
aclDestroyTensor(kroneckerP2);
aclDestroyTensor(out);
aclDestroyTensor(quantScale);
// 7. Release device resources. Modify based on the specific API definition.
aclrtFree(xDeviceAddr);
aclrtFree(kroneckerP1DeviceAddr);
aclrtFree(kroneckerP2DeviceAddr);
aclrtFree(outDeviceAddr);
aclrtFree(quantScaleDeviceAddr);
if (workspaceSize > 0) {
aclrtFree(workspaceAddr);
}
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
return 0;
}