aclnnQuantGroupedMatmulInplaceAdd

📄 View Source Code

Applicable Products

Function

• Description: In micro-batch training scenarios where micro-batch gradient accumulation is required, there are numerous fusion scenarios involving GroupedMatMul followed by InplaceAdd. The QuantGroupedMatmulInplaceAdd operator fuses the above operators to improve network performance. It implements grouped matrix multiplication and addition computations. The basic function is a combination of matrix multiplication and addition, for example, in the T-C quantization scenario: $y_i[m,n]=(x1_i[m,k_i] \times x2_i[k_i,n]) * scale2_i[n] * scale1_i + y_i[m,n], i=1...g$, where g is the number of groups and $m/k_i/n$ are the corresponding dimensions.

  Compared with the aclnnGroupedMatmulV4 API, this API has the following changes:

  • Both input and output parameters are of the aclTensor type.

  • An InplaceAdd computation is added after the GroupedMatMul computation.

  • Only quantization scenarios are supported (1. mx quantization; 2. T-C quantization). For the quantization mode, see Quantization Introduction.

  • Only x1 and x2 inputs of the FLOAT8_E5M2, FLOAT8_E4M3FN, or HIFLOAT8 type are supported.

• Formula:

  • mx quantization:
  $$y_i[m,n] = \sum_{j=0}^{kLoops-1} ((\sum_{k=0}^{gsK-1} (x1Slice_i * x2Slice_i)) * (scale1_i[m, j] * scale2_i[j, n])) + y_i[m,n]$$

  Where gsK represents the quantization block size of the K axis, i.e., 32; $x1Slice_i$ represents the vector of length gsK in the m-th row of $x1_i$; $x2Slice_i$ represents the vector of length gsK in the n-th column of $x2_i$. The K axis starts slicing from $j*gsK$, with j in the range [0, kLoops), where kLoops = ceil($K_i$ / gsK). The case where the last slice length is less than gsK is supported.

  • T-C quantization:
  $$y_i=(x1_i\times x2_i) * scale2_i * scale1_i + y_i$$

Function Prototype

Each operator is divided into a two-phase API. You must first call the "aclnnQuantGroupedMatmulInplaceAddGetWorkspaceSize" API to obtain the input parameters and calculate the required workspace size based on the computation process, and then call the "aclnnQuantGroupedMatmulInplaceAdd" API to perform the computation.

aclnnStatus aclnnQuantGroupedMatmulInplaceAddGetWorkspaceSize(
    const aclTensor *x1,
    const aclTensor *x2,
    const aclTensor *scale1Optional,
    const aclTensor *scale2,
    const aclTensor *groupList,
    aclTensor       *yRef,
    int64_t          groupListType,
    int64_t          groupSize,
    uint64_t        *workspaceSize,
    aclOpExecutor   **executor)

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

aclnnQuantGroupedMatmulInplaceAddGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (shape)	Non-contiguous Tensor
  x1	Input	Device-side aclTensor representing input x1 in the formula.	-	FLOAT8_E4M3FN, FLOAT8_E5M2, HIFLOAT8	ND	2(K, M)	√
  x2	Input	Device-side aclTensor, representing input x2 in the formula.	-	FLOAT8_E4M3FN, FLOAT8_E5M2, HIFLOAT8	ND	2(K, N)	√
  scale1Optional	Optional input	Represents the scale factor introduced by x1 quantization in the quantization parameters, a Device-side aclTensor.	For comprehensive constraints, see Constraints.	FLOAT32, FLOAT8_E8M0	ND	1-3	√
  scale2	Input	Indicates the scale factor introduced by x2 quantization in the quantization parameters. It is an aclTensor on the Device side.	For comprehensive constraints, see Constraints.	FLOAT32, FLOAT8_E8M0	ND	2-3	√
  groupList	Input	Indicates the matmul size distribution along the group axis for input and output, as a Device-side aclTensor.	When groupListType is 0, groupList must be a non-negative monotonic non-decreasing sequence. When groupListType is 1, groupList must be a non-negative sequence. The last value in groupList constrains the valid portion of the output data. The portion not specified in groupList will not be updated.	INT64	ND	1(g,)	√
  yRef	Input/Output	Device-side aclTensor that corresponds to the input/output y in the formula.	When the M axis of x1 or the N axis of x2 is 0, yRef is an empty tensor.	FLOAT32	ND	3(g, M, N)	√
  groupListType	Optional input	Integer parameter. Only supports two values: 0 and 1.	0: The value in groupList is the cumsum result (cumulative sum) of the grouping axis size. 1: The value in groupList is the size of each group on the grouping axis.	INT64	-	-	-
  groupSize	Input	Integer parameter that specifies the quantization group size in the m, n, and k dimensions.	The groupSize input consists of three values — groupSizeM, groupSizeN, and groupSizeK — concatenated together. Each value occupies 16 bits, and together they occupy the lower 48 bits of the int64_t type groupSize (the upper 16 bits of groupSize are invalid). The formula is: groupSize = groupSizeK | groupSizeN << 16 | groupSizeM << 32. Currently, only 0 is supported.	INT64	-	-	-
  workspaceSize	Output	Returns the workspace size to be applied on the Device side.	-	-	-	-	-
  executor	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 Value	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	Indicates that a required input, output, or attribute parameter is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of x1, x2, scale2, groupList, yRef, scale1Optional, groupListType, or groupSize is not within the supported range.

