// matmul_prelu_custom.cpp
#include "kernel_operator.h"
#include "lib/matmul_intf.h"

using namespace matmul;

__aicore__ inline uint32_t Ceiling(uint32_t a, uint32_t b)
{
    return (a + b - 1) / b;
}

__aicore__ inline void CopyTiling(TCubeTiling *tiling, GM_ADDR tilingGM)
{
    uint32_t *ptr = reinterpret_cast<uint32_t *>(tiling);
    auto tiling32 = reinterpret_cast<__gm__ uint32_t *>(tilingGM);

    for (uint32_t i = 0; i < sizeof(TCubeTiling) / sizeof(uint32_t); i++, ptr++) {
        *ptr = *(tiling32 + i);
    }
    return;
}

template <typename aType, typename bType, typename cType, typename biasType> class MatmulPreluKernel {
public:
    __aicore__ inline MatmulPreluKernel(){};
    __aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c, cType alpha, GM_ADDR workspace,
                                const TCubeTiling &tiling, AscendC::TPipe *pipe);
    __aicore__ inline void Process(AscendC::TPipe *pipe);

    __aicore__ inline void MatmulCompute();
    __aicore__ inline void PreluCompute();
    __aicore__ inline void CopyOut(uint32_t count);
    __aicore__ inline void CalcOffset(int32_t blockIdx, const TCubeTiling &tiling, int32_t &offsetA, int32_t &offsetB,
                                      int32_t &offsetC, int32_t &offsetBias);

    Matmul<MatmulType<AscendC::TPosition::GM, CubeFormat::ND, aType>, MatmulType<AscendC::TPosition::GM, CubeFormat::ND, bType>,
           MatmulType<AscendC::TPosition::VECIN, CubeFormat::ND, cType>, MatmulType<AscendC::TPosition::GM, CubeFormat::ND, biasType>>
        matmulObj;

    AscendC::GlobalTensor<aType> aGlobal;
    AscendC::GlobalTensor<bType> bGlobal;
    AscendC::GlobalTensor<cType> cGlobal;
    AscendC::GlobalTensor<biasType> biasGlobal;
    AscendC::LocalTensor<cType> reluOutLocal;
    TCubeTiling tiling;
    AscendC::TQue<AscendC::QuePosition::VECOUT, 1> reluOutQueue;
    cType alpha;
};

template <typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulPreluKernel<aType, bType, cType, biasType>::Init(GM_ADDR a, GM_ADDR b, GM_ADDR bias,
                                                                              GM_ADDR c, cType alpha, GM_ADDR workspace,
                                                                              const TCubeTiling &tiling, AscendC::TPipe *pipe)
{
    aGlobal.Init(a, tiling.M * tiling.Ka, AscendC::TPosition::GM, CubeFormat::ND);
    bGlobal.Init(b, tiling.Kb * tiling.N, AscendC::TPosition::GM, CubeFormat::ND);
    biasGlobal.Init(bias, tiling.N, AscendC::TPosition::GM, CubeFormat::ND);
    cGlobal.Init(c, tiling.M * tiling.N, AscendC::TPosition::GM, CubeFormat::ND);
    this->alpha = alpha;
    this->tiling = tiling;
}

template <typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulPreluKernel<aType, bType, cType, biasType>::Process(AscendC::TPipe *pipe)
{
    int32_t offsetA, offsetB, offsetC, offsetBias;
    CalcOffset(pipe->GetBlockIdx(), tiling, offsetA, offsetB, offsetC, offsetBias);

    aGlobal.SetOffset(offsetA);
    bGlobal.SetOffset(offsetB);
    biasGlobal.SetOffset(offsetBias);
    cGlobal.SetOffset(offsetC);

    MatmulCompute();
    PreluCompute();
    CopyOut(tiling.singleCoreM * tiling.singleCoreN);
}

template <typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulPreluKernel<aType, bType, cType, biasType>::MatmulCompute()
{
    matmulObj.Compute(aGlobal, bGlobal, cGlobal, biasGlobal);
}

template <typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulPreluKernel<aType, bType, cType, biasType>::PreluCompute()
{
    for (int i = 0; i < tiling.singleCoreM * tiling.singleCoreN; ++i) {
        cType value = cGlobal[i];
        cGlobal[i] = value > 0 ? value : value * alpha;
    }
}

