aclnnQuantBatchMatmulInplaceAdd

📄 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 QuantBatchMatmul followed by InplaceAdd. The QuantBatchMatmulInplaceAdd operator fuses the above operators to improve network performance. It implements quantized matrix multiplication and addition computations, with the basic functionality being a combination of matrix multiplication and addition.

• Formula:

  • mx quantization:
  $$y[m,n] = \sum_{j=0}^{kLoops-1} ((\sum_{k=0}^{gsK-1} (x1Slice * x2Slice)) * (scale1[m, j] * scale2[j, n])) + y[m,n]$$

  Where $gsK$ represents the quantization block size of the K axis, i.e., 32; $x1Slice$ represents the vector of length $gsK$ in the m-th row of $x1$; $x2Slice$ represents the vector of length $gsK$ in the n-th column of $x2$. 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.

Function Prototype

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

aclnnStatus aclnnQuantBatchMatmulInplaceAddGetWorkspaceSize(
    const aclTensor *x1,
    const aclTensor *x2,
    const aclTensor *x1Scale,
    const aclTensor *x2Scale,
    aclTensor       *yRef,
    bool            transposeX1,
    bool            transposeX2,
    int64_t         groupSize,
    uint64_t        *workspaceSize,
    aclOpExecutor   **executor)

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

aclnnQuantBatchMatmulInplaceAddGetWorkspaceSize

• Parameters:

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension	Non-contiguous Tensor
  x1	Input	The aclTensor on the device side, which is input x1 in the formula.	-	FLOAT8_E4M3FN, FLOAT8_E5M2	ND	2	√
  x2	Input	The aclTensor on the device side, representing input x2 in the formula.	-	FLOAT8_E3M4FN, FLOAT8_E5M2	ND	2	√
  x1Scale	Optional Input	Indicates the scale factor introduced by x1 quantization in the quantization parameters, an aclTensor on the Device side.	For comprehensive constraints, see Constraints.	FLOAT8_E8M0	ND	3	√
  x2Scale	Input	Indicates the scale factor introduced by x2 quantization in the quantization parameters, an aclTensor on the device side.	For comprehensive constraints, see Constraints.	FLOAT8_E8M0	ND	3	√
  yRef	Input/Output	Device-side aclTensor, corresponding to the input/output y in the formula.	When input x1 is an empty tensor with m=0 or input x2 is an empty tensor with n=0, the output is an empty tensor.	FLOAT32	ND	2	√
  transposeX1	Input	Indicates whether to transpose the input shape of x1.	-	bool	-	-	-
  transposeX2	Input	Indicates whether to transpose the input shape of x2.	-	bool	-	-	-
  groupSize	Input	Integer parameter that specifies the quantization group size in the m, n, and k directions.	Consists of three values concatenated: groupSizeM, groupSizeN, and groupSizeK, each occupying 16 bits, together occupying the lower 48 bits of the int64_t type groupSize (the upper 16 bits of groupSize are invalid). See the formula below the table for details.	INT64	-	-	-
  workspaceSize	Output	Returns the workspace size to be applied for on the Device side.	-	-	-	-	-
  executor	Output	Returns the operator executor, which contains the operator computation flow.	-	-	-	-	-

  • Formula: $$groupSize = groupSizeK | groupSizeN << 16 | groupSizeM << 32$$

• 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 input parameter is a required input, output, or required attribute and is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data types and data formats of x1, x2, x1Scale, x2Scale, yRef, and groupSize are not within the supported range.
  		The shapes of x1, x2, x1Scale, x2Scale, and yRef do not meet the validation conditions.
  		x1, x2, x2Scale, and yRef are empty tensors.
  		The passed groupSize does not meet the validation conditions, or when the passed groupSize is 0, the shape relationship between x1, x2 and x1Scale, x2Scale cannot be used to infer groupSize.

aclnnQuantBatchMatmulInplaceAdd

• Parameters:

  Parameter	Input/Output	Description
  workspace	Input	Specifies the workspace memory address applied for on the Device side.
  workspaceSize	Input	Specifies the workspace size applied for on the Device side, which is obtained from the first-phase API aclnnQuantBatchMatmulInplaceAddGetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation flow.
  stream	Input	Specifies the AscendCL stream for executing tasks.

• Return value

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

Constraints

• Deterministic computation: aclnnQuantBatchMatmulInplaceAdd defaults to a deterministic implementation.

• Currently, only transposeX1 set to true and transposeX2 set to false are supported.

• Constraints on groupSize:

  • The input groupSize is internally decomposed into groupSizeM, groupSizeN, and groupSizeK using the following formula. When one or more of them are 0, groupSizeM, groupSizeN, and groupSizeK are reset based on the input shapes of x1, x2, x1Scale, and x2Scale for computation. Principle: If groupSizeM = 0, it indicates that the quantization group value in the m direction is inferred by the API. The inference formula is groupSizeM = m / scaleM (m must be divisible by scaleM), where m is consistent with the m in the x1 shape, and scaleM is consistent with the m in the x1Scale shape.$$groupSize = groupSizeK | groupSizeN << 16 | groupSizeM << 32$$

