Gemm

#include "kernel_operator.h"
namespace AscendC {
class KernelCubeGEMM {
public:
    __aicore__ inline KernelCubeGEMM() {}
    __aicore__ inline void Init(__gm__ uint8_t* fmGm, __gm__ uint8_t* weGm, __gm__ uint8_t* dstGm, uint32_t mInput,
        uint32_t kInput, uint32_t nInput, bool initVal, LoopMode mode)
    {
        m = mInput;
        k = kInput;
        n = nInput;

        initValue = initVal;
        loopMode = mode;

        featureMapA1Size = m * k;
        weightA1Size = k * n;
        dstCO1Size = m * n;

        roundm = DivCeil(m, 16) * 16;
        roundn = DivCeil(n, 16) * 16;
        roundk = DivCeil(k, c0Size) * c0Size;

        fmGlobal.SetGlobalBuffer((__gm__ half*)fmGm);
        weGlobal.SetGlobalBuffer((__gm__ half*)weGm);
        dstGlobal.SetGlobalBuffer((__gm__ float*)dstGm);

        pipe.InitBuffer(inQueueFmA1, 1, featureMapA1Size * sizeof(half));
        pipe.InitBuffer(inQueueWeB1, 1, weightA1Size * sizeof(half));
        pipe.InitBuffer(outQueueCO1, 1, dstCO1Size * sizeof(float));
        pipe.InitBuffer(outQueueUB, 1, dstCO1Size * sizeof(float));
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyUB();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn()
    {
        LocalTensor<half> featureMapA1 = inQueueFmA1.AllocTensor<half>();
        LocalTensor<half> weightB1 = inQueueWeB1.AllocTensor<half>();

        DataCopy(featureMapA1, fmGlobal, featureMapA1Size);
        DataCopy(weightB1, weGlobal, weightA1Size);

        inQueueFmA1.EnQue(featureMapA1);
        inQueueWeB1.EnQue(weightB1);
    }

    __aicore__ inline void Compute()
    {
        LocalTensor<half> featureMapA1 = inQueueFmA1.DeQue<half>();
        LocalTensor<half> weightB1 = inQueueWeB1.DeQue<half>();
        LocalTensor<float> dstCO1 = outQueueCO1.AllocTensor<float>();

        GemmTiling tilling = GetGemmTiling<half>(m, k, n);
        tilling.loopMode = loopMode;
        // 左矩阵形状为[m,k],右矩阵形状为[k,n]，计算结果搬出至GM，目的矩阵无需初始化
        Gemm(dstCO1, featureMapA1, weightB1, m, k, n, tilling, false, initValue);

        outQueueCO1.EnQue<float>(dstCO1);
        inQueueFmA1.FreeTensor(featureMapA1);
        inQueueWeB1.FreeTensor(weightB1);
    }

    __aicore__ inline void CopyUB()
    {
        LocalTensor<float> dstCO1 = outQueueCO1.DeQue<float>();
        LocalTensor<float> dstUB = outQueueUB.AllocTensor<float>();

        DataCopyParams dataCopyParams;
        dataCopyParams.blockCount = 1;
        dataCopyParams.blockLen = roundm * roundn * sizeof(float) / 1024;
        DataCopyEnhancedParams enhancedParams;
        enhancedParams.blockMode = BlockMode::BLOCK_MODE_MATRIX;

        DataCopy(dstUB, dstCO1, dataCopyParams, enhancedParams);

        outQueueUB.EnQue<float>(dstUB);
        outQueueCO1.FreeTensor(dstCO1);
    }

    __aicore__ inline void CopyOut()
    {
        LocalTensor<float> dstUB = outQueueUB.DeQue<float>();
        DataCopy(dstGlobal, dstUB, roundm * roundn);
        outQueueUB.FreeTensor(dstUB);
    }

private:
    TPipe pipe;
    // feature map queue
    TQue<QuePosition::A1, 1> inQueueFmA1;
    // weight queue
    TQue<QuePosition::B1, 1> inQueueWeB1;
    // dst queue
    TQue<QuePosition::CO1, 1> outQueueCO1;

    TQue<QuePosition::VECOUT, 1> outQueueUB;

    GlobalTensor<half> fmGlobal, weGlobal;
    GlobalTensor<float> dstGlobal;

    uint16_t m;
    uint16_t k;
    uint16_t n;
    uint32_t roundm, roundk, roundn;

    uint32_t c0Size = 16;
    bool initValue = false;
    LoopMode loopMode = LoopMode::MODE_NM;

    uint32_t featureMapA1Size, weightA1Size, dstCO1Size;
};
} // namespace AscendC

extern "C" __global__ __aicore__ void cube_gemm_simple_kernel(__gm__ uint8_t* fmGm, __gm__ uint8_t* weGm,
    __gm__ uint8_t* dstGm, uint32_t m, uint32_t k, uint32_t n, bool initValue, LoopMode mode)
{
    AscendC::KernelCubeGEMM op;
    // 上方示例结果入参为：m = 32, k = 64, n = 32, initValue = false, mode = LoopMode::MODE_NM
    op.Init(fmGm, weGm, dstGm, m, k, n, initValue, mode);
    op.Process();
}