kernel.asc

#include <cmath>
#include "kernel_operator.h"
#include "adv_api/activation/gelu.h"

using namespace AscendC;

/*
 * Gelu Kernel — Ascend C Direct Invocation
 * This file is #included — do NOT add main(), #pragma once, or include guards.
 *
 * TensorGroupInfo / TensorInfo (predefined):
 *   struct TensorInfo { const int64_t* shape; int64_t numDims; int32_t dtype; };
 *   struct TensorGroupInfo { const TensorInfo* tensors; int64_t numTensors; };
 *   dtype: 0=fp32 1=fp16 2=bf16 3=int8 4=int16 5=int32 6=int64 7=uint8 8=uint16 9=uint32 10=uint64 11=bool
 *   Usage: info_x.tensors[0].shape[0]  // first dim of first tensor in input x
 *
 * __global__ __vector__ void gelu_custom(...)
 * {
 *     // TODO: implement your kernel logic
 *     // Available APIs: AscendC::TPipe, AscendC::TQue, AscendC::DataCopy, AscendC::Add, ...
 * }
 */

// 每个 tile 处理的元素数（2201 UB=192KB）
// 动态 tile 策略：大数据（count > BIG_COUNT_THRESHOLD）用大 tile 提高带宽效率与 UB 利用率，
// 中小数据用小 tile 保证多核并行度（避免 tile 过大导致并行核数不足）。
// 阈值采用数据量级固定判断，不依赖 availableCoreNum（其值在真机不可靠）。
constexpr uint32_t SMALL_TILE_NUM = 2048;
constexpr uint32_t BIG_TILE_NUM = 4096;
constexpr uint32_t BIG_COUNT_THRESHOLD = 100000;

// fp16：用 AscendC::Gelu（真机验证通过，误差 < 1e-3）
class KernelGeluHalf
{
public:
    __aicore__ inline KernelGeluHalf() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t offset, uint32_t totalCount, uint32_t tileNum)
    {
        this->offset = offset;
        this->totalCount = totalCount;
        this->tileNum = tileNum;
        xGm.SetGlobalBuffer((__gm__ half *)x);
        yGm.SetGlobalBuffer((__gm__ half *)y);
        pipe.InitBuffer(inQueueX, 2, tileNum * sizeof(half));
        pipe.InitBuffer(outQueueY, 2, tileNum * sizeof(half));
    }

    __aicore__ inline void Process()
    {
        uint32_t loopCount = (totalCount + tileNum - 1) / tileNum;
        if (loopCount == 0)
        {
            return;
        }

        // 预取第一个 tile，让后续 CopyIn 与 CopyOut 的 DataCopy 重叠（MTE2/MTE3 并行）
        uint32_t firstCount = (totalCount < tileNum) ? totalCount : tileNum;
        CopyIn(0, firstCount);

        for (uint32_t i = 0; i < loopCount; i++)
        {
            uint32_t tileOffset = i * tileNum;
            uint32_t count = (totalCount - tileOffset < tileNum) ? (totalCount - tileOffset) : tileNum;

            Compute(count);

            if (i + 1 < loopCount)
            {
                uint32_t nextOffset = (i + 1) * tileNum;
                uint32_t nextCount = (totalCount - nextOffset < tileNum) ? (totalCount - nextOffset) : tileNum;
                CopyIn(nextOffset, nextCount);
            }

            CopyOut(tileOffset, count);
        }
    }

