aclnnTransSparse4to2Para

Applicable Products

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

Function

Operator function:

Performs compression preprocessing on a structured sparse weight matrix, and outputs the compressed sparse matrix and the corresponding index matrix. Assuming that every 4 elements in the original sparse matrix contains at least 2 zeros, the compressed sparse matrix is a matrix with 2 zeros filtered out from every 4 elements. During matrix compression, an index matrix is generated. For every 4 elements in the original matrix, 2 2-bit indices are generated in the index matrix and encoded according to the rules.

Function Prototype

aclnnStatus aclnnTransSparse4to2Para(
    const int8_t* weight,
    aclIntArray*  shape,
    int8_t**      sparseWeight,
    int64_t**     sparseWeightDims,
    uint64_t*     sparseWeightDimsNum,
    uint8_t**     index,
    int64_t**     indexDims,
    uint64_t*     indexDimsNum)

aclnnTransSparse4to2Para

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-Contiguous Tensor
weight	Input	Uncompressed sparse right matrix in the matrix multiplication operation.	The shape can be expressed as (n, k). Every four elements contain at least two zeros.	INT8	ND	2	×
shape	Input	Shape of the uncompressed sparse right matrix in the matrix multiplication operation.	Created by calling the aclCreateIntArray API.	-	-	-	-
sparseWeight	Output	Specifies the compressed right matrix in the matrix multiplication operation.	The memory is released by the caller.	-	-	-	-
sparseWeightDims	Output	Pointer to the first address of the StorageShape array of the compressed right matrix in the matrix multiplication operation.	The memory is released by the caller. StorageShape can be expressed as (ceil(k_half / 32), ceil(n / 16), 16, 32), where k_half = ceil(k / 8) * 8 / 2.	-	-	-	-
sparseWeightDimsNum	Output	Dimension of the StorageShape array of the compressed right matrix in the matrix multiplication operation.	-	-	-	-	-
index	Output	Index matrix corresponding to the compressed right matrix in the matrix multiplication operation.	The memory is released by the caller.	-	-	-	-
indexDims	Output	Pointer to the first address of the index matrix StorageShape array corresponding to the compressed right matrix in the matrix multiplication operation.	The memory is released by the caller. StorageShape can be expressed as (ceil(k_half / 32), ceil(n / 16), 16, 8), where k_half = ceil(k / 8) * 8 / 2.	-	-	-	-
indexDimsNum	Output	Dimension of the index matrix StorageShape array corresponding to the compressed right matrix in the matrix multiplication operation.	-	-	-	-	-

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 input weight is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The shape, format, or data type of weight does not meet the requirements, or the data does not satisfy the condition that at least two out of every four elements are zero.

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: 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 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 a 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 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 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 as required by 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 selection, K-axis is aligned to 8 and then 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 the 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. Obtain output values by copying results from device memory to host memory. Modify as needed based on the specific API definition.
    auto size = GetShapeSize(outShape);
    // bf16 data cannot be printed directly in C. Read 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.
    // Set 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;
}