aclnnAddRmsNormDynamicQuant

Supported Products

Product	Supported
Atlas A3 training series products/Atlas A3 inference series products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference series products	×
Atlas training series products	×

Function

Description: The RmsNorm operator is a normalization operation commonly used in foundation models. Compared with the LayerNorm operator, the RmsNorm operator removes the part of subtracting the mean value. The DynamicQuant operator is used to perform symmetric dynamic quantization on the input tensor. The AddRmsNormDynamicQuant operator fuses the Add operator before RmsNorm and the normalized output of RmsNorm to one or two DynamicQuant operators to reduce move-in and move-out operations.
Formula: $x=x_{1}+x_{2}$ $y = \operatorname{RmsNorm}(x)=\frac{x}{\operatorname{Rms}(\mathbf{x})}\cdot gamma, \quad \text { where } \operatorname{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+epsilon}$
- If smoothScale1Optional and smoothScale2Optional are not input, the outputs y2Out and scale2Out are meaningless. The calculation process is as follows:
$scale1Out=row\_max(abs(y))/127$ $y1Out=round(y/scale1Out)$
- If only smoothScale1Optional is input, the outputs y2Out and scale2Out are meaningless. The calculation process is as follows:
$input = y\cdot smoothScale1Optional$ $scale1Out=row\_max(abs(input))/127$ $y1Out=round(input/scale1Out)$
- If both smoothScale1Optional and smoothScale2Optional are input, all the five outputs of the operator are valid. The calculation process is as follows:
$input1 = y\cdot smoothScale1Optional$ $input2 = y\cdot smoothScale2Optional$ $scale1Out=row\_max(abs(input1))/127$ $scale2Out=row\_max(abs(input2))/127$ $y1Out=round(input1/scale1Out)$ $y2Out=round(input2/scale2Out)$
row_max indicates that the maximum value of each row is calculated.

Prototype

Each operator has two-phase API calls. First, aclnnAddRmsNormDynamicQuantGetWorkspaceSize is called to obtain the input parameters and compute the required workspace size based on the process. Then, aclnnAddRmsNormDynamicQuant is called to perform computation.