template <typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulPreluKernel<aType, bType, cType, biasType>::CopyOut(uint32_t count)
{
    for (int i = 0; i < count; ++i) {
        cGlobal[i] = reluOutLocal[i];
    }
}

template <typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void
MatmulPreluKernel<aType, bType, cType, biasType>::CalcOffset(int32_t blockIdx, const TCubeTiling &tiling,
                                                             int32_t &offsetA, int32_t &offsetB, int32_t &offsetC,
                                                             int32_t &offsetBias)
{
    auto mSingleBlocks = Ceiling(tiling.M, tiling.singleCoreM);
    auto mCoreIndx = blockIdx % mSingleBlocks;
    auto nCoreIndx = blockIdx / mSingleBlocks;

    offsetA = mCoreIndx * tiling.Ka * tiling.singleCoreM;
    offsetB = nCoreIndx * tiling.singleCoreN;
    offsetC = mCoreIndx * tiling.N * tiling.singleCoreM + nCoreIndx * tiling.singleCoreN;
    offsetBias = nCoreIndx * tiling.singleCoreN;
}

extern "C" __global__ __aicore__ void matmul_prelu_custom(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c,
                                                          float alpha, GM_ADDR workspace, GM_ADDR tilingGm)
{
    AscendC::TPipe pipe;
    TCubeTiling tiling;
    CopyTiling(&tiling, tilingGm);

    MatmulPreluKernel<half, half, float, float> kernel;
    kernel.Init(a, b, bias, c, alpha, workspace, tiling, &pipe);
    kernel.Process(&pipe);
}

// matmul_prelu_custom_tiling.cpp
#include <cassert>
#include <fstream>
#include <iostream>
#include <map>
#include <string>

#include "tiling/tiling_api.h"
#include "tiling/platform/platform_ascendc.h"
using namespace matmul_tiling;
using namespace std;

uint8_t *GetTilingBuf(optiling::TCubeTiling *tilingData)
{
    uint32_t tilingSize = tilingData->GetDataSize();
    uint8_t *buf = (uint8_t *)malloc(tilingSize);
    tilingData->SaveToBuffer(buf, tilingSize);
    return buf;
}

uint8_t *GenerateTiling(const char *socVersion)
{
    int M = 1024;
    int N = 640;
    int K = 256;

    TPosition leftPosition = TPosition::GM;
    CubeFormat leftFormat = CubeFormat::ND;
    DataType leftDtype = DataType::DT_FLOAT16;
    bool isTransA = false;

    TPosition rightPosition = TPosition::GM;
    CubeFormat rightFormat = CubeFormat::ND;
    DataType rightDtype = DataType::DT_FLOAT16;
    bool isTransB = false;

    TPosition resultPosition = TPosition::GM;
    CubeFormat resultFormat = CubeFormat::ND;
    DataType resultDtype = DataType::DT_FLOAT;

    TPosition biasPosition = TPosition::GM;
    CubeFormat biasFormat = CubeFormat::ND;
    DataType biasDtype = DataType::DT_FLOAT;
    bool isBias = true;

    int usedCoreNum = 2;
    int baseM = 256;
    int baseN = 128;

    optiling::TCubeTiling tilingData;
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(socVersion);
    MultiCoreMatmulTiling tilingApi(*ascendcPlatform);

    tilingApi.SetLeftMatrix(M, K, leftPosition, leftFormat, leftDtype, isTransA);
    tilingApi.SetRightMatrix(K, N, rightPosition, rightFormat, rightDtype, isTransB);
    tilingApi.SetResultMatrix(M, N, resultPosition, resultFormat, resultDtype);
    tilingApi.SetBiasMatrix(N, biasPosition, biasFormat, biasDtype, isBias);
    tilingApi.SetUsedCoreNum(usedCoreNum);
    tilingApi.SetBaseMN(baseM, baseN);

    int64_t res = tilingApi.GetTiling(tilingData);
    tilingData.set_stepM(1);
    tilingData.set_stepN(1);
    if (res == -1) {
        std::cout << "gen tiling failed" << std::endl;
    }
    return GetTilingBuf(&tilingData);
}