aclnnSparse4to2QuantMatmulWeightNz

Applicable Products

Product	Supported
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√

Function

Operator function:

Performs sparse 4:2 quantized matrix multiplication.
Formula:
$out = x@sparseWeight * sparseWeightScale * xScale + bias$

Function Prototype

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

aclnnStatus aclnnSparse4to2QuantMatmulWeightNzGetWorkspaceSize(
    const aclTensor *x,
    const aclTensor *sparseWeight,
    const aclTensor *index,
    const aclTensor *xScale,
    const aclTensor *sparseWeightScale,
    const aclTensor *biasOptional,
    aclTensor       *out,
    uint64_t        *workspaceSize,
    aclOpExecutor   **executor)

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

aclnnSparse4to2QuantMatmulWeightNzGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-contiguous Tensor
x	Input	The left matrix in the matrix multiplication operation.	Only 2D input is supported, with the shape expressed as (m, k).	INT8	ND	2	×
sparseWeight	Input	The right matrix in the matrix multiplication operation, which is a sparse matrix processed by the aclnnTransSparse4to2Para API.	Only 2D input is supported. ViewShape can be expressed as (n, k_half), where k_half = ceil(k / 8) * 8 / 2. StorageShape is obtained from the output parameters sparseWeightDims and sparseWeightDimsNum of the aclnnTransSparse4to2Para API. The obtained value is 4D, and its relationship with ViewShape is (ceil(k_half / 32), ceil(n / 16), 16, 32).	INT8	FRACTAL_NZ	2	×
Index	Input	Index matrix computed after compression through the aclnnTransSparse4to2Para API.	4-dimensional input. ViewShape can be expressed as (ceil(k_half / 32), ceil(n / 16), 16, 8). StorageShape is obtained from the output parameters indexDims and indexDimsNum of the aclnnTransSparse4to2Para API. The obtained value is 4-dimensional, identical to ViewShape, which is (ceil(k_half / 32), ceil(n / 16), 16, 8).	UINT8	ND	4	×
xScale	Input	Quantization parameter, scaling factor.	The shape can be expressed as (m,).	FLOAT32	ND	1	×
sparseWeightScale	Input	Quantization parameter, scaling factor.	The shape can be expressed as (n,).	FLOAT32	ND	1	×
biasOptional	Input	Optional bias.	shape can be expressed as (n,).	BFLOAT16	ND	1	×
out	Output	Output tensor, which is the out in the formula.	shape can be expressed as (m, n).	BFLOAT16	ND	2	×
workspaceSize	Output	Returns the workspace size to be requested 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 Code	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	The required input x, sparseWeight, index, or out is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type, shape, data format, or dtype of the input or output is not within the supported range.

aclnnSparse4to2QuantMatmulWeightNz

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, obtained from the first-phase API aclnnSparse4to2QuantMatmulGetWorkspaceSize.
executor	Input	Operator executor, which 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

The value of k in the last dimension of x, i.e., in the shape description, cannot exceed 65535.
Currently, only the scenario where neither sparseWeightScale nor xScale is nullptr is supported.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:

#include <iostream>
#include <memory>
#include <vector>
#include <stdlib.h>

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

#define CREATE_TENSOR(hostData, shape, deviceAddr, dtype, tensor)                                        \
    ret = CreateAclTensor(hostData, shape, &deviceAddr, dtype, &tensor);                                 \
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> tensor##Ptr(tensor, aclDestroyTensor); \
    std::unique_ptr<void, aclError (*)(void*)> deviceAddr##Ptr(deviceAddr, aclrtFree);                  \
    CHECK_RET(ret == ACL_SUCCESS, return ret)