aclnnQuantGroupedMatmulInplaceAdd

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Memory address of the workspace applied for on the Device side.
  workspaceSize	Input	Size of the workspace applied for on the Device side, obtained through the first-phase API aclnnQuantGroupedMatmulInplaceAddGetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation flow.
  stream	Input	Specifies the AscendCL stream that executes the task.

• Return Value

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

Constraints

• Deterministic computation: aclnnQuantGroupedMatmulInplaceAdd defaults to a deterministic implementation.

• The size of each dimension of x1 and x2, after 32-byte alignment, must be less than the maximum int32 value 2147483647, and the inner axis size must be less than 2097152.

  • The supported input types for the dynamic quantization (T-C quantization) scenario are:

    • The supported data type combinations for non-empty parameters must satisfy the following table:

      x1	x2	scale2	scale1Optional	yRef
      HIFLOAT8	HIFLOAT8	FLOAT32	FLOAT32	FLOAT32

    • scale1Optional/scale2 must meet the following constraints (where g is the number of matmul groups, i.e., the number of groups):

      Parameter	Shape Constraint
      scale1Optional	2D tensor or 1D tensor, with shape (g, 1) or (g,)
      scale2	2D tensor, with shape (g, N)

  • Data types supported in the dynamic quantization (MX quantization) scenario:

    • The data type combinations must meet the following table:

      x1	x2	scale2	scale1Optional	yRef
      FLOAT8_E5M2/FLOAT8_E4M3FN	FLOAT8_E5M2/FLOAT8_E4M3FN	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT32

    • scale1Optional/scale2 must satisfy the following constraints (where g is the number of matmul groups, i.e., the group count, and g_i is the i-th group (index starting from 0)):

      Parameter	Shape Constraint
      scale1Optional	A 3D tensor with shape ((K / 64) + g, M, 2). The start address offset of scale_i is ((K_0 + K_1 + ...+ K_{i-1})/ 64 + g_i) M 2, i.e., the start address offset of scale_0 is 0, the start address offset of scale_1 is (K_0 / 64 + 1) M 2, the start address offset of scale_2 is ((K_0 + K_1) / 64 + 2) M 2, and so on.
      scale2	A 3D tensor with shape ((K / 64) + g, N, 2). The start address offset follows the same rule as scale1Optional.

