Documentation
Download
文档中心
aclnnInplaceAddRmsNorm

aclnnInplaceAddRmsNorm

📄 View Source Code

Applicable Products

ProductSupported
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

  • 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:

    xi=x1Refi+x2Refix_i=x1Ref_i+x2Ref_i RmsNorm⁡(xi)=1Rms⁡(x)∗xi∗gammai, where Rms⁡(x)=1n∑i=1nxi2+epsilon\operatorname{RmsNorm}(x_i)=\frac{1}{\operatorname{Rms}(\mathbf{x})} * x_i * gamma_i, \quad \text { where } \operatorname{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+epsilon}

Function Prototype

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

aclnnStatus aclnnInplaceAddRmsNormGetWorkspaceSize(
  aclTensor       *x1Ref,
  aclTensor       *x2Ref,
  const aclTensor *gamma,
  double           epsilon,
  const aclTensor *rstdOut,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)
aclnnStatus aclnnInplaceAddRmsNorm(
  void          *workspace,
  uint64_t       workspaceSize,
  aclOpExecutor *executor,
  aclrtStream    stream)

aclnnInplaceAddRmsNormGetWorkspaceSize

  • Parameters

    Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-contiguous Tensor
    x1Ref (aclTensor*) Input/Output Indicates the first input for the Add computation and the final output result after normalization. Corresponds to `x1Ref` and `RmsNorm(x)` in the formula, respectively. Null Not Supported. FLOAT32, FLOAT16, BFLOAT16 ND 1-8 ×
    x2Ref (aclTensor*) Input/Output Represents the second input for the Add operation and the result of the Add operation. Corresponds to `x2Ref` and `x` in the formula, respectively.
    • Null tensor is not supported.
    • The shape must be consistent with x1Ref.
    		 FLOAT32, FLOAT16, BFLOAT16 ND 1-8 ×
    gamma (aclTensor*) Input Represents the scaling factor (weight) of RmsNorm. Corresponds to `gamma` in the formula.
    • Null tensor is not supported.
    • The shape must be consistent with the last several dimensions of `x1Ref`, which are the dimensions of `x1Ref` to be normalized.
    		 FLOAT32, FLOAT16, BFLOAT16 ND 1-8 ×
    epsilon (double) Input A value added to the denominator to ensure numerical stability. Corresponds to `epsilon` in the corresponding formula. The recommended value is 1e-6. - - - -
    rstdOut (aclTensor*) Output Represents the standard deviation after normalization. Corresponds to `Rms(x)` in the formula.
    • Null Not Supported.
    • The shape is consistent with the first few dimensions of `x1Ref`, which represent the dimensions that do not require normalization. Example of the relationship between rstdOut shape, x1Ref shape, and gamma shape: if x1Ref shape is (2, 3, 4, 8) and gamma shape is (8), then rstdOut shape is (2, 3, 4, 1); if x1Ref shape is (2, 3, 4, 8) and gamma shape is (4, 8), then rstdOut shape is (2, 3, 1, 1).
    		 FLOAT32 ND 1-8 ×
    workspaceSize (uint64_t*) Output Returns the workspace size to be requested on the Device side. - - - - -
    executor (aclOpExecutor**) Output Returns the op executor, which contains the operator computation flow. - - - - -

    • Atlas inference products:

      • The data type of the parameters x1Ref, x2Ref, and gamma does not support BFLOAT16.

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

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

aclnnInplaceAddRmsNorm

  • Parameters

    Parameter Input/Output Description
    workspace Input Specifies the workspace memory address applied for on the Device side.
    workspaceSize Input Specifies the workspace size applied for on the Device side, obtained through the first-phase API aclnnInplaceAddRmsNormGetWorkspaceSize.
    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

  • Boundary value scenario description:

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

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

  • Deterministic computation:

    • aclnnInplaceAddRmsNorm 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_inplace_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);

    // Copy the host-side data to the device memory by calling 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 by calling aclCreateTensor.
    *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.
    // Set deviceId based on the actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle based on your own check requirements.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct the input and output. Custom construction is required based on the API interface.
    std::vector<int64_t> xShape = {2, 16};
    std::vector<int64_t> gammaShape = {16};
    std::vector<int64_t> rstdShape = {2, 1};
    void* x1RefDeviceAddr = nullptr;
    void* x2RefDeviceAddr = nullptr;
    void* gammaDeviceAddr = nullptr;
    void* rstdDeviceAddr = nullptr;
    aclTensor* x1Ref = nullptr;
    aclTensor* x2Ref = nullptr;
    aclTensor* gamma = nullptr;
    aclTensor* rstd = nullptr;
    std::vector<float> x1RefHostData = {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> x2RefHostData = {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> rstdHostData = {1, 2};
    float epsilon = 1e-6;

    // Create the x1Ref aclTensor.
    ret = CreateAclTensor(x1RefHostData, xShape, &x1RefDeviceAddr, aclDataType::ACL_FLOAT, &x1Ref);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the x2Ref aclTensor.
    ret = CreateAclTensor(x2RefHostData, xShape, &x2RefDeviceAddr, aclDataType::ACL_FLOAT, &x2Ref);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the gamma aclTensor.
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the rstd aclTensor.
    ret = CreateAclTensor(rstdHostData, rstdShape, &rstdDeviceAddr, aclDataType::ACL_FLOAT, &rstd);
    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 aclnnInplaceAddRmsNorm.
    ret = aclnnInplaceAddRmsNormGetWorkspaceSize(x1Ref, x2Ref, gamma, epsilon, rstd, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplaceAddRmsNormGetWorkspaceSize 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 aclnnInplaceAddRmsNorm.
    ret = aclnnInplaceAddRmsNorm(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplaceAddRmsNorm failed. ERROR: %d\n", ret); return ret);
    // 4. (Boilerplate) Wait for the task execution to finish synchronously.
    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(xShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(
        resultData.data(), resultData.size() * sizeof(resultData[0]), x1RefDeviceAddr, 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("x1Ref result[%ld] is: %f\n", i, resultData[i]);
    }

    std::vector<float> resultData1(size, 0);
    ret = aclrtMemcpy(
        resultData1.data(), resultData1.size() * sizeof(resultData1[0]), x2RefDeviceAddr, size * sizeof(resultData1[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("x2Ref result[%ld] is: %f\n", i, resultData1[i]);
    }

    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(x1Ref);
    aclDestroyTensor(x2Ref);
    aclDestroyTensor(gamma);
    aclDestroyTensor(rstd);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(x1RefDeviceAddr);
    aclrtFree(x2RefDeviceAddr);
    aclrtFree(gammaDeviceAddr);
    aclrtFree(rstdDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}
This site uses cookies to improve your browsing experience. By continuing to browse this site, you agree to our use of cookies.View Details