aclnnAddRmsNormQuantV2

📄 View Source Code

Applicable Products

Function

• Description: RmsNorm is a commonly used normalization operation in large models. Compared with LayerNorm, it removes the part of subtracting the mean. The AddRmsNormQuant operator fuses the Add operator before RmsNorm and the output of RmsNorm normalization into one or two Quantize operators, reducing data movement. Compared with AddRmsNormQuant, the AddRmsNormQuantV2 operator adds the bias term betaOptional parameter during the RmsNorm computation, which is the beta in the formula.

• Formula:

• When divMode is true:

  $$y1Out=round((y/scales1)+zeroPoints1Optional)$$ $$y2Out=round((y/scales2)+zeroPoints2Optional)$$

• When divMode is false:

  $$y1Out=round((y*scales1)+zeroPoints1Optional)$$ $$y2Out=round((y*scales2)+zeroPoints2Optional)$$

Function Prototype

Each operator uses a two-phase API. You must first call the aclnnAddRmsNormQuantV2GetWorkspaceSize API to obtain the input parameters and calculate the required workspace size based on the computation process, and then call the aclnnAddRmsNormQuantV2 API to perform the computation.

aclnnStatus aclnnAddRmsNormQuantV2GetWorkspaceSize(
  const aclTensor *x1,
  const aclTensor *x2,
  const aclTensor *gamma,
  const aclTensor *scales1,
  const aclTensor *scales2Optional,
  const aclTensor *zeroPoints1Optional,
  const aclTensor *zeroPoints2Optional,
  const aclTensor *betaOptional,
  int64_t          axis,
  double           epsilon,
  bool             divMode,
  aclTensor       *y1Out,
  aclTensor       *y2Out,
  aclTensor       *xOut,
  aclTensor       *rmsNormOut,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnAddRmsNormQuantV2GetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (shape)	Non-contiguous Tensor
  x1 (aclTensor*)	Input	Represents the source data tensor during normalization. Corresponds to `x1` in the formula.	Supports null tensor.	FLOAT16, BFLOAT16	ND	1-8	√
  x2 (aclTensor*)	Input	Source data tensor during normalization. Corresponds to `x2` in the formula.	Supports null tensor. The shape and data type needs to be consistent with x1.	FLOAT16, BFLOAT16	ND	1-8	√
  gamma (aclTensor*)	Input	Represents the weight tensor during normalization. Corresponds to `gamma` in the formula.	Supports null tensor. The data type needs to remain consistent with `x1`. The shape needs to remain consistent with the dimension of `x1` to be normalized.	FLOAT16, BFLOAT16	ND	1-8	√
  scales1 (aclTensor*)	Input	Represents the scales tensor for y1Out obtained during quantization, corresponding to `scales1` in the formula.	Supports null tensor. The shape needs to be consistent with `gamma`. When the value of the divMode parameter is True, the value of this parameter cannot be 0.	FLOAT32, BFLOAT16	ND	1-8	√
  scales2Optional (aclTensor*)	Input	Represents the scales tensor used to obtain y2Out during quantization. Corresponds to `scales2` in the formula.	Supports null tensor. Optional parameter that supports passing in a null pointer. The shape and data type need to remain consistent with `scales1`. When the `divMode` parameter is True, the value of this parameter cannot be 0.	FLOAT32, BFLOAT16	ND	1-8	√
  zeroPoints1Optional (aclTensor*)	Input	Represents the offset tensor for y1Out obtained during quantization. Corresponds to `zeroPoints1Optional` in the formula.	Supports null tensor. Optional parameter, supports passing in a null pointer. The shape needs to be consistent with `gamma`.	INT32, BFLOAT16	ND	1-8	√
  zeroPoints2Optional (aclTensor*)	Input	Represents the offset tensor for y2Out obtained during quantization. Corresponds to `zeroPoints2Optional` in the formula.	Supports empty tensors. Optional parameter. Supports passing in a null pointer. The data type needs to remain consistent with `zeroPoints1Optional`. The shape needs to remain consistent with `gamma`.	INT32, BFLOAT16	ND	1-8	√
  betaOptional (aclTensor*)	Input	Bias term during normalization. Corresponds to `beta` in the formula.	Supports null tensor. Optional parameter. Supports passing in a null pointer. The shape and data type need to be consistent with `gamma`.	FLOAT16, BFLOAT16	-	-	-
  axis (int64_t)	Input	Indicates the element-wise axis for quantization. Other axes are broadcast. The specified axis cannot exceed the number of dimensions of the input `x1`.	Currently only -1 is supported. Passing any other value has no effect.	-	-	-	-
  epsilon (double)	Input	Represents the value used to prevent division-by-zero errors, corresponding to `epsilon` in the formula.	It is recommended to pass in a small positive number.	-	-	-	-
  divMode (bool)	Input	Indicates the parameter that determines whether to use division in the quantization formula, corresponding to `divMode` in the formula.	-	-	-	-	-
  y1Out (aclTensor*)	Output	Represents the quantized output tensor, corresponding to `y1Out` in the formula.	Supports null tensor. The shape needs to be consistent with the input `x1`/`x2`.	INT8	ND	1-8	√
  y2Out (aclTensor*)	Output	Represents the quantized output tensor, corresponding to `y2Out` in the formula.	Supports null tensor. Optional output. The shape needs to be consistent with the input `x1`/`x2`. When `scales2Optional` is null, the value of this output is invalid.	INT8	ND	1-8	√
  xOut (aclTensor*)	Output	Represents the sum of x1 and x2, corresponding to `x` in the formula.	Supports null tensor. The shape and data type need to remain consistent with the input `x1`/`x2`.	FLOAT16, BFLOAT16	ND	1-8	√
  rmsNormOut (aclTensor*)	Output	Represents the result after RmsNorm, corresponding to `rmsNormOut` in the formula.	Supports null tensor. Optional output. The shape and data type need to be consistent with the input `x1`/`x2`.	FLOAT16, BFLOAT16	ND	1-8	√
  workspaceSize (uint64_t*)	Output	Returns the workspace size to be applied on the Device side.	-	-	-	-	-
  executor (aclOpExecutor**)	Output	Returns the op executor, which contains the operator computation flow.	-	-	-	-	-

  • Atlas inference products: The data type of parameters x1, x2, gamma, scales1, scales2Optional, zeroPoints1Optional, zeroPoints2Optional, betaOptional, xOut, and rmsNormOut does not support BFLOAT16.

• 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	If the passed parameter is a required input, output, or required attribute and is a null pointer, returns 161001.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of the input or output is not within the supported range.

aclnnAddRmsNormQuantV2

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Memory address of the workspace applied on the Device side.
  workspaceSize	Input	Size of the workspace applied on the Device side, obtained from the first-phase API aclnnAddRmsNormQuantV2GetWorkspaceSize.
  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

• Atlas inference products: The number of data elements in the dimension of x1 and x2 that requires normalization must not be less than 32. The number of data elements in gamma, betaOptional, scales1, scales2Optional, zeroPoints1Optional, and zeroPoints2Optional must not be less than 32.

• The supported scenarios and combinations for the inputs gamma, scales1, scales2Optional, zeroPoints1Optional, zeroPoints2Optional, betaOptional, divMode, y1Out, y2Out, xOut, and rmsNormOut are as follows:

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

    gamma	scales1	scales2Optional	zeroPoints1Optional	zeroPoints2Optional	betaOptional	divMode	y1Out	y2Out	xOut	rmsNormOut
    Shape is [the last dimension of x1] or [1, the last dimension of x1]	Shape is [1]	Null pointer	Required, shape is [1]	Null pointer	Required and shape remains consistent with gamma	True	Required	Invalid output	Required	Null pointer
    Shape is [the last dimension of x1] or [1, the last dimension of x1]	Shape is [1]	Null pointer	Required, shape is [1]	Null pointer	Null pointer	True	Required	Invalid output	Null pointer	Required
    Shape is consistent with the dimension of x1 that requires normalization	Shape remains consistent with gamma	Optional, shape remains consistent with gamma	Optional, shape remains consistent with gamma	Optional, shape remains consistent with gamma	Optional, shape remains consistent with gamma	True/False	Required	When scales2Optional is null, this output is invalid; when scales2Optional is non-null, this output is valid	Required	Null pointer

  • Atlas inference products:

    gamma	scales1	scales2Optional	zeroPoints1Optional	zeroPoints2Optional	betaOptional	divMode	y1Out	y2Out	xOut	rmsNormOut
    The shape is consistent with the dimension of x1 that needs normalization.	The shape is consistent with gamma.	Optional. The shape is consistent with gamma.	Optional. The shape is consistent with gamma.	Optional. The shape is consistent with gamma.	Optional. The shape is consistent with gamma.	True/False	Required.	When scales2Optional is null, this output is invalid. When scales2Optional is non-null, this output is valid.	Required.	Null pointer.

• Boundary value scenario description:

  • Atlas inference products: The input does not support containing inf and NaN.

  • Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products: When the input is inf, the output is inf. When the input is NaN, the output is NaN.

• Dimension boundary description:

  The size of each dimension in the shape of parameters x1, x2, gamma, scales1, scales2Optional, zeroPoints1Optional, zeroPoints2Optional, betaOptional, y1Out, y2Out, xOut, and rmsNormOut must not exceed the maximum INT32 value of 2147483647.

• Data format description:

  It is recommended to use the ND format for all input and output tensors. Other data formats are converted to the ND format by the framework by default for processing.

• Supported data types by product model:

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

    x1 Data Type	x2 Data Type	gamma Data Type	scales1 Data Type	scales2Optional Data Type	zeroPoints1Optional Data Type	zeroPoints2Optional Data Type	betaOptional Data Type	y1Out Data Type	y2Out Data Type	xOut Data Type	rmsNormOut Data Type
    FLOAT16	FLOAT16	FLOAT16	FLOAT32	FLOAT32	INT32	INT32	FLOAT16	INT8	INT8	FLOAT16	FLOAT16
    BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	INT8	BFLOAT16	BFLOAT16

  • Atlas inference products:

    x1 Data Type	x2 Data Type	gamma Data Type	scales1 Data Type	scales2Optional Data Type	zeroPoints1Optional Data Type	zeroPoints2Optional Data Type	betaOptional Data Type	y1Out Data Type	y2Out Data Type	xOut Data Type	rmsNormOut Data Type
    FLOAT16	FLOAT16	FLOAT16	FLOAT32	FLOAT32	INT32	INT32	FLOAT16	INT8	INT8	FLOAT16	FLOAT16

• Deterministic computation:

  • aclnnAddRmsNormQuantV2 defaults to a deterministic implementation.

Example

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

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

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

int Init(int32_t deviceId, aclrtStream* stream)
{
    // Boilerplate: resource initialization.
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);

    return 0;
}

