aclnnWeightQuantPreprocess

📄 查看源码

产品支持情况

• Ascend 950PR/Ascend 950DT：支持

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：不支持

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：不支持

• Atlas 200I/500 A2 推理产品：不支持

• Atlas 推理系列产品：不支持

• Atlas 训练系列产品：不支持

功能说明

完成伪量化Matmul（包括QuantBatchMatmulV5、GroupedMatmul-伪量化、WeightQuantBatchMatmulV2）的参数预处理：主要将weight从ND格式转换为FRACTAL_NZ格式（MM_A16S4数据流的pertensor及perchannel/pergroup转置场景为ND直拷透传），并在需要时对weightScale、weightOffsetOptional、biasOptional进行同步处理。

函数原型

每个算子分为两段式接口：

1. 调用aclnnWeightQuantPreprocessGetWorkspaceSize获取workspace大小及执行器；
2. 调用aclnnWeightQuantPreprocess执行计算。

注意：用户需自行构造输出张量，参考约束说明中的shape计算公式。

aclnnStatus aclnnWeightQuantPreprocessGetWorkspaceSize(
    const aclTensor  *weight,
    const aclTensor  *weightScale,
    const aclTensor  *weightOffsetOptional,
    const aclTensor  *biasOptional,
    const aclDataType xDtype,
    const aclDataType xScaleDtype,
    const int64_t     kGroupSize,
    aclTensor        *outWeight,
    aclTensor        *outWeightScale,
    aclTensor        *outWeightOffsetOptional,
    aclTensor        *outBiasOptional,
    uint64_t         *workspaceSize,
    aclOpExecutor   **executor)

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

aclnnWeightQuantPreprocessGetWorkspaceSize

• 参数说明

  参数名	输入/输出	描述	使用说明	数据类型	数据格式	维度（shape）	非连续Tensor
  weight(aclTensor *)	输入	Matmul的权重矩阵	不支持空tensor	float4_e2m1/int4	ND	2-3	仅转置场景支持
  weightScale(aclTensor *)	输入	权重的反量化scale参数	不支持空tensor	float8_e8m0/float16/bfloat16	ND/NCL/NCHW	1-4	仅转置场景支持
  weightOffsetOptional(aclTensor *)	可选输入	权重的反量化offset参数	MM_A16S4数据流支持透传（要求与weightScale同viewShape、同数据类型）；其余数据流不支持，必须为nullptr	float16/bfloat16	ND	1-2	仅转置场景支持
  biasOptional(aclTensor *)	可选输入	Matmul的偏置矩阵	不支持空tensor，必须contiguous	float16/bfloat16	ND	1-2	不支持
  xDtype(aclDataType)	输入	Matmul的激活矩阵的数据类型	-	aclDataType	-	-	-
  xScaleDtype(aclDataType)	输入	激活的量化scale参数的数据类型	-	aclDataType	-	-	-
  kGroupSize(int64_t)	输入	权重在pergroup量化时K维度的group的大小	-	int64	-	-	-
  outWeight(aclTensor *)	输出	预处理后的weight	-	int8/int4/fp8_e4m3/hif8/fp4_e2m1	ND/NZ	2-5	仅转置场景支持
  outWeightScale(aclTensor *)	输出	预处理后的weightScale	-	float16/bfloat16/fp8_e8m0	ND/NCL/NCHW	1-4	仅转置场景支持
  outWeightOffsetOptional(aclTensor *)	输出	预处理后的weightOffset	weightOffsetOptional非空时不可为nullptr，且要求与weightOffsetOptional同viewShape、同storageShape、同数据类型；其余场景必须为nullptr	float16/bfloat16	ND	1-2	仅转置场景支持
  outBiasOptional(aclTensor *)	输出	预处理后的bias	必须contiguous	float16/bfloat16	ND	1-2	不支持
  workspaceSize(uint64_t *)	输出	计算所需的workspace大小	-	uint64*	-	-	-
  executor(aclOpExecutor **)	输出	包含算子计算流程的执行器	-	aclOpExecutor**	-	-	-