• The first dimension of groupList supports a maximum of 1024, i.e., a maximum of 1024 groups.

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 <memory>
#include <utility>
#include "acl/acl.h"
#include "aclnnop/aclnn_quant_grouped_matmul_inplace_add.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: 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 CreateTransposeAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                             aclDataType dataType, aclTensor** tensor) {
    std::vector<int64_t> view_shape = {shape[1], shape[0]};

    auto size = GetShapeSize(view_shape) * sizeof(T);
    // Apply for device-side memory by calling aclrtMalloc.
    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 the device-side 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);
    // 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];
    }
    // Swap strides.
    std::swap(strides[0], strides[1]);
    // Create an aclTensor by calling aclCreateTensor.
    *tensor = aclCreateTensor(view_shape.data(), view_shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    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 memory on the Device side.
    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 the Host side to the 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 the 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 aclnnQuantGroupedMatmulInplaceAddTest(int32_t deviceId, aclrtStream &stream) {
    auto ret = Init(deviceId, &stream);
    // Handle the check based on your own needs.
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl 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> x1Shape = {2, 3};
    std::vector<int64_t> x2Shape= {2, 3};
    std::vector<int64_t> scale2Shape = {2, 3};
    std::vector<int64_t> yShape = {2, 3, 3};
    std::vector<int64_t> scale1Shape = {2, 1};
    std::vector<int64_t> groupListShape = {2};
    void* x1DeviceAddr = nullptr;
    void* x2DeviceAddr = nullptr;
    void* scale2DeviceAddr = nullptr;
    void* scale1DeviceAddr = nullptr;
    void* yDeviceAddr = nullptr;
    void* groupListDeviceAddr = nullptr;
    aclTensor* x1 = nullptr;
    aclTensor* x2 = nullptr;
    aclTensor* groupList = nullptr;
    aclTensor* scale2 = nullptr;
    aclTensor* yRef = nullptr;
    aclTensor* scale1 = nullptr;
    aclTensor* out = nullptr;
    int64_t groupListType = 0;
    int64_t groupSize = 0;
    std::vector<uint8_t> xData(GetShapeSize(x1Shape), 0X10); // hifloat8 2.0 converted to hexadecimal 0X10
    std::vector<int64_t> groupListData = {1, 3};
    std::vector<float> scale2Data(GetShapeSize(scale2Shape), 1);
    std::vector<float> yData(GetShapeSize(yShape), 1);
    std::vector<float> scale1Data(GetShapeSize(scale1Shape), 1);
    // Create the x1 aclTensor.
    ret = CreateTransposeAclTensor<uint8_t>(xData, x1Shape, &x1DeviceAddr, aclDataType::ACL_HIFLOAT8, &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<uint8_t>(xData, x2Shape, &x2DeviceAddr, aclDataType::ACL_HIFLOAT8, &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 scale2 aclTensor.
    ret = CreateAclTensor<float>(scale2Data, scale2Shape, &scale2DeviceAddr, aclDataType::ACL_FLOAT, &scale2);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> scale2TensorPtr(scale2, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> scale2DeviceAddrPtr(scale2DeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the y aclTensor.
    ret = CreateAclTensor<float>(yData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &yRef);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> yTensorPtr(yRef, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> yDeviceAddrPtr(yDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the group_list aclTensor.
    ret = CreateAclTensor<int64_t>(groupListData, groupListShape, &groupListDeviceAddr, aclDataType::ACL_INT64, &groupList);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> groupListTensorPtr(groupList, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> groupListDeviceAddrPtr(groupListDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the scale1 aclTensor.
    ret = CreateAclTensor<float>(scale1Data, scale1Shape, &scale1DeviceAddr, aclDataType::ACL_FLOAT, &scale1);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> scale1TensorPtr(scale1, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> scale1DeviceAddrPtr(scale1DeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // 3. Call the CANN operator library API.
    // Call the first-phase API of aclnnQuantGroupedMatmulInplaceAdd.
    ret = aclnnQuantGroupedMatmulInplaceAddGetWorkspaceSize(x1, x2, scale1, scale2, groupList, yRef, groupListType, groupSize, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantGroupedMatmulInplaceAddGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    void* workspaceAddr = nullptr;
    std::unique_ptr<void, aclError (*)(void *)> workspaceAddrPtr(nullptr, aclrtFree);
    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);
        workspaceAddrPtr.reset(workspaceAddr);
    }
    // Call the second-phase API of aclnnQuantGroupedMatmulInplaceAdd.
    ret = aclnnQuantGroupedMatmulInplaceAdd(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantGroupedMatmulInplaceAdd failed. ERROR: %d\n", ret); return ret);
    // 4. (Boilerplate) Synchronously 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 side memory to the Host side. Modify this based on the specific API definition.
    auto size = GetShapeSize(yShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), size * sizeof(uint32_t), yDeviceAddr,
                      size * sizeof(uint32_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 j = 0; j < size; j++) {
        LOG_PRINT("result[%ld] is: %f\n", j, resultData[j]);
    }
    return ACL_SUCCESS;
}
int main()
{
    // 1. (Boilerplate) Initialize the device/stream. Refer to the AscendCL external API list.
    // Set deviceId based on the actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = aclnnQuantGroupedMatmulInplaceAddTest(deviceId, stream);
    CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantGroupedMatmulInplaceAddTest failed. ERROR: %d\n", ret); return ret);
    Finalize(deviceId, stream);
    return 0;
}