aclnnConvertWeightToINT4Pack

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	√
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function

Preprocesses the input weight data to convert low-bit data from sparse storage to dense storage layout. When the data format of the output weightInt4Pack is declared as FRACTAL_NZ, this operator converts the data format from ND to FRACTAL_NZ.

Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products: Packs INT32 type weight input data into densely arranged INT4 data.
Ascend 950PR/Ascend 950DT: Packs INT32-type weight into densely arranged INT4 type, and packs FLOAT-type weight into densely arranged FLOAT4_E2M1 type.

Function Prototype

Each operator is divided into a two-phase API. You must first call the "aclnnConvertWeightToINT4PackGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation flow, and then call the "aclnnConvertWeightToINT4Pack" API to execute the computation.

aclnnStatus aclnnConvertWeightToINT4PackGetWorkspaceSize(
  const aclTensor *weight,
  aclTensor       *weightInt4Pack,
  uint64_t        *workspaceSize,
  aclOpExecutor   **executor)

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

aclnnConvertWeightToINT4PackGetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-Contiguous Tensor
weight	Input	Low-bit quantized weight of Matmul-type operators, with 32-bit type carrying 4-bit weight values.	-	INT32	ND,FRACTAL_NZ	2-3	-
weightInt4Pack	Output	The low-bit quantized weight of Matmul-type operators, with weight values being 4-bit and densely arranged.	-	INT4,INT32	ND,FRACTAL_NZ	-	-
workspaceSize	Output	Returns the workspace size required on the Device side.	-	-	-	-	-
executor	Output	Returns the operator 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	A required input, output, or attribute passed in is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The shape dimensions of weight or weightInt4Pack passed in do not meet the requirements.
		The data types of weight and weightInt4Pack are not within the supported range.
		The shapes of weight and weightInt4Pack do not meet the constraint requirements.
		An empty tensor is passed in.
		The format of the input tensor is not ND.
ACLNN_ERR_RUNTIME_ERROR	361001	Exception when copying data from the host side to the device side.
ACLNN_ERR_RUNTIME_ERROR	361001	Exception when copying data from the device side to the host side.

aclnnConvertWeightToINT4Pack

Parameters

Parameter	Input/Output	Description
workspace	Input	Workspace memory address applied on the Device side.
workspaceSize	Input	Workspace size applied on the Device side, obtained from the first-phase API aclnnConvertWeightToINT4PackGetWorkspaceSize.
executor	Input	Operator executor that 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

Deterministic description:
- Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products, Ascend 950PR/Ascend 950DT: aclnnConvertWeightToINT4Pack defaults to a deterministic implementation.

The relationships between data types and data formats among parameters are as follows:

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

Weight Data Type	Weight Data Format	WeightInt4Pack Data Type	WeightInt4Pack Data Format	Weight Shape	WeightInt4Pack View Shape	weightInt4Pack storage shape
INT32 (carries INT4 type data, with a data representation range of [-8, 7]).	ND	INT4	ND	The last dimension is aligned to 2.	Consistent with the input weight, i.e., (dim0, dim1).	Same as the view shape.
INT32 (carrying INT4 type data, with a data representation range of [-8, 7]).	ND	INT4	FRACTAL_NZ	The last dimension is aligned to 2.	Consistent with the input weight, i.e., (dim0, dim1).	(⌈dim1/64⌉, ⌈dim0/16⌉, 16, 64)
INT32 (carries INT4 type data, with a data representation range of [-8, 7])	ND	INT32 (one INT32 data stores eight INT4 data)	ND	The last dimension is aligned to 8.	The last dimension is 1/8 of the last dimension of weight, i.e., (dim0, dim1/8).	Same as view shape.
INT32 (carrying INT4 type data, with a data representation range of [-8, 7]).	ND	INT32 (one INT32 data stores eight INT4 data).	FRACTAL_NZ	The last dimension is 8-aligned.	The last dimension is 1/8 of the last dimension of weight, i.e., (dim0, dim1/8).	(⌈dim1/64⌉, ⌈dim0/16⌉, 16, 8)

Ascend 950PR/Ascend 950DT:

Weight Data Type	Weight Data Format	weightInt4Pack Data Type	weightInt4Pack Data Format	Weight Shape	weightInt4Pack view shape	weightInt4Pack storage shape
INT32 (carries INT4 type data, data representation range [-8, 7]) or FLOAT (carries FLOAT4_E2M1 type data, data representation range [-6.0, 6.0])	ND	INT4 or FLOAT4_E2M1	ND	The last dimension is aligned to 2.	Consistent with the input weight, i.e., (dim0, dim1).	Same as the view shape.
INT32 (carries INT4 type data, with a data representation range of [-8, 7]) or FLOAT (carries FLOAT4_E2M1 type data, with a data representation range of [-6.0, 6.0]).	ND	INT4 or FLOAT4_E2M1	FRACTAL_NZ	The last dimension is aligned to 2.	Consistent with the input weight, i.e., (dim0, dim1).	(⌈dim1/16⌉, ⌈dim0/16⌉, 16, 16)
INT32 (carrying INT4 type data, with a data representation range of [-8, 7]) or FLOAT (carrying FLOAT4_E2M1 type data, with a data representation range of [-6.0, 6.0])	ND	INT32 (one INT32 datum stores eight INT4 data) or FLOAT (one FLOAT datum stores eight FLOAT4_E2M1 data)	ND	The last dimension is aligned to 8.	The last dimension is 1/8 of the last dimension of weight, i.e., (dim0, dim1/8).	Same as the view shape.
INT32 (carrying INT4 type data, with a data representation range of [-8, 7]) or FLOAT (carrying FLOAT4_E2M1 type data, with a data representation range of [-6.0, 6.0]).	ND	INT32 (one INT32 data stores eight INT4 data) or FLOAT (one FLOAT data stores eight FLOAT4_E2M1 data)	FRACTAL_NZ	The last dimension is 8-aligned.	The last dimension is 1/8 of the last dimension of weight, i.e., (dim0, dim1/8).	(⌈dim1/16⌉, ⌈dim0/16⌉, 16, 2)
INT32 (carries INT4 type data, with a data representation range of [-8, 7]).	ND	INT4 (one INT32 data stores eight INT4 data).	FRACTAL_NZ	The last dimension is aligned to 8.	The last dimension is 1/8 of the last dimension of weight, i.e., (dim0, dim1/8).	(⌈dim1/32⌉, ⌈dim0/16⌉, 16, 32)

Example

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

The following is sample code for reference only. For the specific compilation and execution process, see Compile and Run Samples. For pseudo-quantization, the aclnnWeightQuantBatchMatmulV2 and aclnnWeightQuantBatchMatmulV3 APIs are available. Here, aclnnWeightQuantBatchMatmulV2 is used as an example.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_cast.h"
#include "aclnnop/aclnn_weight_quant_batch_matmul_v2.h"

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

#define LOG_PRINT(message, ...)     \
  do {                              \
    printf(message, ##__VA_ARGS__); \
  } while (0)

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

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

extern "C" aclnnStatus aclnnConvertWeightToINT4PackGetWorkspaceSize(const aclTensor *weight, aclTensor *weightInt4Pack,
    uint64_t *workspaceSize, aclOpExecutor **executor);

extern "C" aclnnStatus aclnnConvertWeightToINT4Pack(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
    aclrtStream stream);

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 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 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);
  // Copy host-side data to device-side memory using 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);

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

  // Create an aclTensor using the aclCreateTensor API.
  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
  return 0;
}

template <typename T>
int CreateAclTensorInt4(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor, aclFormat format) {
  auto size = hostData.size() * sizeof(T);
  // Apply for device-side memory using 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 device-side memory using 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);

  // 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.
  if (format == aclFormat::ACL_FORMAT_ND) {
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
  } else {
    std::vector<int64_t> nzShape;
    if (dataType == aclDataType::ACL_INT4) {
        nzShape = {CEIL_DIV(shape[1], 64), CEIL_DIV(shape[0], 16), 16, 64};
    } else {
        nzShape = {CEIL_DIV(shape[1], 64), CEIL_DIV(shape[0], 16), 16, 8};
    }
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0,
                              aclFormat::ACL_FORMAT_FRACTAL_NZ, nzShape.data(), nzShape.size(), *deviceAddr);
  }

  return 0;
}