• Dynamic quantization (mx quantization) scenario constraints:

  • The input and output support the following data type combinations:

    x1	x2	x1Scale	x2Scale	outRef
    FLOAT8_E5M2/FLOAT8_E4M3FN	FLOAT8_E5M2/FLOAT8_E4M3FN	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT32

  • Relationship among the x1 data type, x2 data type, x1, x2, x1Scale, x2Scale, and groupSize:

    x1 Data Type	x2 Data Type	x1 Shape	x2 Shape	x1Scale Shape	x2Scale Shape	yRef Shape	[gsM, gsN, gsK]	groupSize
    FLOAT8_E5M2/FLOAT8_E4M3FN	FLOAT8_E5M2/FLOAT8_E4M3FN	(k, m)	(k, n)	(ceil(k / 64), m, 2)	(ceil(k / 64), n, 2)	(m, n)	[1, 1, 32]	32

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 "acl/acl.h"
#include "aclnnop/aclnn_quant_batch_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 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 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);

    // 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;
}

template <typename T1, typename T2>
auto CeilDiv(T1 a, T2 b) -> T1
{
    if (b == 0) {
        return a;
    }
    return (a + b - 1) / b;
}

void Finalize(int32_t deviceId, aclrtStream stream)
{
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
}

int AclnnQuantBatchMatmulInplaceAddTest(int32_t deviceId, 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. Custom construction based on the API is required.
    int64_t M = 8;
    int64_t K = 16;
    int64_t N = 8;

    std::vector<int64_t> x1Shape = {K, M};
    std::vector<int64_t> x2Shape = {K, N};
    std::vector<int64_t> x2ScaleShape = {CeilDiv(K, 64), N, 2};
    std::vector<int64_t> yInputShape = {M, N};
    std::vector<int64_t> x1ScaleShape = {CeilDiv(K, 64), M, 2};
    std::vector<int64_t> yOutShape = {M, N};

    void* x1DeviceAddr = nullptr;
    void* x2DeviceAddr = nullptr;
    void* x2ScaleDeviceAddr = nullptr;
    void* yInputDeviceAddr = nullptr;
    void* x1ScaleDeviceAddr = nullptr;
    void* yOutputDeviceAddr = nullptr;

    aclTensor* x1 = nullptr;
    aclTensor* x2 = nullptr;
    aclTensor* x2Scale = nullptr;
    aclTensor* yInput = nullptr;
    aclTensor* x1Scale = nullptr;
    aclTensor* yOutput = nullptr;

    std::vector<uint8_t> x1HostData(M * K, 1);                 // 0b00111000 represents 1.0 in fp8_e4m3fn.
    std::vector<uint8_t> x2HostData(N * K, 1); // 0b0010 represents 1.0 in fp4_e2m1. Here, uint8 is used to represent two fp4 values.
    std::vector<uint8_t> x2ScaleHostData(CeilDiv(K, 64) * N * 2, 1);
    std::vector<float> yInputHostData(M * N, 1);                        // 1.0 in fp32.
    std::vector<uint8_t> x1ScaleHostData(M * CeilDiv(K, 64) * 2, 1);
    std::vector<float> yOutputHostData(M * N, 1);

    // Create 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 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 x2Scale aclTensor.
    ret = CreateAclTensor(x2ScaleHostData, x2ScaleShape, &x2ScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &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 yInput aclTensor.
    ret = CreateAclTensor(yInputHostData, yInputShape, &yInputDeviceAddr, aclDataType::ACL_FLOAT, &yInput);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> yInputTensorPtr(yInput, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void*)> yInputDeviceAddrPtr(yInputDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // Create x1Scale aclTensor.
    ret = CreateAclTensor(x1ScaleHostData, x1ScaleShape, &x1ScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &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);

    bool transposeX1 = true;
    bool transposeX2 = false;
    int64_t groupSize = 32;

    // 3. Call the CANN operator library API. Replace with the specific API name.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    void* workspaceAddr = nullptr;

    // Call the first-phase API of aclnnQuantBatchMatmulInplaceAdd.
    ret = aclnnQuantBatchMatmulInplaceAddGetWorkspaceSize(x1, x2, x1Scale, x2Scale, yInput, transposeX1, transposeX2, groupSize, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantBatchMatmulInplaceAddGetWorkspaceSize failed. ERROR: %d\n", ret);
              return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    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 aclnnTransQuantParamV2.
    ret = aclnnQuantBatchMatmulInplaceAdd(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantBatchMatmulInplaceAdd 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(yInputShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), size * sizeof(uint32_t), yInputDeviceAddr,
                    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 ACL API manual.
    // Enter the actual device ID for deviceId.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = AclnnQuantBatchMatmulInplaceAddTest(deviceId, stream);
    CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("AclnnQuantBatchMatmulInplaceAddTest failed. ERROR: %d\n", ret); return ret);

    Finalize(deviceId, stream);
    return 0;
}