aclnnAddRmsNorm

📄 View Source Code

Applicable Products

Function

• Description: The RMSNorm operator is a commonly used normalization operation in large models. Compared with the LayerNorm operator, it removes the part of subtracting the mean. The AddRmsNorm operator fuses the Add operator before RMSNorm to reduce data movement.

• Formula:

Function Prototype

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

aclnnStatus aclnnAddRmsNormGetWorkspaceSize(
  const aclTensor *x1,
  const aclTensor *x2,
  const aclTensor *gamma,
  double           epsilon,
  aclTensor       *yOut,
  aclTensor       *rstdOut,
  aclTensor       *xOut,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnAddRmsNormGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (shape)	Non-contiguous Tensor
  x1 (aclTensor*)	Input	Represents the first input for the Add computation. Corresponds to `x1` in the formula.	Supports null Tensor.	FLOAT32, FLOAT16, BFLOAT16	ND	1-8	√
  x2 (aclTensor*)	Input	Represents the second input for the Add operation. Corresponds to `x2` in the corresponding formula.	Supports null Tensor. The shape and data type must be consistent with the shape and data type of x1.	FLOAT32, FLOAT16, BFLOAT16	ND	1-8	√
  gamma (aclTensor*)	Input	Represents the scaling factor (weight) of RMSNorm. Corresponds to `gamma` in the corresponding formula.	Supports null Tensor. The data type must be consistent with that of `x1`.	FLOAT32, FLOAT16, BFLOAT16	ND	1-8	√
  epsilon (double)	Input	A value added to the denominator to ensure numerical stability. Corresponds to `epsilon` in the formula.	The value of epsilon must be greater than or equal to zero. The recommended value is 1e-6.	-	-	-	-
  yOut (aclTensor*)	Output	Represents the final output. Corresponds to `RmsNorm(x)` in the formula.	Supports null tensor. The shape and data type must be consistent with those of input `x1`.	FLOAT32, FLOAT16, BFLOAT16	ND	1-8	√
  rstdOut (aclTensor*)	Output	Indicates the reciprocal of the standard deviation after normalization. Corresponds to the reciprocal of `Rms(x)` in the corresponding formula.	Supports null Tensor. The shape is consistent with the first few dimensions of `x1`, which represent the dimensions that do not require normalization. Example of the relationship between rstdOut shape, x1 shape, and gamma shape: if x1 shape: (2, 3, 4, 8), gamma shape: (8), rstdOut shape: (2, 3, 4, 1); if x1 shape: (2, 3, 4, 8), gamma shape: (4, 8), rstdOut shape: (2, 3, 1, 1). When nullptr is passed, the final output of this parameter is invalid.	FLOAT32	ND	1-8	√
  xOut (aclTensor*)	Output	Represents the result of the Add computation. Corresponds to `x` in the formula.	Supports null tensor. The shape and data type must be consistent with those of the input `x1`. When `xOut` is passed as nullptr, the final output of this parameter is invalid.	FLOAT32, FLOAT16, BFLOAT16	ND	1-8	√
  workspaceSize (uint64_t*)	Output	Returns the workspace size to be allocated on the Device side.	-	-	-	-	-
  executor (aclOpExecutor**)	Output	Returns the operator executor, which contains the operator computation process.	-	-	-	-	-

  • Atlas inference products:

    • The data type of parameters x1, x2, gamma, yOut, and xOut does not support BFLOAT16.

    • The parameter rstdOut is invalid in the current product usage scenario.

  • Ascend 950PR/Ascend 950DT:

    The shape of the parameter gamma must be consistent with the last few dimensions of x1, where the last few dimensions are the dimensions of x1 to be normalized.

• 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	x1, x2, gamma, and yOut are null pointers.
  		When the preset value of rstdOut is not nullptr, the preset value of xOut is nullptr.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of the input or output is not within the supported range.
  		The input and output parameters do not meet the constraints in Parameters.

aclnnAddRmsNorm

• Parameters

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

• Boundary value scenarios

• When the input is Inf, the output is Inf.

• When the input is NaN, the output is NaN.

• The supported combinations of inputs x1, x2, gamma, yOut, rstdOut, and xOut are as follows:

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

    x1	x2	gamma	yOut	rstdOut	xOut
    FLOAT32	FLOAT32	shape must be consistent with the last few dimensions of x1, where the last few dimensions are the dimensions of x1 to be normalized; FLOAT32	FLOAT32	Required	Required, FLOAT32
    FLOAT16	FLOAT16	shape must be consistent with the last few dimensions of x1, where the last few dimensions are the dimensions of x1 to be normalized; FLOAT16	FLOAT16	Required	Required, FLOAT16
    BFLOAT16	BFLOAT16	shape must be consistent with the last few dimensions of x1, where the last few dimensions are the dimensions of x1 to be normalized; BFLOAT16	BFLOAT16	Required	Required, BFLOAT16
    FLOAT16	FLOAT16	shape is [1, the last dimension of x1]; FLOAT16	FLOAT16	null pointer	null pointer, FLOAT16
    BFLOAT16	BFLOAT16	shape is [1, the last dimension of x1]; BFLOAT16	BFLOAT16	null pointer	null pointer, BFLOAT16
    FLOAT16	FLOAT16	shape is [1, the last dimension of x1]; FLOAT16	FLOAT16	null pointer	Required, FLOAT16
    BFLOAT16	BFLOAT16	shape is [1, the last dimension of x1]; BFLOAT16	BFLOAT16	null pointer	Required, BFLOAT16

  • Atlas inference products:

    x1	x2	gamma	yOut	rstdOut	xOut
    FLOAT32	FLOAT32	shape must be consistent with the last few dimensions of x1, where the last few dimensions are the dimensions of x1 to be normalized; FLOAT32	FLOAT32	Required	Required, FLOAT32
    FLOAT16	FLOAT16	shape must be consistent with the last few dimensions of x1, where the last few dimensions are the dimensions of x1 to be normalized; FLOAT16	FLOAT16	Required	Required, FLOAT16
    FLOAT16	FLOAT16	shape is [1, the last dimension of x1]; FLOAT16	FLOAT16	null pointer	null pointer, FLOAT16
    FLOAT16	FLOAT16	shape is [1, the last dimension of x1]; FLOAT16	FLOAT16	null pointer	Required, FLOAT16

• Deterministic computation:

  • aclnnAddRmsNorm 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.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 for 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 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 data from host to 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 strides of a contiguous tensor.
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }

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

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);
    // Handle the check based on your needs.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct inputs and outputs. Customize the construction based on the API.
    std::vector<int64_t> xShape = {2, 16};
    std::vector<int64_t> gammaShape = {16};
    std::vector<int64_t> yShape = {2, 16};
    std::vector<int64_t> rstdShape = {2, 1};
    void* x1DeviceAddr = nullptr;
    void* x2DeviceAddr = nullptr;
    void* gammaDeviceAddr = nullptr;
    void* yDeviceAddr = nullptr;
    void* rstdDeviceAddr = nullptr;
    void* xDeviceAddr = nullptr;
    aclTensor* x1 = nullptr;
    aclTensor* x2 = nullptr;
    aclTensor* gamma = nullptr;
    aclTensor* y = nullptr;
    aclTensor* rstd = nullptr;
    aclTensor* x = nullptr;
    std::vector<float> x1HostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
                                     0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    std::vector<float> x2HostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
                                     0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    std::vector<float> gammaHostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    std::vector<float> yHostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
                                    0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    std::vector<float> rstdHostData = {1, 2};
    std::vector<float> xHostData = {0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
                                    0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7};
    float epsilon = 1e-6;

    // Create x1 aclTensor
    ret = CreateAclTensor(x1HostData, xShape, &x1DeviceAddr, aclDataType::ACL_FLOAT, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create x2 aclTensor
    ret = CreateAclTensor(x2HostData, xShape, &x2DeviceAddr, aclDataType::ACL_FLOAT, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create gamma aclTensor
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT, &gamma);
    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 rstd aclTensor
    ret = CreateAclTensor(rstdHostData, rstdShape, &rstdDeviceAddr, aclDataType::ACL_FLOAT, &rstd);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &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 aclnnAddRmsNorm.
    ret = aclnnAddRmsNormGetWorkspaceSize(x1, x2, gamma, epsilon, y, rstd, x, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNormGetWorkspaceSize 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 aclnnAddRmsNorm.
    ret = aclnnAddRmsNorm(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddRmsNorm failed. ERROR: %d\n", ret); return ret);
    // 4. (Boilerplate) Synchronously 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 value and copy the result from the device memory to the host memory. 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(float),
        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("y result[%ld] is: %f\n", i, resultData[i]);
    }

    // 6. Release aclTensor. Modify based on the specific API definition.
    aclDestroyTensor(x1);
    aclDestroyTensor(x2);
    aclDestroyTensor(gamma);
    aclDestroyTensor(y);
    aclDestroyTensor(rstd);
    aclDestroyTensor(x);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(x1DeviceAddr);
    aclrtFree(x2DeviceAddr);
    aclrtFree(xDeviceAddr);
    aclrtFree(gammaDeviceAddr);
    aclrtFree(yDeviceAddr);
    aclrtFree(rstdDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}