#define CREATE_SPARSE_TENSOR(hostData, weightShape, storageShape, deviceAddr, dataType, tensor)          \
    ret = CreateSparseTensor(hostData, weightShape, storageShape, &deviceAddr, dataType, &tensor);       \
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> tensor##Ptr(tensor, aclDestroyTensor); \
    std::unique_ptr<void, aclError (*)(void*)> deviceAddr##Ptr(deviceAddr, aclrtFree);                  \
    CHECK_RET(ret == ACL_SUCCESS, return ret)

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: 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 CreateSparseTensor(
    const T* sparseWeightData, const std::vector<int64_t>& viewShape, const std::vector<int64_t>& storageShape,
    void** deviceAddr, aclDataType dataType, aclTensor** tensor)
{
    auto size = static_cast<uint64_t>(GetShapeSize(storageShape)) * 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);
    // Call aclrtMemcpy to copy host data to device memory.
    ret = aclrtMemcpy(*deviceAddr, size, sparseWeightData, size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);

    // Calculate strides of the contiguous tensor.
    std::vector<int64_t> strides(viewShape.size(), 1);
    for (int64_t i = viewShape.size() - 2; i >= 0; i--) {
        strides[i] = viewShape[i + 1] * strides[i + 1];
    }

    // Call aclCreateTensor to create an aclTensor.
    *tensor = aclCreateTensor(
        viewShape.data(), viewShape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, storageShape.data(),
        storageShape.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 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);
    // Call aclrtMemcpy to copy host data to device memory.
    if (hostData.size() > 0) {
        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 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 aclCreateTensor 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();
}

void GenRandomMask(std::vector<size_t>& masks)
{
    masks[0] = random() % 4;
    masks[1] = random() % 4;
    while (masks[1] == masks[0]) {
        masks[1] = random() % 4;
    }
}

void GenRandomSparseData(std::vector<int8_t>& weightHostData)
{
    srandom(233U);
    std::vector<size_t> masks(2, 0UL);
    constexpr size_t step = 4UL;
    for (size_t i = 0; i < weightHostData.size(); i += step) {
        GenRandomMask(masks);
        for (auto mask : masks) {
            weightHostData[i + mask] = 0;
        }
    }
}

std::vector<int64_t> GenStorageShape(int64_t* dims, uint64_t dimsNum)
{
    std::vector<int64_t> storageShape;
    for (uint64_t i = 0UL; i < dimsNum; i++) {
        storageShape.push_back(dims[i]);
    }
    return storageShape;
}