• 返回值

  aclnnStatus：返回状态码，具体参见aclnn返回码。

  第一段接口完成入参校验，出现以下场景时报错：

  返回值	错误码	描述
  ACLNN_ERR_PARAM_NULLPTR	161001	weight、weightScale、outWeight或outWeightScale是空指针；或biasOptional非空但outBiasOptional是空指针。
  ACLNN_ERR_PARAM_INVALID	161002	输入的数据类型组合不支持，无法匹配当前支持的任何数据流（MM_MX_A8W4/GMM_MX_A8W4/MM_A16S4）。
  		weight、weightScale、outWeight或outWeightScale是空tensor；或biasOptional/outBiasOptional在提供时为空tensor。
  		weight、weightScale、biasOptional、outWeight、outWeightScale或outBiasOptional的数据类型和数据格式不在支持的范围之内。
  		weight、weightScale、biasOptional、outWeight、outWeightScale或outBiasOptional的shape或storage shape不满足校验条件。
  		weight或weightScale的stride不满足转置要求，或biasOptional/outBiasOptional在提供时不连续。
  		weightOffsetOptional或outWeightOffsetOptional非空（MM_MX_A8W4/GMM_MX_A8W4数据流；MM_A16S4数据流支持透传，要求与weightScale同viewShape、同数据类型），或kGroupSize不等于32（MM_MX_A8W4/GMM_MX_A8W4数据流）；MM_A16S4数据流还要求kGroupSize与weightScale的量化粒度匹配。
  ACLNN_ERR_RUNTIME_ERROR	361001	产品型号不支持。
  ACLNN_ERR_INNER_CREATE_EXECUTOR	561101	内部错误，执行器创建失败。
  ACLNN_ERR_INNER_NULLPTR	561103	workspaceSize或executor是空指针，或API内部构图接口返回空指针。

aclnnWeightQuantPreprocess

• 参数说明

  参数名	输入/输出	描述
  workspace	输入	在Device侧申请的workspace内存地址。
  workspaceSize	输入	在Device侧申请的workspace大小，由第一段接口aclnnWeightQuantPreprocessGetWorkspaceSize获取。
  executor	输入	op执行器，包含了算子计算流程。
  stream	输入	指定执行任务的Stream。

• 返回值

  aclnnStatus：返回状态码，具体参见aclnn返回码。

约束说明

• 确定性说明：aclnnWeightQuantPreprocess默认确定性实现。

MM_MX_A8W4数据流（MM表示Matmul；MX_A8W4表示x的数据类型为FLOAT8_E4M3FN，weight的数据类型为FLOAT4_E2M1，Mx量化模式）

• weight

  • 数据类型：FLOAT4_E2M1
  • 格式：ND
  • K % kGroupSize == 0
  • view shape：2-D {K, N}
  • storage shape：{N, K}（transposed）
  • stride：[1, K]（最后两维transposed）
  • 不支持空tensor

• weightScale

  • 数据类型：FLOAT8_E8M0
  • 格式：ND/NCL
  • view shape：3-D {ceildiv(K, 64), N, 2}
  • storage shape：{N, ceildiv(K, 64), 2}（transposed）
  • stride：[2, 2*ceildiv(K,64), 1]（维度0和1交换）
  • 不支持空tensor

• weightOffsetOptional

  • 当前不支持，必须为nullptr
  • outWeightOffsetOptional也必须为nullptr

• biasOptional

  • 数据类型：float16/bfloat16
  • 格式：ND
  • 必须为contiguous
  • 不支持空tensor（若提供）

• kGroupSize

  • 必须等于 32

• xDtype

  • FLOAT8_E4M3FN

• xScaleDtype

  • FLOAT8_E8M0

• outWeight

  • 数据类型：与weight相同
  • 格式：FRACTAL_NZ_C0_32
  • view shape：与weight view shape相同 {K, N}
  • storage shape：4-D {ceildiv(K, 32), ceildiv(N, 16), 16, 32}

• outWeightScale

  • 数据类型：与weightScale相同
  • 格式：ND
  • view shape：与weightScale view shape相同
  • storage shape：与weightScale storage shape相同