template <typename T>
int CreateAclTensor(
    const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType,
    aclTensor** tensor)
{
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to allocate memory on the device side.
    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 data from the host side to the 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);

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

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);
    // Handle the check based on your needs.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct the inputs and outputs. Customize the construction based on the API interface.
    std::vector<int64_t> xShape = {64, 32};
    std::vector<int64_t> gammaShape = {32};
    std::vector<int64_t> yShape = {64, 32};
    long long xShapeSize = GetShapeSize(xShape);
    long long gammaShapeSize = GetShapeSize(gammaShape);
    void* x1DeviceAddr = nullptr;
    void* x2DeviceAddr = nullptr;
    void* gammaDeviceAddr = nullptr;
    void* betaDeviceAddr = nullptr;
    void* scales1DeviceAddr = nullptr;
    void* zeroPoints1DeviceAddr = nullptr;
    void* y1DeviceAddr = nullptr;
    void* y2DeviceAddr = nullptr;
    void* xDeviceAddr = nullptr;
    aclTensor* x1 = nullptr;
    aclTensor* x2 = nullptr;
    aclTensor* gamma = nullptr;
    aclTensor* beta = nullptr;
    aclTensor* scales1 = nullptr;
    aclTensor* zeroPoints1 = nullptr;
    aclTensor* y1 = nullptr;
    aclTensor* y2 = nullptr;
    aclTensor* x = nullptr;
    std::vector<int16_t> x1HostData(xShapeSize, 0);
    std::vector<int16_t> x2HostData(xShapeSize, 0);
    std::vector<int16_t> gammaHostData(gammaShapeSize, 0);
    std::vector<int16_t> betaHostData(gammaShapeSize, 0);
    std::vector<float> scales1HostData(gammaShapeSize, 1);
    std::vector<int32_t> zeroPoints1HostData(gammaShapeSize, 100);
    std::vector<int8_t> y1HostData(xShapeSize, 0);
    std::vector<int8_t> y2HostData(xShapeSize, 0);
    std::vector<int16_t> xHostData(xShapeSize, 0);
    float epsilon = 1e-6;
    int64_t axis = -1;
    bool divMode = true;
    // Create the x1 aclTensor.
    ret = CreateAclTensor(x1HostData, xShape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create x2 aclTensor.
    ret = CreateAclTensor(x2HostData, xShape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create gamma aclTensor.
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT16, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create beta aclTensor.
    ret = CreateAclTensor(betaHostData, gammaShape, &betaDeviceAddr, aclDataType::ACL_FLOAT16, &beta);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create scales1 aclTensor.
    ret = CreateAclTensor(scales1HostData, gammaShape, &scales1DeviceAddr, aclDataType::ACL_FLOAT, &scales1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create zeroPoints1 aclTensor.
    ret =
        CreateAclTensor(zeroPoints1HostData, gammaShape, &zeroPoints1DeviceAddr, aclDataType::ACL_INT32, &zeroPoints1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create y1 aclTensor.
    ret = CreateAclTensor(y1HostData, yShape, &y1DeviceAddr, aclDataType::ACL_INT8, &y1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create y2 aclTensor.
    ret = CreateAclTensor(y2HostData, yShape, &y2DeviceAddr, aclDataType::ACL_INT8, &y2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT16, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 3. Call the CANN operator library API. Replace with the specific API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnAddRmsNormQuantV2.
    ret = aclnnAddRmsNormQuantV2GetWorkspaceSize(
        x1, x2, gamma, scales1, nullptr, zeroPoints1, nullptr, beta, axis, epsilon, divMode, y1, y2, x, nullptr,
        &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormQuantV2GetWorkspaceSize failed. ERROR: %d\n", ret);
              return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    void* workspaceAddr = nullptr;
    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 aclnnAddRmsNormQuantV2.
    ret = aclnnAddRmsNormQuantV2(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormQuantV2 failed. ERROR: %d\n", ret); return ret);
    // 4. (Boilerplate) Synchronize and wait for task execution to complete.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
    // 5. Obtain the output values and copy the results from the device memory to the host memory. Modify based on the specific API definition.
    auto size = GetShapeSize(yShape);
    std::vector<int8_t> resultData(size, 0);
    ret = aclrtMemcpy(
        resultData.data(), resultData.size() * sizeof(resultData[0]), y1DeviceAddr, size * sizeof(int8_t),
        ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("result[%ld] is: %d\n", i, resultData[i]);
    }

    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(x1);
    aclDestroyTensor(x2);
    aclDestroyTensor(gamma);
    aclDestroyTensor(beta);
    aclDestroyTensor(scales1);
    aclDestroyTensor(zeroPoints1);
    aclDestroyTensor(y1);
    aclDestroyTensor(y2);
    aclDestroyTensor(x);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(x1DeviceAddr);
    aclrtFree(x2DeviceAddr);
    aclrtFree(gammaDeviceAddr);
    aclrtFree(betaDeviceAddr);
    aclrtFree(scales1DeviceAddr);
    aclrtFree(zeroPoints1DeviceAddr);
    aclrtFree(y1DeviceAddr);
    aclrtFree(y2DeviceAddr);
    aclrtFree(xDeviceAddr);

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