int aclnnSparse4to2QuantMatmulTest(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 inputs and outputs. Customize based on the API.
    int64_t m = 64L;
    int64_t k = 512L;
    int64_t n = 128L;
    std::vector<int64_t> xShape = {m, k};
    std::vector<int64_t> weightShape = {n, k};
    std::vector<int64_t> indexShape = {n, (k + 7) / 8};
    std::vector<int64_t> biasShape = {n};
    std::vector<int64_t> xScaleShape = {m};
    std::vector<int64_t> weightScaleShape = {n};
    std::vector<int64_t> outShape = {m, n};
    void* xDeviceAddr = nullptr;
    void* sparseWeightDeviceAddr = nullptr;
    void* indexDeviceAddr = nullptr;
    void* biasDeviceAddr = nullptr;
    void* xScaleDeviceAddr = nullptr;
    void* weightScaleDeviceAddr = nullptr;
    void* outDeviceAddr = nullptr;
    aclTensor* x = nullptr;
    aclTensor* sparseWeight = nullptr;
    aclTensor* index = nullptr;
    aclTensor* bias = nullptr;
    aclTensor* xScale = nullptr;
    aclTensor* weightScale = nullptr;
    aclTensor* out = nullptr;
    std::vector<int8_t> xHostData(GetShapeSize(xShape), 1);
    std::vector<int8_t> weightHostData(GetShapeSize(weightShape), 1);
    std::vector<uint16_t> biasHostData(GetShapeSize(biasShape), 1); // Actually in bfloat16 half-precision format.
    std::vector<float> xScaleHostData(GetShapeSize(xScaleShape), 1);
    std::vector<float> weightScaleHostData(GetShapeSize(weightScaleShape), 1);
    GenRandomSparseData(weightHostData);

    int8_t* sparseWeightHostData = nullptr;
    uint8_t* indexHostData = nullptr;
    int64_t* sparseWeightDims = nullptr;
    uint64_t sparseWeightDimsNum = 0UL;
    int64_t* indexDims = nullptr;
    uint64_t indexDimsNum = 0UL;
    aclIntArray* weightShapeArray = aclCreateIntArray(weightShape.data(), weightShape.size());
    ret = aclnnTransSparse4to2Para(
        weightHostData.data(), weightShapeArray, &sparseWeightHostData, &sparseWeightDims, &sparseWeightDimsNum,
        &indexHostData, &indexDims, &indexDimsNum);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTransSparse4to2Para failed. ERROR: %d\n", ret); return ret);
    std::unique_ptr<int8_t[]> sparseWeightHostDataPtr(sparseWeightHostData);
    std::unique_ptr<uint8_t[]> indexHostDataPtr(indexHostData);
    std::unique_ptr<int64_t[]> sparseWeightDimsPtr(sparseWeightDims);
    std::unique_ptr<int64_t[]> indexDimsPtr(indexDims);

    CREATE_TENSOR(xHostData, xShape, xDeviceAddr, aclDataType::ACL_INT8, x);

    weightShape.back() = (weightShape.back() + 7) / 8 * 8 / 2; // After 4-to-2 sparsification, K dimension is aligned to 8 and halved.
    auto sparseWeightStorageShape = GenStorageShape(sparseWeightDims, sparseWeightDimsNum);
    CREATE_SPARSE_TENSOR(
        sparseWeightHostData, weightShape, sparseWeightStorageShape, sparseWeightDeviceAddr, aclDataType::ACL_INT8,
        sparseWeight);

    auto indexStorageShape = GenStorageShape(indexDims, indexDimsNum);
    CREATE_SPARSE_TENSOR(indexHostData, indexShape, indexStorageShape, indexDeviceAddr, aclDataType::ACL_UINT8, index);
    CREATE_TENSOR(biasHostData, biasShape, biasDeviceAddr, aclDataType::ACL_BF16, bias);
    CREATE_TENSOR(xScaleHostData, xScaleShape, xScaleDeviceAddr, aclDataType::ACL_FLOAT, xScale);
    CREATE_TENSOR(weightScaleHostData, weightScaleShape, weightScaleDeviceAddr, aclDataType::ACL_FLOAT, weightScale);
    CREATE_TENSOR(std::vector<uint16_t>(), outShape, outDeviceAddr, aclDataType::ACL_BF16, out);

    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    void* workspaceAddr = nullptr;

    // Call the first-phase API of aclnnSparse4to2QuantMatmul.
    ret = aclnnSparse4to2QuantMatmulWeightNzGetWorkspaceSize(
        x, sparseWeight, index, xScale, weightScale, bias, out, &workspaceSize, &executor);

    CHECK_RET(
        ret == ACL_SUCCESS, LOG_PRINT("aclnnSparse4to2QuantMatmulWeightNzGetWorkspaceSize failed. ERROR: %d\n", ret);
        return ret);
    // Allocate device memory based on workspaceSize computed by the first-phase API.
    std::unique_ptr<void, aclError (*)(void*)> workspaceAddrPtrTrans(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);
        workspaceAddrPtrTrans.reset(workspaceAddr);
    }
    // Call the second-phase API of aclnnSparse4to2QuantMatmul.
    ret = aclnnSparse4to2QuantMatmulWeightNz(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSparse4to2QuantMatmulWeightNz failed. ERROR: %d\n", ret); return ret);

    // 4. (Boilerplate) Synchronize and wait for task completion.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // 5. Retrieve the output value. Copy the result from device memory to host memory. Modify based on the specific API definition.
    auto size = GetShapeSize(outShape);
    // bf16 data cannot be printed directly in C. Read it as uint16 and convert to bf16 via binary.
    std::vector<uint16_t> resultData(size, 0);
    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);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("result[%ld] is: %u\n", i, resultData[i]);
    }
    return ACL_SUCCESS;
}

int main()
{
    // 1. (Boilerplate) Initialize device/stream. Refer to the ACL API manual.
    // Fill in deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = aclnnSparse4to2QuantMatmulTest(deviceId, stream);
    CHECK_FREE_RET(
        ret == ACL_SUCCESS, LOG_PRINT("aclnnSparse4to2QuantMatmulTest failed. ERROR: %d\n", ret); return ret);

    Finalize(deviceId, stream);
    return 0;
}