aclnnStatus aclnnAddRmsNormDynamicQuantGetWorkspaceSize(
  const aclTensor *x1,
  const aclTensor *x2,
  const aclTensor *gamma,
  const aclTensor *smoothScale1Optional,
  const aclTensor *smoothScale2Optional,
  double           epsilon,
  aclTensor       *y1Out,
  aclTensor       *y2Out,
  aclTensor       *xOut,
  aclTensor       *scale1Out,
  aclTensor       *scale2Out,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnAddRmsNormDynamicQuantGetWorkspaceSize

Parameters:

Name	Input/Output	Description	Precaution	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
x1	Input	Source data tensor in the normalization process. It corresponds to `x1` in the formula.	Empty tensors are supported.	FLOAT16, BFLOAT16	ND	2-8	√
x2	Input	Source data tensor in the normalization process. It corresponds to `x2` in the formula.	Empty tensors are supported. The shape and data type must be the same as those of `x1`.	FLOAT16, BFLOAT16	ND	2-8	√
gamma	Input	Weight tensor in the normalization process. It corresponds to `gamma` in the formula.	Empty tensors are supported. The data type must be the same as that of `x1`. The shape must be the same as the last dimension of `x1`.	FLOAT16, BFLOAT16	ND	1	√
smoothScale1Optional	Input	smoothScale tensor used to obtain `y1Out` during quantization. It corresponds to `smoothScale1Optional` in the formula.	Empty tensors are supported. (Optional) A null pointer can be passed. The shape and data type must be the same as those of `gamma`.	FLOAT16, BFLOAT16	ND	1	√
smoothScale2Optional	Input	smoothScale tensor used to obtain `y2Out` during quantization. It corresponds to `smoothScale2Optional` in the formula.	Empty tensors are supported. (Optional) A null pointer can be passed. It must be used together with smoothScale1Optional. The shape and data type must be the same as those of `gamma`.	FLOAT16, BFLOAT16	ND	1	√
epsilon	Input	`epsilon` in the formula, which is used to prevent division-by-zero errors.	You are advised to pass a small positive number, for example, 1e-6.	DOUBLE	-	-	-
y1Out	Output	Quantized output tensor, which corresponds to `y1Out` in the formula.	Empty tensors are supported. The shape must be the same as that of input `x1`/`x2`. Alternatively, the shape is two-dimensional, the first dimension is equal to the product of the dimensions except the last dimension of `x1`, and the second dimension is equal to the last dimension of `x1`.	INT8, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND	2-8	√
y2Out	Output	Quantized output tensor, which corresponds to `y2Out` in the formula.	Empty tensors are supported. If smoothScale2Optional does not exist, this output is meaningless. If the output of `y2Out` is meaningless, the shape is [1]. In other scenarios, the shape must be the same as that of `y1Out`.	INT8, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND	2-8	√
xOut	Output	Sum of x1 and x2, which corresponds to `x` in the formula.	Empty tensors are supported. The shape and data type must be the same as those of input `x1`/`x2`.	FLOAT16, BFLOAT16	ND	2-8	√
scale1Out	Output	Quantized output of the first channel, which corresponds to `scale1Out` in the formula.	Empty tensors are supported. The shape must be the same as that of input `x1` except the last dimension, or the same as the product of the dimensions of `x1` except the last dimension.	FLOAT32	ND	1-8	√
scale2Out	Output	Quantized output of the second channel, which corresponds to `scale2Out` in the formula.	Empty tensors are supported. If smoothScale2Optional does not exist, this output is meaningless. The shape must be the same as that of `scale1Out`.	FLOAT32	ND	1-8	√
workspaceSize	Output	Size of the workspace required to be allocated on the device.	-	-	-	-	-
executor	Output	Operator executor, containing the operator computation process.	-	-	-	-	-

Atlas A3 training series products/Atlas A3 inference series products and Atlas A2 training products/Atlas A2 inference products:
- The output parameters y1Out and y2Out support only INT8.

Returns:

aclnnStatus: status code. For details, see aclnn Return Codes.

The first-phase API implements input parameter verification. The following errors may be thrown.

Return	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	The input parameter is a required input, output, or attribute, and is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The input or output data type is not supported.
ACLNN_ERR_INNER_TILING_ERROR	561002	smoothScale2Optional is input, but smoothScale1Optional is not input.
ACLNN_ERR_INNER_TILING_ERROR	561002	The shape relationship between the input and output does not meet the expectation.

aclnnAddRmsNormDynamicQuant

Parameters:

Name	Input/Output	Description
workspace	Input	Address of the workspace to be allocated on the device.
workspaceSize	Input	Size of the workspace to be allocated on the device, obtained by calling the first-phase API aclnnAddRmsNormDynamicQuantGetWorkspaceSize.
executor	Input	Operator executor, containing the operator computation process.
stream	Input	Stream for executing the task.

Returns:

aclnnStatus: status code. For details, see aclnn Return Codes.

Constraints

Description of supported types

Empty tensors: Empty input and output are supported.
Data format description

The ND format is recommended for all input and output tensors. If other data formats are used, the framework converts them into the ND format by default for processing.

Description of data types supported by different product models

Atlas A3 training series products/Atlas A3 inference series products and Atlas A2 training products/Atlas A2 inference products:

`x1`	`x2`	`gamma`	`smoothScale1Optional`	`smoothScale2Optional`	`y1Out`	`y2Out`	`xOut`	`scale1Out`	`scale2Out`
FLOAT16	FLOAT16	FLOAT16	FLOAT16	FLOAT16	INT8	INT8	FLOAT16	FLOAT32	FLOAT32
BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	INT8	BFLOAT16	FLOAT32	FLOAT32

Deterministic compute:
- aclnnAddRmsNormDynamicQuant defaults to a deterministic implementation.

Example

The following example is for reference only. For details, see Compilation and Running Sample.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_add_rms_norm_dynamic_quant.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)
{
    // (Fixed writing) 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 allocate memory on the device.
    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 the data from the host to the device.
    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 the contiguous tensor.
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }

    // Call aclCreateTensor to create an aclTensor.
    *tensor = aclCreateTensor(
        shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(),
        *deviceAddr);
    return 0;
}