• outBiasOptional

  • 数据类型：与biasOptional相同
  • 格式：ND
  • 必须为contiguous
  • view shape：与biasOptional相同
  • storage shape：与biasOptional相同

GMM_MX_A8W4数据流（GMM表示GroupedMatmul；MX_A8W4表示x的数据类型为FLOAT8_E4M3FN，weight的数据类型为FLOAT4_E2M1，Mx量化模式）

• weight

  • 数据类型：FLOAT4_E2M1
  • 格式：ND
  • K % kGroupSize == 0
  • view shape：3-D {G, K, N}
  • storage shape：{G, N, K}（transposed，最后两维交换）
  • stride：[K*N, 1, K]（维度1和2 transposed）
  • 不支持空tensor

• weightScale

  • 数据类型：FLOAT8_E8M0
  • 格式：ND/NCL/NCHW
  • view shape：4-D {G, ceildiv(K, 64), N, 2}
  • storage shape：{G, N, ceildiv(K, 64), 2}（transposed，维度2和3交换）
  • stride：[2*ceildiv(K,64)*N, 2, 2*ceildiv(K,64), 1]（维度2和3交换）
  • 不支持空tensor

• weightOffsetOptional

  • 当前不支持，必须为nullptr
  • outWeightOffsetOptional也必须为nullptr

• biasOptional

  • 数据类型：float16/bfloat16
  • 格式：ND
  • 必须为contiguous
  • 不支持空tensor（若提供）

• kGroupSize

  • 必须等于 32

• xDtype

  • FLOAT8_E4M3FN

• xScaleDtype

  • FLOAT8_E8M0

• outWeight

  • 数据类型：与weight相同
  • 格式：FRACTAL_NZ_C0_32
  • view shape：与weight view shape相同 {G, K, N}
  • storage shape：5-D {G, ceildiv(K, 32), ceildiv(N, 16), 16, 32}

• outWeightScale

  • 数据类型：与weightScale相同
  • 格式：ND
  • view shape：与weightScale view shape相同
  • storage shape：与weightScale storage shape相同

• outBiasOptional

  • 数据类型：与biasOptional相同
  • 格式：ND
  • 必须为contiguous
  • view shape：与biasOptional相同
  • storage shape：与biasOptional相同

MM_A16S4数据流（MM表示Matmul；A16S4表示x的数据类型为FLOAT16/BFLOAT16，weight的数据类型为INT4（紧凑4-bit，每字节打包2个值），支持pertensor/perchannel/pergroup量化模式）

• weight

  • 数据类型：INT4
  • 格式：ND
  • view shape：2-D {K, N}（K为reduction维，N为输出维）
  • 支持连续或末两维严格转置（stride [1, K]）
  • 紧凑4-bit打包维（连续维）必须为偶数：非转置时N为偶数，转置时K为偶数
  • 不支持空tensor

• weightScale

  • 数据类型：FLOAT16/BFLOAT16
  • 格式：ND
  • view shape按量化粒度区分：
    • pertensor：单元素，{1}或{1, 1}
    • perchannel：{N}或{1, N}
    • pergroup：2-D {ceildiv(K, kGroupSize), N}，且分组数大于1；weight转置输入时，weightScale需与weight布局一致（同为转置视图）
  • 不支持空tensor

• weightOffsetOptional

  • 可选；若提供，数据类型与view shape必须与weightScale相同，格式为ND，不支持空tensor
  • outWeightOffsetOptional同步要求：与weightOffsetOptional同view shape、同storage shape、同数据类型

• biasOptional

  • 格式：ND
  • view shape：{N}或{1, N}
  • 必须为contiguous
  • 不支持空tensor（若提供）

• kGroupSize

  • pertensor/perchannel：必须等于 0
  • pergroup：必须大于 0，且weightScale第0维等于ceildiv(K, kGroupSize)

• xDtype

  • FLOAT16/BFLOAT16

• xScaleDtype

  • 必须为ACL_DT_UNDEFINED