int main() {
  // 1. (Boilerplate) Initialize the device/stream. See the ACL API manual.
  // Fill in the deviceId based on your actual device.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
  aclDataType weightInt4PackDtype = aclDataType::ACL_INT4;
  aclFormat weightFormat = aclFormat::ACL_FORMAT_FRACTAL_NZ;
  bool isWeightTransposed = true;

  // 2. Construct the inputs and outputs. Custom construction is required based on the API definition.
  int64_t m = 16;
  int64_t k = 72;
  int64_t n = 17;
  int64_t weightDim0 = k;
  int64_t weightDim1 = n;
  if (isWeightTransposed) {
    weightDim0 = n;
    weightDim1 = k;
  }
  std::vector<int64_t> xShape = {m, k};
  std::vector<int64_t> weightShape = {weightDim0, weightDim1};
  std::vector<int64_t> weightInt4PackShape;
  if (weightInt4PackDtype == aclDataType::ACL_INT4) {
    weightInt4PackShape = {weightDim0, weightDim1};
  } else {
    weightInt4PackShape = {weightDim0, weightDim1/8};
  }
  std::vector<int64_t> yShape = {m, n};
  void* xDeviceAddr = nullptr;
  void* weightDeviceAddr = nullptr;
  void* weightInt4PackDeviceAddr = nullptr;
  void* yDeviceAddr = nullptr;
  aclTensor* x = nullptr;
  aclTensor* weight = nullptr;
  aclTensor* weightInt4Pack = nullptr;
  aclTensor* y = nullptr;
  std::vector<float> xHostData(m * k, 1);
  std::vector<int32_t> weightHostData(k * n, 1);
  std::vector<float> yHostData(m * n, 0);

  std::vector<int64_t> antiquantScaleShape = {n};
  void* antiquantScaleDeviceAddr = nullptr;
  aclTensor* antiquantScale = nullptr;
  std::vector<float> antiquantScaleHostData(n, 1);

  // Create x aclTensor.
  ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create weight aclTensor.
  ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_INT32, &weight);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  if (weightInt4PackDtype == aclDataType::ACL_INT4) {
    std::vector<int8_t> weightInt4PackHostData(n * k / 2, 0); //One int8 value stores two int4 values, so divide by 2 here.
    if (weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
      weightInt4PackHostData.resize(CEIL_ALIGN(weightDim1/2, 32) * CEIL_ALIGN(weightDim0, 16), 0);
    }
    // Create weightInt4Pack aclTensor.
    ret = CreateAclTensorInt4(weightInt4PackHostData, weightInt4PackShape, &weightInt4PackDeviceAddr,
                              weightInt4PackDtype, &weightInt4Pack, weightFormat);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  } else {
    std::vector<int32_t> weightInt4PackHostData(n * k / 8, 1); //One int32 value stores eight int4 values, so divide by 8 here.
    if (weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
      weightInt4PackHostData.resize(CEIL_ALIGN(weightDim1/8, 8) * CEIL_ALIGN(weightDim0, 16), 0);
      ret = CreateAclTensorInt4(weightInt4PackHostData, weightInt4PackShape, &weightInt4PackDeviceAddr,
                                weightInt4PackDtype, &weightInt4Pack, weightFormat);
    } else {
        // Create the weightInt4Pack aclTensor.
        ret = CreateAclTensor(weightInt4PackHostData, weightInt4PackShape, &weightInt4PackDeviceAddr,
                              weightInt4PackDtype, &weightInt4Pack);
    }
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  }
  // Create the y aclTensor.
  ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &y);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the antiquantScale aclTensor.
  ret = CreateAclTensor(antiquantScaleHostData, antiquantScaleShape, &antiquantScaleDeviceAddr, aclDataType::ACL_FLOAT, &antiquantScale);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // Create the xFp16 aclTensor.
  void* xFp16DeviceAddr = nullptr;
  aclTensor* xFp16 = nullptr;
  ret = CreateAclTensor(xHostData, xShape, &xFp16DeviceAddr, aclDataType::ACL_FLOAT16, &xFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the antiquantScale aclTensor.
  void* antiquantScaleFp16DeviceAddr = nullptr;
  aclTensor* antiquantScaleFp16 = nullptr;
  ret = CreateAclTensor(antiquantScaleHostData, antiquantScaleShape, &antiquantScaleFp16DeviceAddr, aclDataType::ACL_FLOAT16, &antiquantScaleFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create yFp16 aclTensor.
  void* yFp16DeviceAddr = nullptr;
  aclTensor* yFp16 = nullptr;
  ret = CreateAclTensor(yHostData, yShape, &yFp16DeviceAddr, aclDataType::ACL_FLOAT16, &yFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

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

  // Convert weight from int32 to int4pack.
  ret = aclnnConvertWeightToINT4PackGetWorkspaceSize(weight, weightInt4Pack, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnConvertWeightToINT4PackGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  ret = aclnnConvertWeightToINT4Pack(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnConvertWeightToINT4Pack failed. ERROR: %d\n", ret); return ret);

  // When weight is in transpose scenario and weightInt4Pack shape is NZ, call aclInitTensor to convert to a non-contiguous tensor.
  if (isWeightTransposed && weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
    std::vector<int64_t> strides(weightInt4PackShape.size(), 1);
    for (int64_t i = weightInt4PackShape.size() - 2; i >= 0; i--) {
        strides[i] = weightInt4PackShape[i + 1] * strides[i + 1];
    }
    std::swap(strides[0], strides[1]);
    std::swap(weightInt4PackShape[0], weightInt4PackShape[1]);
    std::vector<int64_t> nzShape = {CEIL_DIV(k, 64), CEIL_DIV(n, 16), 16, 8};
    if (weightInt4PackDtype == aclDataType::ACL_INT4) {
        nzShape[3] = 64;
    }
    aclInitTensor(weightInt4Pack, weightInt4PackShape.data(), weightInt4PackShape.size(), weightInt4PackDtype, strides.data(), 0,
                  weightFormat, nzShape.data(), nzShape.size(), weightInt4PackDeviceAddr);
  }

  // Call cast to generate FP16 input.
  ret = aclnnCastGetWorkspaceSize(x, aclDataType::ACL_FLOAT16, xFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize0 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast0 failed. ERROR: %d\n", ret); return ret);

  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  ret = aclnnCastGetWorkspaceSize(antiquantScale, aclDataType::ACL_FLOAT16, antiquantScaleFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize1 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast1 failed. ERROR: %d\n", ret); return ret);

  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // Call the first-phase API of aclnnWeightQuantBatchMatmulV2.
  ret = aclnnWeightQuantBatchMatmulV2GetWorkspaceSize(xFp16, weightInt4Pack, antiquantScaleFp16, nullptr, nullptr, nullptr, nullptr, 0, yFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantBatchMatmulV2GetWorkspaceSize 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 aclnnWeightQuantBatchMatmulV2.
  ret = aclnnWeightQuantBatchMatmulV2(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantBatchMatmulV2 failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Wait synchronously for task execution to finish.
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

// Convert the output to FP32.
  ret = aclnnCastGetWorkspaceSize(yFp16, aclDataType::ACL_FLOAT, y, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize2 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast2 failed. ERROR: %d\n", ret); return ret);

  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 device-side memory to the host side. Modify based on the specific API definition.
  auto size = GetShapeSize(yShape);
  std::vector<float> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), yDeviceAddr,
                    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: %f\n", i, resultData[i]);
  }

  // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
  aclDestroyTensor(x);
  aclDestroyTensor(weight);
  aclDestroyTensor(weightInt4Pack);
  aclDestroyTensor(antiquantScale);
  aclDestroyTensor(y);
  aclDestroyTensor(xFp16);
  aclDestroyTensor(antiquantScaleFp16);
  aclDestroyTensor(yFp16);

  // 7. Release device resources.
  aclrtFree(xDeviceAddr);
  aclrtFree(weightDeviceAddr);
  aclrtFree(weightInt4PackDeviceAddr);
  aclrtFree(antiquantScaleDeviceAddr);
  aclrtFree(yDeviceAddr);
  aclrtFree(xFp16DeviceAddr);
  aclrtFree(antiquantScaleFp16DeviceAddr);
  aclrtFree(yFp16DeviceAddr);

  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();

  return 0;
}

Ascend 950PR/Ascend 950DT：

The sample code is as follows (INT32 input). For reference only. For the specific compilation and execution process, see Compile and Run Samples. Pseudo-quantization includes the aclnnWeightQuantBatchMatmulV2 and aclnnWeightQuantBatchMatmulV3 APIs. Here, aclnnWeightQuantBatchMatmulV2 is used as an example.


#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_cast.h"
#include "aclnnop/aclnn_weight_quant_batch_matmul_v2.h"

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

#define LOG_PRINT(message, ...)     \
  do {                              \
    printf(message, ##__VA_ARGS__); \
  } while (0)

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

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

extern "C" aclnnStatus aclnnConvertWeightToINT4PackGetWorkspaceSize(const aclTensor *weight, aclTensor *weightInt4Pack,
    uint64_t *workspaceSize, aclOpExecutor **executor);

extern "C" aclnnStatus aclnnConvertWeightToINT4Pack(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
    aclrtStream stream);

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 apply for on-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-side data to on-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 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 T>
int CreateAclTensorInt4(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor, aclFormat format) {
  auto size = hostData.size() * sizeof(T);
  // Call aclrtMalloc to apply for on-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-side data to on-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 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.
  if (format == aclFormat::ACL_FORMAT_ND) {
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
  } else {
    std::vector<int64_t> nzShape;
    if (dataType == aclDataType::ACL_INT4) {
        nzShape = {CEIL_DIV(shape[1], 16), CEIL_DIV(shape[0], 16), 16, 16};
    } else {
        nzShape = {CEIL_DIV(shape[1], 2), CEIL_DIV(shape[0], 16), 16, 2};
    }
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0,
                              aclFormat::ACL_FORMAT_FRACTAL_NZ, nzShape.data(), nzShape.size(), *deviceAddr);
  }

  return 0;
}

int main() {
  // 1. (Boilerplate) Initialize device/stream. See the ACL API manual.
  // Fill in the deviceId based on the actual device.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
  aclDataType weightInt4PackDtype = aclDataType::ACL_INT4;
  aclFormat weightFormat = aclFormat::ACL_FORMAT_FRACTAL_NZ;
  bool isWeightTransposed = false;

  // 2. Construct inputs and outputs. Customize the construction based on the API definition.
  int64_t m = 16;
  int64_t k = 64;
  int64_t n = 64;
  int64_t weightDim0 = k;
  int64_t weightDim1 = n;
  if (isWeightTransposed) {
    weightDim0 = n;
    weightDim1 = k;
  }
  std::vector<int64_t> xShape = {m, k};
  std::vector<int64_t> weightShape = {weightDim0, weightDim1};
  std::vector<int64_t> weightInt4PackShape;
  if (weightInt4PackDtype == aclDataType::ACL_INT4) {
    weightInt4PackShape = {weightDim0, weightDim1};
  } else {
    weightInt4PackShape = {weightDim0, weightDim1/8};
  }
  std::vector<int64_t> yShape = {m, n};
  void* xDeviceAddr = nullptr;
  void* weightDeviceAddr = nullptr;
  void* weightPackDeviceAddr = nullptr;
  void* yDeviceAddr = nullptr;
  aclTensor* x = nullptr;
  aclTensor* weight = nullptr;
  aclTensor* weightPacked = nullptr;
  aclTensor* y = nullptr;
  std::vector<float> xHostData(m * k, 1);
  std::vector<int32_t> weightHostData(k * n, 1);
  std::vector<float> yHostData(m * n, 0);

  std::vector<int64_t> antiquantScaleShape = {n};
  void* antiquantScaleDeviceAddr = nullptr;
  aclTensor* antiquantScale = nullptr;
  std::vector<float> antiquantScaleHostData(n, 1);

  // Create x aclTensor.
  ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create weight aclTensor.
  ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_INT32, &weight);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  if (weightInt4PackDtype == aclDataType::ACL_INT4) {
    std::vector<int8_t> weightInt4PackHostData(n * k / 2, 0); // One int8 value stores two int4 values, so divided by 2 here.
    if (weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
      weightInt4PackHostData.resize(CEIL_ALIGN(weightDim1/2, 8) * CEIL_ALIGN(weightDim0, 16), 0);
    }
    // Create weightPacked aclTensor.
    ret = CreateAclTensorInt4(weightInt4PackHostData, weightInt4PackShape, &weightPackDeviceAddr,
                              weightInt4PackDtype, &weightPacked, weightFormat);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  } else {
    std::vector<int32_t> weightInt4PackHostData(n * k / 8, 1); // One int32 value stores eight int4 values, so divided by 8 here.
    if (weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
      weightInt4PackHostData.resize(CEIL_ALIGN(weightDim1/8, 2) * CEIL_ALIGN(weightDim0, 16), 0);
      ret = CreateAclTensorInt4(weightInt4PackHostData, weightInt4PackShape, &weightPackDeviceAddr,
                                weightInt4PackDtype, &weightPacked, weightFormat);
    } else {
        // Create weightPacked aclTensor.
        ret = CreateAclTensor(weightInt4PackHostData, weightInt4PackShape, &weightPackDeviceAddr,
                              weightInt4PackDtype, &weightPacked);
    }
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  }
  // Create y aclTensor.
  ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &y);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create antiquantScale aclTensor.
  ret = CreateAclTensor(antiquantScaleHostData, antiquantScaleShape, &antiquantScaleDeviceAddr, aclDataType::ACL_FLOAT, &antiquantScale);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // Create xFp16 aclTensor.
  void* xFp16DeviceAddr = nullptr;
  aclTensor* xFp16 = nullptr;
  ret = CreateAclTensor(xHostData, xShape, &xFp16DeviceAddr, aclDataType::ACL_FLOAT16, &xFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create antiquantScale aclTensor.
  void* antiquantScaleFp16DeviceAddr = nullptr;
  aclTensor* antiquantScaleFp16 = nullptr;
  ret = CreateAclTensor(antiquantScaleHostData, antiquantScaleShape, &antiquantScaleFp16DeviceAddr, aclDataType::ACL_FLOAT16, &antiquantScaleFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create yFp16 aclTensor.
  void* yFp16DeviceAddr = nullptr;
  aclTensor* yFp16 = nullptr;
  ret = CreateAclTensor(yHostData, yShape, &yFp16DeviceAddr, aclDataType::ACL_FLOAT16, &yFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

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

  // Convert weight from int32 to int4pack.
  ret = aclnnConvertWeightToINT4PackGetWorkspaceSize(weight, weightPacked, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnConvertWeightToINT4PackGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  ret = aclnnConvertWeightToINT4Pack(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnConvertWeightToINT4Pack failed. ERROR: %d\n", ret); return ret);

  // When weight is transposed and weightPacked shape is NZ, call aclInitTensor to convert to a non-contiguous tensor.
  if (isWeightTransposed && weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
    std::vector<int64_t> strides(weightInt4PackShape.size(), 1);
    for (int64_t i = weightInt4PackShape.size() - 2; i >= 0; i--) {
        strides[i] = weightInt4PackShape[i + 1] * strides[i + 1];
    }
    std::swap(strides[0], strides[1]);
    std::swap(weightInt4PackShape[0], weightInt4PackShape[1]);
    std::vector<int64_t> nzShape = {CEIL_DIV(k, 16), CEIL_DIV(n, 16), 16, 2};
    if (weightInt4PackDtype == aclDataType::ACL_INT4) {
        nzShape[3] = 16;
    }
    aclInitTensor(weightPacked, weightInt4PackShape.data(), weightInt4PackShape.size(), weightInt4PackDtype, strides.data(), 0,
                  weightFormat, nzShape.data(), nzShape.size(), weightPackDeviceAddr);
  }

  // Call cast to generate FP16 input.
  ret = aclnnCastGetWorkspaceSize(x, aclDataType::ACL_FLOAT16, xFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize0 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast0 failed. ERROR: %d\n", ret); return ret);

  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  ret = aclnnCastGetWorkspaceSize(antiquantScale, aclDataType::ACL_FLOAT16, antiquantScaleFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize1 failed. ERROR: %d\n", ret); return ret);
  // Allocate 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast1 failed. ERROR: %d\n", ret); return ret);

  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // Call the first-phase API of aclnnWeightQuantBatchMatmulV2.
  ret = aclnnWeightQuantBatchMatmulV2GetWorkspaceSize(xFp16, weightPacked, antiquantScaleFp16, nullptr, nullptr, nullptr, nullptr, 0, yFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantBatchMatmulV2GetWorkspaceSize 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 aclnnWeightQuantBatchMatmulV2.
  ret = aclnnWeightQuantBatchMatmulV2(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantBatchMatmulV2 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);

// Convert the output to FP32.
  ret = aclnnCastGetWorkspaceSize(yFp16, aclDataType::ACL_FLOAT, y, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize2 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast2 failed. ERROR: %d\n", ret); return ret);

  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 on-device memory to the host side. Modify based on the specific API interface definition.
  auto size = GetShapeSize(yShape);
  std::vector<float> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), yDeviceAddr,
                    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: %f\n", i, resultData[i]);
  }

  // 6. Release aclTensor and aclScalar. Modify based on the specific API interface definition.
  aclDestroyTensor(x);
  aclDestroyTensor(weight);
  aclDestroyTensor(weightPacked);
  aclDestroyTensor(antiquantScale);
  aclDestroyTensor(y);
  aclDestroyTensor(xFp16);
  aclDestroyTensor(antiquantScaleFp16);
  aclDestroyTensor(yFp16);

  // 7. Release device resources.
  aclrtFree(xDeviceAddr);
  aclrtFree(weightDeviceAddr);
  aclrtFree(weightPackDeviceAddr);
  aclrtFree(antiquantScaleDeviceAddr);
  aclrtFree(yDeviceAddr);
  aclrtFree(xFp16DeviceAddr);
  aclrtFree(antiquantScaleFp16DeviceAddr);
  aclrtFree(yFp16DeviceAddr);

  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();

  return 0;
}

Ascend 950PR/Ascend 950DT: The sample code is as follows (FLOAT input). For reference only. For the specific compilation and execution process, see Compile and Run Samples. Pseudo-quantization includes the aclnnWeightQuantBatchMatmulV2 and aclnnWeightQuantBatchMatmulV3 APIs. Here, aclnnWeightQuantBatchMatmulV2 is used as an example.


#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_cast.h"
#include "aclnnop/aclnn_weight_quant_batch_matmul_v2.h"

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

#define LOG_PRINT(message, ...)     \
  do {                              \
    printf(message, ##__VA_ARGS__); \
  } while (0)

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

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

extern "C" aclnnStatus aclnnConvertWeightToINT4PackGetWorkspaceSize(const aclTensor *weight, aclTensor *weightInt4Pack,
    uint64_t *workspaceSize, aclOpExecutor **executor);

extern "C" aclnnStatus aclnnConvertWeightToINT4Pack(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
    aclrtStream stream);

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 apply for on-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-side data to on-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 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 T>
int CreateAclTensorB4(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor, aclFormat format) {
  auto size = hostData.size() * sizeof(T);
  // Call aclrtMalloc to apply for on-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-side data to on-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 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.
  if (format == aclFormat::ACL_FORMAT_ND) {
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
  } else {
    std::vector<int64_t> nzShape;
    if (dataType == aclDataType::ACL_INT4 || dataType == aclDataType::ACL_FLOAT4_E2M1) {
        nzShape = {CEIL_DIV(shape[1], 16), CEIL_DIV(shape[0], 16), 16, 16};
    } else {
        nzShape = {CEIL_DIV(shape[1], 2), CEIL_DIV(shape[0], 16), 16, 2};
    }
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0,
                              aclFormat::ACL_FORMAT_FRACTAL_NZ, nzShape.data(), nzShape.size(), *deviceAddr);
  }

  return 0;
}

int main() {
  // 1. (Boilerplate) Initialize device/stream. Refer to the acl API manual.
  // Fill in deviceId based on the actual device.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
  aclDataType weightPackedDtype = aclDataType::ACL_FLOAT4_E2M1; // Optional: ACL_FLOAT type.
  aclFormat weightFormat = aclFormat::ACL_FORMAT_FRACTAL_NZ; // Optional: ACL_FORMAT_ND.
  bool isWeightTransposed = false; // ND: supports true/false. NZ: supports false.

  // 2. Construct inputs and outputs. Custom construction is required based on the API interface definition.
  int64_t m = 16;
  int64_t k = 64;
  int64_t n = 64;
  int64_t antiquantGroupSize = 32;
  int64_t weightDim0 = k;
  int64_t weightDim1 = n;
  if (isWeightTransposed) {
    weightDim0 = n;
    weightDim1 = k;
  }
  std::vector<int64_t> xShape = {m, k};
  std::vector<int64_t> weightShape = {weightDim0, weightDim1};
  std::vector<int64_t> weightPackedShape;
  if (weightPackedDtype == aclDataType::ACL_FLOAT4_E2M1) {
    weightPackedShape = {weightDim0, weightDim1};
  } else {
    weightPackedShape = {weightDim0, weightDim1 / 8};
  }
  std::vector<int64_t> yShape = {m, n};
  void* xDeviceAddr = nullptr;
  void* weightDeviceAddr = nullptr;
  void* weightPackDeviceAddr = nullptr;
  void* yDeviceAddr = nullptr;
  aclTensor* x = nullptr;
  aclTensor* weight = nullptr;
  aclTensor* weightPacked = nullptr;
  aclTensor* y = nullptr;
  std::vector<float> xHostData(m * k, 1);
  std::vector<float> weightHostData(k * n, 1);
  std::vector<float> yHostData(m * n, 0);

  std::vector<int64_t> antiquantScaleShape = {k / antiquantGroupSize, n};
  std::vector<uint8_t> antiquantScaleHostData(n * k / antiquantGroupSize, 1); // Use uint8 to carry float8_e8m0 data.

  // Create the x aclTensor.
  ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the weight aclTensor.
  ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_FLOAT, &weight);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  if (weightPackedDtype == aclDataType::ACL_FLOAT4_E2M1) {
    std::vector<int8_t> weightB4PackHostData(n * k / 2, 0); // One B8 data stores two B4 data, so divided by 2 here.
    if (weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
      weightB4PackHostData.resize(CEIL_ALIGN(weightDim1 / 2, 8) * CEIL_ALIGN(weightDim0, 16), 0);
    }
    // Create weightPacked aclTensor.
    ret = CreateAclTensorB4(weightB4PackHostData, weightPackedShape, &weightPackDeviceAddr,
                              weightPackedDtype, &weightPacked, weightFormat);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  } else {
    std::vector<int32_t> weightB4PackHostData(n * k / 8, 1); // One int32 data stores eight int4 data, so divided by 8 here.
    if (weightFormat == aclFormat::ACL_FORMAT_FRACTAL_NZ) {
      weightB4PackHostData.resize(CEIL_ALIGN(weightDim1 / 8, 2) * CEIL_ALIGN(weightDim0, 16), 0);
      ret = CreateAclTensorB4(weightB4PackHostData, weightPackedShape, &weightPackDeviceAddr,
                                weightPackedDtype, &weightPacked, weightFormat);
    } else {
        // Create weightPacked aclTensor.
        ret = CreateAclTensor(weightB4PackHostData, weightPackedShape, &weightPackDeviceAddr,
                              weightPackedDtype, &weightPacked);
    }
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  }
  // Create y aclTensor.
  ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &y);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // Create xFp16 aclTensor.
  void* xFp16DeviceAddr = nullptr;
  aclTensor* xFp16 = nullptr;
  ret = CreateAclTensor(xHostData, xShape, &xFp16DeviceAddr, aclDataType::ACL_FLOAT16, &xFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create antiquantScale aclTensor.
  void* antiquantScaleFp8DeviceAddr = nullptr;
  aclTensor* antiquantScaleFp8 = nullptr;
  ret = CreateAclTensor(antiquantScaleHostData, antiquantScaleShape, &antiquantScaleFp8DeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &antiquantScaleFp8);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create yFp16 aclTensor.
  void* yFp16DeviceAddr = nullptr;
  aclTensor* yFp16 = nullptr;
  ret = CreateAclTensor(yHostData, yShape, &yFp16DeviceAddr, aclDataType::ACL_FLOAT16, &yFp16);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

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

  // Convert weight from int32 to int4pack.
  ret = aclnnConvertWeightToINT4PackGetWorkspaceSize(weight, weightPacked, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnConvertWeightToINT4PackGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  ret = aclnnConvertWeightToINT4Pack(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnConvertWeightToINT4Pack failed. ERROR: %d\n", ret); return ret);

  // When weight is transposed and weightPacked shape is ND, call aclInitTensor to convert to a non-contiguous tensor.
  if (isWeightTransposed && weightFormat == aclFormat::ACL_FORMAT_ND) {
    weightPackedShape = {k, n};
    std::vector<int64_t> strides = {1, k};
    std::vector<int64_t> storage_shape = {n, k};
    if (weightPackedDtype == aclDataType::ACL_FLOAT) {
      weightPackedShape = {k / 8, n};
      strides = {1, k / 8};
      storage_shape = {n, k / 8};
    }
    aclInitTensor(weightPacked, weightPackedShape.data(), weightPackedShape.size(), weightPackedDtype, strides.data(), 0,
                  weightFormat, storage_shape.data(), storage_shape.size(), weightPackDeviceAddr);
  }

  // Call cast to generate FP16 input.
  ret = aclnnCastGetWorkspaceSize(x, aclDataType::ACL_FLOAT16, xFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize0 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast0 failed. ERROR: %d\n", ret); return ret);

  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // Call the first-phase API of aclnnWeightQuantBatchMatmulV2.
  ret = aclnnWeightQuantBatchMatmulV2GetWorkspaceSize(xFp16, weightPacked, antiquantScaleFp8, nullptr, nullptr, nullptr, nullptr, antiquantGroupSize, yFp16, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantBatchMatmulV2GetWorkspaceSize 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 aclnnWeightQuantBatchMatmulV2.
  ret = aclnnWeightQuantBatchMatmulV2(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnWeightQuantBatchMatmulV2 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);

// Convert the output to FP32.
  ret = aclnnCastGetWorkspaceSize(yFp16, aclDataType::ACL_FLOAT, y, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCastGetWorkspaceSize2 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);
  }
  ret = aclnnCast(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCast2 failed. ERROR: %d\n", ret); return ret);

  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 on-device memory to the host side. Modify based on the specific API interface definition.
  auto size = GetShapeSize(yShape);
  std::vector<float> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), yDeviceAddr,
                    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: %f\n", i, resultData[i]);
  }

  // 6. Release aclTensor and aclScalar. Modify based on the specific API's interface definition.
  aclDestroyTensor(x);
  aclDestroyTensor(weight);
  aclDestroyTensor(weightPacked);
  aclDestroyTensor(y);
  aclDestroyTensor(xFp16);
  aclDestroyTensor(antiquantScaleFp8);
  aclDestroyTensor(yFp16);

  // 7. Release device resources.
  aclrtFree(xDeviceAddr);
  aclrtFree(weightDeviceAddr);
  aclrtFree(weightPackDeviceAddr);
  aclrtFree(yDeviceAddr);
  aclrtFree(xFp16DeviceAddr);
  aclrtFree(antiquantScaleFp8DeviceAddr);
  aclrtFree(yFp16DeviceAddr);

  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();

  return 0;
}