int main()
{
    // 1. (Fixed writing) Initialize the device and stream. For details, see the ACL API manual.
    // Set the device ID in use.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle the check as required.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct the input and output based on the API.
    std::vector<int64_t> xShape = {2, 8};
    std::vector<int64_t> gammaShape = {8};
    std::vector<int64_t> reduceShape = {2, 1};

    void* x1DeviceAddr = nullptr;
    void* x2DeviceAddr = nullptr;
    void* gammaDeviceAddr = nullptr;
    void* smooth1DeviceAddr = nullptr;
    void* smooth2DeviceAddr = nullptr;

    void* y1DeviceAddr = nullptr;
    void* y2DeviceAddr = nullptr;
    void* xDeviceAddr = nullptr;
    void* scale1DeviceAddr = nullptr;
    void* scale2DeviceAddr = nullptr;

    aclTensor* x1 = nullptr;
    aclTensor* x2 = nullptr;
    aclTensor* gamma = nullptr;
    aclTensor* smooth1 = nullptr;
    aclTensor* smooth2 = nullptr;
    aclTensor* y1 = nullptr;
    aclTensor* y2 = nullptr;
    aclTensor* x = nullptr;
    aclTensor* scale1 = nullptr;
    aclTensor* scale2 = nullptr;

    int64_t xShapeSize = GetShapeSize(xShape);
    int64_t gammaShapeSize = GetShapeSize(gammaShape);
    int64_t reduceShapeSize = GetShapeSize(reduceShape);

    std::vector<short> x1HostData(xShapeSize, 0x3800);
    std::vector<short> x2HostData(xShapeSize, 0x3800);
    std::vector<short> gammaHostData(gammaShapeSize, 0x3e00);
    std::vector<short> smooth1HostData(gammaShapeSize, 0x3e00);
    std::vector<short> smooth2HostData(gammaShapeSize, 0x3e00);

    std::vector<short> y1HostData(xShapeSize, 0);
    std::vector<short> y2HostData(xShapeSize, 0);
    std::vector<short> xHostData(xShapeSize, 0);
    std::vector<short> scale1HostData(reduceShapeSize, 0);
    std::vector<short> scale2HostData(reduceShapeSize, 0);

    float epsilon = 1e-6;

    // Create an x1 aclTensor.
    ret = CreateAclTensor(x1HostData, xShape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create an x2 aclTensor.
    ret = CreateAclTensor(x2HostData, xShape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create a gamma aclTensor.
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT16, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create a smooth1 aclTensor.
    ret = CreateAclTensor(smooth1HostData, gammaShape, &smooth1DeviceAddr, aclDataType::ACL_FLOAT16, &smooth1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create a smooth2 aclTensor.
    ret = CreateAclTensor(smooth2HostData, gammaShape, &smooth2DeviceAddr, aclDataType::ACL_FLOAT16, &smooth2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // Create a y1 aclTensor.
    ret = CreateAclTensor(y1HostData, xShape, &y1DeviceAddr, aclDataType::ACL_INT8, &y1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create a y2 aclTensor.
    ret = CreateAclTensor(y2HostData, xShape, &y2DeviceAddr, aclDataType::ACL_INT8, &y2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create an x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT16, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create an outScale1 aclTensor.
    ret = CreateAclTensor(scale1HostData, reduceShape, &scale1DeviceAddr, aclDataType::ACL_FLOAT, &scale1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create an outScale1 aclTensor.
    ret = CreateAclTensor(scale2HostData, reduceShape, &scale2DeviceAddr, aclDataType::ACL_FLOAT, &scale2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API, which needs to be replaced with the actual API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnAddRmsNormDynamicQuant.
    ret = aclnnAddRmsNormDynamicQuantGetWorkspaceSize(
        x1, x2, gamma, smooth1, smooth2, epsilon, y1, y2, x, scale1, scale2, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnRmsNormGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Allocate device memory based on 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 aclnnAddRmsNormDynamicQuant.
    ret = aclnnAddRmsNormDynamicQuant(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormDynamicQuant failed. ERROR: %d\n", ret); return ret);
    // 4. (Fixed writing) Wait until the task execution is complete.
    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 the device to the host. Modify the code based on the API definition.
    auto size = GetShapeSize(xShape);
    std::vector<int8_t> y1Ret(size, 0);
    ret = aclrtMemcpy(
        y1Ret.data(), y1Ret.size() * sizeof(y1Ret[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, y1Ret[i]);
    }

    // 6. Release aclTensor and aclScalar. Modify the configuration based on the API definition.
    aclDestroyTensor(x1);
    aclDestroyTensor(x2);
    aclDestroyTensor(gamma);
    aclDestroyTensor(smooth1);
    aclDestroyTensor(smooth2);
    aclDestroyTensor(y1);
    aclDestroyTensor(y2);
    aclDestroyTensor(x);
    aclDestroyTensor(scale1);
    aclDestroyTensor(scale2);

    // 7. Release device resources. Modify the configuration based on the API definition.
    aclrtFree(x1DeviceAddr);
    aclrtFree(x2DeviceAddr);
    aclrtFree(gammaDeviceAddr);
    aclrtFree(smooth1DeviceAddr);
    aclrtFree(smooth2DeviceAddr);
    aclrtFree(y1DeviceAddr);
    aclrtFree(y2DeviceAddr);
    aclrtFree(xDeviceAddr);
    aclrtFree(scale1DeviceAddr);
    aclrtFree(scale2DeviceAddr);

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