• outWeight

  • 数据类型：与weight相同
  • 格式：pertensor（转置/非转置）与perchannel/pergroup转置场景为ND直拷（物理透传）；perchannel/pergroup非转置场景为FRACTAL_NZ_C0_16（ND→NZ分形转换）
  • view shape：与weight view shape相同 {K, N}
  • storage shape：FRACTAL_NZ_C0_16时为4-D {ceildiv(N, 16), ceildiv(K, 16), 16, 16}（N块在前，区别于A8W4 NZ_C0_32的K块在前）；ND直拷场景与输入weight保持相同的view shape/strides（属纯透传）

• outWeightScale

  • 数据类型：与weightScale相同
  • 格式：ND
  • view shape：与weightScale view shape相同
  • storage shape：与weightScale storage shape相同

• outBiasOptional

  • 数据类型：与biasOptional相同
  • 格式：ND
  • 必须为contiguous
  • view shape：与biasOptional相同
  • storage shape：与biasOptional相同

• 除上述数据流外的数据类型与shape组合为预留接口，当前调用将返回ACLNN_ERR_PARAM_INVALID

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。

注意：用户需自行计算并构造输出张量shape，参考约束说明中的公式。以下示例为MM_MX_A8W4数据流：

• outWeight viewShape：与weight viewShape相同
• outWeight storageShape：{CeilDiv(K, 32), CeilDiv(N, 16), 16, 32}
• outWeight format：ACL_FORMAT_FRACTAL_NZ_C0_32

#include <iostream>
#include <memory>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_weight_quant_preprocess.h"

#define CHECK_RET(cond, return_expr) \
    do {                             \
        if (!(cond)) {               \
            return_expr;             \
        }                            \
    } while (0)

#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))

int64_t GetShapeSize(const std::vector<int64_t>& shape)
{
    int64_t size = 1;
    for (auto d : shape)
        size *= d;
    return size;
}

class AclRuntimeGuard {
public:
    explicit AclRuntimeGuard(int32_t deviceId) : deviceId_(deviceId) {}

    ~AclRuntimeGuard()
    {
        if (stream_ != nullptr) {
            aclrtDestroyStream(stream_);
            stream_ = nullptr;
        }
        if (deviceSet_) {
            aclrtResetDevice(deviceId_);
            deviceSet_ = false;
        }
        if (aclInited_) {
            aclFinalize();
            aclInited_ = false;
        }
    }

    int Init(aclrtStream* stream)
    {
        auto ret = aclInit(nullptr);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        aclInited_ = true;

        ret = aclrtSetDevice(deviceId_);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        deviceSet_ = true;

        ret = aclrtCreateStream(stream);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        stream_ = *stream;
        return ACL_SUCCESS;
    }

private:
    int32_t deviceId_;
    aclrtStream stream_ = nullptr;
    bool aclInited_ = false;
    bool deviceSet_ = false;
};