private:
    __aicore__ inline void CopyIn(uint32_t tileOffset, uint32_t count)
    {
        LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        DataCopy(xLocal, xGm[offset + tileOffset], count);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(uint32_t count)
    {
        LocalTensor<half> xLocal = inQueueX.DeQue<half>();
        LocalTensor<half> yLocal = outQueueY.AllocTensor<half>();
        Gelu<half>(yLocal, xLocal, count);
        outQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(uint32_t tileOffset, uint32_t count)
    {
        LocalTensor<half> yLocal = outQueueY.DeQue<half>();
        DataCopy(yGm[offset + tileOffset], yLocal, count);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<TPosition::VECIN, 2> inQueueX;
    TQue<TPosition::VECOUT, 2> outQueueY;
    GlobalTensor<half> xGm;
    GlobalTensor<half> yGm;
    uint32_t offset = 0;
    uint32_t totalCount = 0;
    uint32_t tileNum = 0;
};

// fp32：手写精确 erf（P/Q 有理近似，来自 CANN erf 实现），满足 < 1e-4
// GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2)))
class KernelGeluFloat
{
public:
    __aicore__ inline KernelGeluFloat() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t offset, uint32_t totalCount, uint32_t tileNum)
    {
        this->offset = offset;
        this->totalCount = totalCount;
        this->tileNum = tileNum;
        xGm.SetGlobalBuffer((__gm__ float *)x);
        yGm.SetGlobalBuffer((__gm__ float *)y);
        pipe.InitBuffer(inQueueX, 2, tileNum * sizeof(float));
        pipe.InitBuffer(calcQueue, 4, tileNum * sizeof(float));
        pipe.InitBuffer(outQueueY, 2, tileNum * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        uint32_t loopCount = (totalCount + tileNum - 1) / tileNum;
        if (loopCount == 0)
        {
            return;
        }

        // 预取第一个 tile，让后续 CopyIn 与 CopyOut 的 DataCopy 重叠（MTE2/MTE3 并行）
        uint32_t firstCount = (totalCount < tileNum) ? totalCount : tileNum;
        CopyIn(0, firstCount);

        for (uint32_t i = 0; i < loopCount; i++)
        {
            uint32_t tileOffset = i * tileNum;
            uint32_t count = (totalCount - tileOffset < tileNum) ? (totalCount - tileOffset) : tileNum;

            Compute(count);

            if (i + 1 < loopCount)
            {
                uint32_t nextOffset = (i + 1) * tileNum;
                uint32_t nextCount = (totalCount - nextOffset < tileNum) ? (totalCount - nextOffset) : tileNum;
                CopyIn(nextOffset, nextCount);
            }

            CopyOut(tileOffset, count);
        }
    }

private:
    __aicore__ inline void CopyIn(uint32_t tileOffset, uint32_t count)
    {
        LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        DataCopy(xLocal, xGm[offset + tileOffset], count);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(uint32_t count)
    {
        LocalTensor<float> x = inQueueX.DeQue<float>();
        LocalTensor<float> b0 = calcQueue.AllocTensor<float>(); // x2
        LocalTensor<float> b1 = calcQueue.AllocTensor<float>(); // t = clip(x/sqrt2)
        LocalTensor<float> b2 = calcQueue.AllocTensor<float>(); // P -> erf
        LocalTensor<float> b3 = calcQueue.AllocTensor<float>(); // Q

        // t = clip(x / sqrt(2), -3.92, 3.92)
        Muls<float>(b1, x, 0.70710678118654752440f, count);
        Maxs<float>(b1, b1, -3.92f, count);
        Mins<float>(b1, b1, 3.92f, count);

        // x2 = t * t
        Mul<float>(b0, b1, b1, count);

        // P(x2) = (((p4*x2+p3)*x2+p2)*x2+p1)*x2 + p0 （低阶有理近似，精度 ~2.5e-5 满足 1e-4）
        Muls<float>(b2, b0, -0.000028792113376f, count);
        Adds<float>(b2, b2, -0.00068839401572f, count);
        Mul<float>(b2, b2, b0, count);
        Adds<float>(b2, b2, 0.039120181773f, count);
        Mul<float>(b2, b2, b0, count);
        Adds<float>(b2, b2, 0.11699635615f, count);
        Mul<float>(b2, b2, b0, count);
        Adds<float>(b2, b2, 1.1283820967f, count);

        // Q(x2) = (((q4*x2+q3)*x2+q2)*x2+q1)*x2 + 1
        Muls<float>(b3, b0, -0.00033548044601f, count);
        Adds<float>(b3, b3, 0.0065209478643f, count);
        Mul<float>(b3, b3, b0, count);
        Adds<float>(b3, b3, 0.080147219569f, count);
        Mul<float>(b3, b3, b0, count);
        Adds<float>(b3, b3, 0.43707684811f, count);
        Mul<float>(b3, b3, b0, count);
        Adds<float>(b3, b3, 1.0f, count);

        // erf = t * P / Q
        Mul<float>(b2, b1, b2, count);
        Div<float>(b2, b2, b3, count);

        // y = 0.5 * x * (1 + erf)
        Adds<float>(b2, b2, 1.0f, count);

        LocalTensor<float> y = outQueueY.AllocTensor<float>();
        Muls<float>(y, x, 0.5f, count);
        Mul<float>(y, y, b2, count);

        calcQueue.FreeTensor(b0);
        calcQueue.FreeTensor(b1);
        calcQueue.FreeTensor(b2);
        calcQueue.FreeTensor(b3);
        outQueueY.EnQue(y);
        inQueueX.FreeTensor(x);
    }

    __aicore__ inline void CopyOut(uint32_t tileOffset, uint32_t count)
    {
        LocalTensor<float> yLocal = outQueueY.DeQue<float>();
        DataCopy(yGm[offset + tileOffset], yLocal, count);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<TPosition::VECIN, 2> inQueueX;
    TQue<TPosition::VECCALC, 4> calcQueue;
    TQue<TPosition::VECOUT, 2> outQueueY;
    GlobalTensor<float> xGm;
    GlobalTensor<float> yGm;
    uint32_t offset = 0;
    uint32_t totalCount = 0;
    uint32_t tileNum = 0;
};

// fp16 kernel 入口
extern "C" __global__ __vector__ void gelu_custom(GM_ADDR input_x, GM_ADDR output, uint32_t totalCount, uint32_t perBlock, uint32_t tileNum)
{
    uint32_t offset = (uint32_t)GetBlockIdx() * perBlock;
    if (offset >= totalCount)
    {
        return;
    }
    uint32_t count = (totalCount - offset < perBlock) ? (totalCount - offset) : perBlock;

    KernelGeluHalf op;
    op.Init(input_x, output, offset, count, tileNum);
    op.Process();
}

// fp32 kernel 入口
extern "C" __global__ __vector__ void gelu_custom_float(GM_ADDR input_x, GM_ADDR output, uint32_t totalCount, uint32_t perBlock, uint32_t tileNum)
{
    uint32_t offset = (uint32_t)GetBlockIdx() * perBlock;
    if (offset >= totalCount)
    {
        return;
    }
    uint32_t count = (totalCount - offset < perBlock) ? (totalCount - offset) : perBlock;

    KernelGeluFloat op;
    op.Init(input_x, output, offset, count, tileNum);
    op.Process();
}

extern "C" void run_kernel(GM_ADDR input_x, const TensorGroupInfo &info_input_x, GM_ADDR output, const TensorGroupInfo &info_output, int64_t availableCoreNum, aclrtStream stream)
{
    // 计算输入张量元素总数（沿所有维度相乘）
    const TensorInfo &inInfo = info_input_x.tensors[0];
    int64_t total = 1;
    for (int64_t d = 0; d < inInfo.numDims; ++d)
    {
        total *= inInfo.shape[d];
    }
    if (total <= 0)
    {
        return;
    }

    uint32_t count = (uint32_t)total;

    // 动态 tile：大数据用大 tile 提高带宽效率，中小数据用小 tile 保证并行核数（固定数据量阈值）
    uint32_t tileNum = (count > BIG_COUNT_THRESHOLD) ? BIG_TILE_NUM : SMALL_TILE_NUM;

    // 多核分块：按 tile 粒度均分，每个 block 的 GM 偏移为 tileNum 的整数倍（32B 对齐）
    uint32_t totalTiles = (count + tileNum - 1) / tileNum;
    uint32_t blockNum = 1;
    if (availableCoreNum > 1 && totalTiles > 1)
    {
        blockNum = (uint32_t)availableCoreNum;
        if (blockNum > totalTiles)
        {
            blockNum = totalTiles;
        }
    }
    uint32_t tilesPerBlock = (totalTiles + blockNum - 1) / blockNum;
    uint32_t perBlock = tilesPerBlock * tileNum;

    // dtype: 0=fp32 1=fp16，按类型分发到对应 kernel
    if (inInfo.dtype == 0)
    {
        gelu_custom_float<<<blockNum, nullptr, stream>>>(input_x, output, count, perBlock, tileNum);
    }
    else
    {
        gelu_custom<<<blockNum, nullptr, stream>>>(input_x, output, count, perBlock, tileNum);
    }
}