aclnnClippedSwiglu

📄 View Source Code

Applicable Products

Function

• Description: A Swish gated linear unit activation function with clipping, which implements the SwiGlu computation on x. Compared with aclnnSwiGlu, this API adds the following input parameters: groupIndex, alpha, limit, bias, and interleaved, to support the SwiGlu variant used by the GPT-OSS model and the grouping scenario used by the MoE model.

• Formula:

  For a given input tensor x with dimensions [a,b,c,d,e,f,g,…], aclnnClippedSwiglu performs the following computation:

  1. Concatenate the axes of x based on the input parameter dim. After concatenation, the dimensions become [pre, cut, after]. The cut axis is the axis that needs to be split into two tensors after concatenation. The split method can be front-back split or odd-even split. pre and after can be 1. For example, when dim is 3, the dimensions of x after concatenation are [abc, d, efg*…]. In addition, since the elements of the after axis are stored contiguously and the computation is element-wise, the cut axis and the after axis are merged, resulting in x with dimensions [pre, cut].

  2. Filter the pre axis of x based on the input parameter group_index. The formula is as follows:

     $$sum = \text{Sum}(group\_index)$$ $$x = x[ : sum, : ]$$

     Here, sum represents the sum of all elements in group_index. When group_index is not input, this step is skipped.

  3. Split x based on the input parameter interleaved using the following formula:

     When interleaved is true, an even-odd split is performed:

     $$A = x[ : , : : 2]$$ $$B = x[ : , 1 : : 2]$$

     When interleaved is false, a front-back split is performed:

     $$h = x.shape[1] // 2$$ $$A = x[ : , : h]$$ $$B = x[ : , h : ]$$

  4. Perform variant SwiGlu computation based on the input parameters alpha, limit, and bias, using the following formula:

     $$A = A.clamp(min=None, max=limit)$$ $$B = B.clamp(min=-limit, max=limit)$$ $$y\_glu = A * sigmoid(alpha * A)$$ $$y = y\_glu * (B + bias)$$

  5. Reshape the output tensor y to have the same number of dimensions as x before concatenation, with the size on the dim axis being half of that of x, and other dimensions identical to x.

Function Prototype

Each operator is divided into a two-phase API. The "aclnnClippedSwigluGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnClippedSwiglu" API is called to perform the computation.

aclnnStatus aclnnClippedSwigluGetWorkspaceSize(
    const aclTensor *x,
    const aclTensor *groupIndexOptional,
    int64_t          dim,
    double           alpha,
    double           limit,
    double           bias,
    bool             interleaved,
    const aclTensor *out,
    uint64_t        *workspaceSize,
    aclOpExecutor   **executor)

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

aclnnClippedSwigluGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-contiguous Tensor
  x (aclTensor*)	Input	The input x in the formula.	Null pointers are not supported. The dimension must be greater than 0, and the shape must be an even number in the dimension corresponding to the input parameter dim.	FLOAT, FLOAT16, BFLOAT16	ND	1-8	-
  groupIndexOptional (aclTensor*)	Input	Input group_index in the formula, indicating the grouping status.	Supports null pointers. If not null, the dimension must be 1D, and elements must be greater than or equal to 0. The i-th element represents the number of batches of x to be processed by the i-th group.	INT64	ND	1	-
  dim (int64_t)	Input	The input dim in the formula, indicating the dimension index along which x is concatenated and split.	The value range is [-x.dim(), x.dim()-1].	-	-	-	-
  alpha (double)	Input	Input alpha in the formula, representing the parameter used by the SwiGlu variant.	The recommended value is 1.702.	-	-	-	-
  limit (double)	Input	The input limit in the formula, representing the threshold value used by the SwiGlu variant.	The recommended value is 7.0.	-	-	-	-
  bias (double)	Input	The input bias in the formula, which represents the bias parameter used by the SwiGlu variant.	The recommended value is 1.0.	-	-	-	-
  interleaved (bool)	Input	Indicates whether to split x in an odd-even manner.	If set to true, performs an even-odd split on x. If set to false, performs a front-back split on x.	-	-	-	-
  out (aclTensor*)	Output	Output y in the formula.	Null pointer is not supported. The shape is half of x in the dimension corresponding to the input parameter dim, and is the same as x in other dimensions.	FLOAT, 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 op 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	The passed x or y is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of the input or output is not within the supported range.
  		The dimension of the input or output parameter is not within the supported range.
  		dim is not within the specified value range.

aclnnClippedSwiglu

• Parameters:

  Parameter	Input/Output	Description
  workspace	Input	Address of the workspace memory applied on the device side.
  workspaceSize	Input	Size of the workspace applied on the device side, obtained through the first-phase API aclnnClippedSwigluGetWorkspaceSize.
  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

Deterministic computation: aclnnClippedSwiglu defaults to a deterministic implementation. Non-deterministic implementation is not supported, and configuring deterministic computation does not take effect.

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_clipped_swiglu.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 shapeSize = 1;
    for (auto i : shape) {
        shapeSize *= i;
    }
    return shapeSize;
}

int Init(int32_t deviceId, aclrtStream* stream)
{
    // Boilerplate: ACL 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 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 the host 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 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 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.
    // Set 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);

    // 2. Construct the inputs and outputs. Customize the construction based on the API.
    std::vector<int64_t> xShape = {2, 32};
    std::vector<int64_t> groupIndexShape = {1};
    std::vector<int64_t> outShape = {2, 16};
    void* xDeviceAddr = nullptr;
    void* groupIndexDeviceAddr = nullptr;
    void* outDeviceAddr = nullptr;
    aclTensor* x = nullptr;
    aclTensor* groupIndex = nullptr;
    aclTensor* out = nullptr;
    std::vector<float> xHostData = {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21,
                             22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43,
                             44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63};
    std::vector<int64_t> groupIndexData = {1};
    std::vector<float> outHostData = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

    int dim = -1;
    float alpha = 1.0;
    float limit = 7.0;
    float bias = 1.702;
    bool interleaved = true;
    // Create the x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the groupIndex aclTensor.
    ret = CreateAclTensor(groupIndexData, groupIndexShape, &groupIndexDeviceAddr, aclDataType::ACL_INT64, &groupIndex);
    // Create the out aclTensor.
    ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
    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;
    // Call the first-phase API of aclnnClippedSwiglu.
    ret = aclnnClippedSwigluGetWorkspaceSize(
        x, groupIndex, dim, alpha, limit, bias, interleaved, out, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnClippedSwigluGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize computed 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 aclnnClippedSwiglu.
    ret = aclnnClippedSwiglu(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnClippedSwiglu 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. Modify based on the specific API definition.
    auto size = GetShapeSize(outShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(
        resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, 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(groupIndex);
    aclDestroyTensor(out);
    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(xDeviceAddr);
    aclrtFree(groupIndexDeviceAddr);
    aclrtFree(outDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}