int main()
{
    int32_t deviceId = 0;
    aclrtStream stream = nullptr;
    AclRuntimeGuard aclGuard(deviceId);
    auto ret = aclGuard.Init(&stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Init failed" << std::endl; return ret);

    // weight: FLOAT4_E2M1, transposed (MM_MX_A8W4)
    int64_t k = 64;
    int64_t n = 128;
    int64_t C0 = 32; // FLOAT4_E2M1 对应 C0=32

    std::vector<int64_t> weightViewShape = {k, n};
    std::vector<int64_t> weightStorageShape = {n, k};
    std::vector<int64_t> weightStrides = {1, k};
    int64_t weightStorageSize = GetShapeSize(weightStorageShape);
    int64_t weightBytes = weightStorageSize / 2; // FP4: 4 bits = 0.5 bytes per element

    std::vector<int8_t> weightHostData(weightBytes, 0);
    void* weightDeviceAddr = nullptr;
    ret = aclrtMalloc(&weightDeviceAddr, weightBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc weight failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> weightDeviceAddrPtr(weightDeviceAddr, aclrtFree);
    ret = aclrtMemcpy(weightDeviceAddr, weightBytes, weightHostData.data(), weightBytes, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Memcpy weight failed" << std::endl; return ret);
    aclTensor* weight = aclCreateTensor(weightViewShape.data(), weightViewShape.size(), ACL_FLOAT4_E2M1,
                                        weightStrides.data(), 0, ACL_FORMAT_ND, weightStorageShape.data(),
                                        weightStorageShape.size(), weightDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> weightPtr(weight, aclDestroyTensor);
    CHECK_RET(weight != nullptr, std::cout << "Create weight tensor failed" << std::endl; return ACL_ERROR_FAILURE);

    // weightScale: FLOAT8_E8M0, 3-D transposed (MM_MX_A8W4)
    // viewShape: {ceildiv(K,64), N, 2} = {1, 128, 2}
    // storageShape: {N, ceildiv(K,64), 2} = {128, 1, 2}
    // transposed stride: {2, 2, 1} (dim0 <-> dim1)
    std::vector<int64_t> scaleViewShape = {k / 64, n, 2};
    std::vector<int64_t> scaleStorageShape = {n, k / 64, 2};
    std::vector<int64_t> scaleStrides = {2, 2, 1};
    int64_t scaleStorageSize = GetShapeSize(scaleStorageShape);
    int64_t scaleBytes = scaleStorageSize; // FP8: 1 byte per element

    std::vector<int8_t> scaleHostData(scaleBytes, 0);
    void* scaleDeviceAddr = nullptr;
    ret = aclrtMalloc(&scaleDeviceAddr, scaleBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc weightScale failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> scaleDeviceAddrPtr(scaleDeviceAddr, aclrtFree);
    ret = aclrtMemcpy(scaleDeviceAddr, scaleBytes, scaleHostData.data(), scaleBytes, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Memcpy weightScale failed" << std::endl; return ret);
    aclTensor* weightScale = aclCreateTensor(scaleViewShape.data(), scaleViewShape.size(), ACL_FLOAT8_E8M0,
                                             scaleStrides.data(), 0, ACL_FORMAT_ND, scaleStorageShape.data(),
                                             scaleStorageShape.size(), scaleDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> weightScalePtr(weightScale, aclDestroyTensor);
    CHECK_RET(weightScale != nullptr, std::cout << "Create weightScale tensor failed" << std::endl;
              return ACL_ERROR_FAILURE);

    // 用户自行构造 outWeight (FRACTAL_NZ_C0_32)
    // viewShape 与 weight viewShape 相同，storageShape 按公式计算
    std::vector<int64_t> outWeightViewShape = {k, n};
    std::vector<int64_t> outWeightStorageShape = {CEIL_DIV(k, C0), CEIL_DIV(n, 16), 16, C0};
    int64_t outWeightStorageSize = GetShapeSize(outWeightStorageShape);
    int64_t outWeightBytes = outWeightStorageSize / 2; // FP4

    void* outWeightDeviceAddr = nullptr;
    ret = aclrtMalloc(&outWeightDeviceAddr, outWeightBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc outWeight failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> outWeightDeviceAddrPtr(outWeightDeviceAddr, aclrtFree);
    aclTensor* outWeight = aclCreateTensor(outWeightViewShape.data(), outWeightViewShape.size(), ACL_FLOAT4_E2M1,
                                           nullptr, 0, ACL_FORMAT_FRACTAL_NZ_C0_32, outWeightStorageShape.data(),
                                           outWeightStorageShape.size(), outWeightDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outWeightPtr(outWeight, aclDestroyTensor);
    CHECK_RET(outWeight != nullptr, std::cout << "Create outWeight tensor failed" << std::endl;
              return ACL_ERROR_FAILURE);

    // 构造 outWeightScale (viewShape 和 storageShape 都与 weightScale 相同)
    // 根据实现要求：outWeightScale 的 viewShape 和 storageShape 必须都与 weightScale 相同
    void* outScaleDeviceAddr = nullptr;
    ret = aclrtMalloc(&outScaleDeviceAddr, scaleBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc outWeightScale failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> outScaleDeviceAddrPtr(outScaleDeviceAddr, aclrtFree);
    aclTensor* outWeightScale = aclCreateTensor(scaleViewShape.data(), scaleViewShape.size(), ACL_FLOAT8_E8M0,
                                                scaleStrides.data(), 0, ACL_FORMAT_ND, scaleStorageShape.data(),
                                                scaleStorageShape.size(), outScaleDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outWeightScalePtr(outWeightScale, aclDestroyTensor);
    CHECK_RET(outWeightScale != nullptr, std::cout << "Create outWeightScale tensor failed" << std::endl;
              return ACL_ERROR_FAILURE);

    aclDataType xDtype = ACL_FLOAT8_E4M3FN;
    aclDataType xScaleDtype = ACL_FLOAT8_E8M0;
    int64_t kGroupSize = 32;

    // 1. 获取 workspace 与执行器
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    ret = aclnnWeightQuantPreprocessGetWorkspaceSize(
        weight, weightScale, nullptr, nullptr, // weightOffsetOptional, biasOptional
        xDtype, xScaleDtype, kGroupSize, outWeight, outWeightScale, nullptr, nullptr, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "GetWorkspaceSize failed" << std::endl; return ret);

    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, std::cout << "Malloc workspace failed" << std::endl; return ret);
        workspaceAddrPtr.reset(workspaceAddr);
    }

    // 2. 执行计算
    ret = aclnnWeightQuantPreprocess(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Preprocess failed" << std::endl; return ret);

    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Synchronize failed" << std::endl; return ret);

    // 3. 释放资源
    workspaceAddrPtr.reset();
    outWeightScalePtr.reset();
    outWeightPtr.reset();
    weightScalePtr.reset();
    weightPtr.reset();
    outScaleDeviceAddrPtr.reset();
    outWeightDeviceAddrPtr.reset();
    scaleDeviceAddrPtr.reset();
    weightDeviceAddrPtr.reset();
    return 0;
}

MM_A16S4数据流调用示例如下（x为FLOAT16，weight为INT4紧凑排布；包含perchannel非转置NZ_C0_16输出与pertensor ND直拷两个场景）：

• outWeight viewShape：与weight viewShape相同
• outWeight storageShape：FRACTAL_NZ_C0_16输出场景为{CeilDiv(N, 16), CeilDiv(K, 16), 16, 16}（N块在前，区别于MM_MX_A8W4的{CeilDiv(K, 32), CeilDiv(N, 16), 16, 32}）；ND直拷场景与输入weight保持相同的view shape/strides
• outWeight format：perchannel/pergroup非转置场景为ACL_FORMAT_FRACTAL_NZ_C0_16，pertensor及perchannel/pergroup转置场景为ACL_FORMAT_ND

#include <iostream>
#include <memory>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_weight_quant_preprocess.h"

#define CHECK_RET(cond, return_expr) \
    do {                             \
        if (!(cond)) {               \
            return_expr;             \
        }                            \
    } while (0)

#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))

int64_t GetShapeSize(const std::vector<int64_t>& shape)
{
    int64_t size = 1;
    for (auto d : shape)
        size *= d;
    return size;
}

class AclRuntimeGuard {
public:
    explicit AclRuntimeGuard(int32_t deviceId) : deviceId_(deviceId) {}

    ~AclRuntimeGuard()
    {
        if (stream_ != nullptr) {
            aclrtDestroyStream(stream_);
            stream_ = nullptr;
        }
        if (deviceSet_) {
            aclrtResetDevice(deviceId_);
            deviceSet_ = false;
        }
        if (aclInited_) {
            aclFinalize();
            aclInited_ = false;
        }
    }

    int Init(aclrtStream* stream)
    {
        auto ret = aclInit(nullptr);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        aclInited_ = true;

        ret = aclrtSetDevice(deviceId_);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        deviceSet_ = true;

        ret = aclrtCreateStream(stream);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        stream_ = *stream;
        return ACL_SUCCESS;
    }

private:
    int32_t deviceId_;
    aclrtStream stream_ = nullptr;
    bool aclInited_ = false;
    bool deviceSet_ = false;
};

// 设备侧 Tensor 持有者：析构时按创建逆序自动释放 tensor 与 device 内存
class DeviceTensor {
public:
    // bytesPerElem=0.5 表示 4-bit 紧凑打包（INT4/FP4，每字节 2 个值），内存按 numel/2 字节申请
    DeviceTensor(const std::vector<int64_t>& viewShape, const std::vector<int64_t>& storageShape,
                 const std::vector<int64_t>& strides, aclDataType dtype, aclFormat format, double bytesPerElem)
    {
        int64_t storageSize = GetShapeSize(storageShape);
        bytes_ = static_cast<int64_t>(storageSize * bytesPerElem);

        std::vector<int8_t> hostData(bytes_, 0);
        auto ret = aclrtMalloc(&deviceAddr_, bytes_, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc device memory failed" << std::endl; return);
        ret = aclrtMemcpy(deviceAddr_, bytes_, hostData.data(), bytes_, ACL_MEMCPY_HOST_TO_DEVICE);
        CHECK_RET(ret == ACL_SUCCESS, std::cout << "Memcpy H2D failed" << std::endl; return);
        const int64_t* stridesData = strides.empty() ? nullptr : strides.data();
        tensor_ = aclCreateTensor(viewShape.data(), viewShape.size(), dtype, stridesData, 0, format,
                                  storageShape.data(), storageShape.size(), deviceAddr_);
        CHECK_RET(tensor_ != nullptr, std::cout << "Create tensor failed" << std::endl; return);
    }

    ~DeviceTensor()
    {
        if (tensor_ != nullptr) {
            aclDestroyTensor(tensor_);
            tensor_ = nullptr;
        }
        if (deviceAddr_ != nullptr) {
            aclrtFree(deviceAddr_);
            deviceAddr_ = nullptr;
        }
    }

    aclTensor* Get() const { return tensor_; }
    bool IsValid() const { return tensor_ != nullptr && deviceAddr_ != nullptr; }

private:
    aclTensor* tensor_ = nullptr;
    void* deviceAddr_ = nullptr;
    int64_t bytes_ = 0;
};

int RunPreprocess(aclTensor* weight, aclTensor* weightScale, aclTensor* weightOffset, aclDataType xDtype,
                  int64_t kGroupSize, aclTensor* outWeight, aclTensor* outWeightScale, aclTensor* outWeightOffset,
                  aclrtStream stream)
{
    // A16W4 数据流无 xScale 语义，xScaleDtype 固定传 ACL_DT_UNDEFINED
    aclDataType xScaleDtype = ACL_DT_UNDEFINED;

    // 1. 获取 workspace 与执行器
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    auto ret = aclnnWeightQuantPreprocessGetWorkspaceSize(weight, weightScale, weightOffset, nullptr, // biasOptional
                                                          xDtype, xScaleDtype, kGroupSize, outWeight, outWeightScale,
                                                          outWeightOffset, nullptr, // outBiasOptional
                                                          &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "GetWorkspaceSize failed, ret=" << ret << std::endl; return ret);

    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, std::cout << "Malloc workspace failed" << std::endl; return ret);
        workspaceAddrPtr.reset(workspaceAddr);
    }

    // 2. 执行计算
    ret = aclnnWeightQuantPreprocess(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Preprocess failed, ret=" << ret << std::endl; return ret);

    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Synchronize failed" << std::endl; return ret);
    return ACL_SUCCESS;
}

// 场景一：MM_A16S4_PERCHANNEL（INT4 perchannel，非转置 weight -> FRACTAL_NZ_C0_16 输出，offset 直拷透传）
int TestA16S4PerChannelNz(aclrtStream stream)
{
    int64_t k = 256;
    int64_t n = 256;

    // weight: INT4，非转置连续 [K, N]；紧凑 4-bit 打包维（N）须为偶数
    DeviceTensor weight({k, n}, {k, n}, {n, 1}, ACL_INT4, ACL_FORMAT_ND, 0.5);
    CHECK_RET(weight.IsValid(), return ACL_ERROR_FAILURE);

    // weightScale: perchannel {N}，FP16；perchannel 场景 kGroupSize 必须为 0
    DeviceTensor weightScale({n}, {n}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(weightScale.IsValid(), return ACL_ERROR_FAILURE);

    // weightOffsetOptional: A16S4 各数据流支持透传，要求与 weightScale 同形同 dtype
    DeviceTensor weightOffset({n}, {n}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(weightOffset.IsValid(), return ACL_ERROR_FAILURE);

    // outWeight: 用户自行构造，viewShape 与 weight 相同；NZ_C0_16 的 storageShape 为
    // {ceildiv(N, 16), ceildiv(K, 16), 16, 16}（N 块在前，区别于 A8W4 NZ_C0_32 的 K 块在前）
    std::vector<int64_t> outStorageShape = {CEIL_DIV(n, 16), CEIL_DIV(k, 16), 16, 16};
    DeviceTensor outWeight({k, n}, outStorageShape, {}, ACL_INT4, ACL_FORMAT_FRACTAL_NZ_C0_16, 0.5);
    CHECK_RET(outWeight.IsValid(), return ACL_ERROR_FAILURE);

    // outWeightScale/outWeightOffset: viewShape、storageShape、dtype 均与输入相同（直拷）
    DeviceTensor outWeightScale({n}, {n}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(outWeightScale.IsValid(), return ACL_ERROR_FAILURE);
    DeviceTensor outWeightOffset({n}, {n}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(outWeightOffset.IsValid(), return ACL_ERROR_FAILURE);

    auto ret = RunPreprocess(weight.Get(), weightScale.Get(), weightOffset.Get(), ACL_FLOAT16, 0, outWeight.Get(),
                             outWeightScale.Get(), outWeightOffset.Get(), stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "TestA16S4PerChannelNz failed" << std::endl; return ret);
    std::cout << "TestA16S4PerChannelNz success" << std::endl;
    return ACL_SUCCESS;
}

// 场景二：MM_A16S4_PERTENSOR（INT4 pertensor，一律 ND 直拷，与转置无关）
int TestA16S4PerTensorNd(aclrtStream stream)
{
    int64_t k = 256;
    int64_t n = 256;

    // weight: INT4，连续 [K, N]（pertensor 对转置/非转置均按字节直拷）
    DeviceTensor weight({k, n}, {k, n}, {n, 1}, ACL_INT4, ACL_FORMAT_ND, 0.5);
    CHECK_RET(weight.IsValid(), return ACL_ERROR_FAILURE);

    // weightScale: pertensor 单元素 {1}，FP16；kGroupSize 必须为 0
    DeviceTensor weightScale({1}, {1}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(weightScale.IsValid(), return ACL_ERROR_FAILURE);
    DeviceTensor weightOffset({1}, {1}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(weightOffset.IsValid(), return ACL_ERROR_FAILURE);

    // outWeight: ND 直拷，format 为 ND，storageShape 与 viewShape 相同
    DeviceTensor outWeight({k, n}, {k, n}, {}, ACL_INT4, ACL_FORMAT_ND, 0.5);
    CHECK_RET(outWeight.IsValid(), return ACL_ERROR_FAILURE);

    DeviceTensor outWeightScale({1}, {1}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(outWeightScale.IsValid(), return ACL_ERROR_FAILURE);
    DeviceTensor outWeightOffset({1}, {1}, {1}, ACL_FLOAT16, ACL_FORMAT_ND, 2);
    CHECK_RET(outWeightOffset.IsValid(), return ACL_ERROR_FAILURE);

    auto ret = RunPreprocess(weight.Get(), weightScale.Get(), weightOffset.Get(), ACL_FLOAT16, 0, outWeight.Get(),
                             outWeightScale.Get(), outWeightOffset.Get(), stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "TestA16S4PerTensorNd failed" << std::endl; return ret);
    std::cout << "TestA16S4PerTensorNd success" << std::endl;
    return ACL_SUCCESS;
}

// WeightQuantPreprocess currently supports Ascend 950 only.
int main()
{
    int32_t deviceId = 0;
    aclrtStream stream = nullptr;
    AclRuntimeGuard aclGuard(deviceId);
    auto ret = aclGuard.Init(&stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Init failed" << std::endl; return ret);

    ret = TestA16S4PerChannelNz(stream);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = TestA16S4PerTensorNd(stream);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    std::cout << "All A16S4 examples run success" << std::endl;
    return 0;
}