/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/* !
 * \file matmul_splitk.asc
 */

#include "data_utils.h"
#include "kernel_tiling/kernel_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/tiling_api.h"
#include "acl/acl.h"
#include "kernel_operator.h"
#define ASCENDC_CUBE_ONLY
#include "lib/matmul_intf.h"

__aicore__ __inline__ constexpr uint32_t DivCeil(uint32_t a, uint32_t b) { return (a + b - 1) / b; }

constexpr uint32_t M = 16;
constexpr uint32_t N = 16;
constexpr uint32_t K = 1024;
constexpr bool isTransA = false;
constexpr bool isTransB = false;
constexpr bool isBias = false;

__aicore__ inline void CopyTiling(TCubeTiling* tiling, __gm__ uint8_t* 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 MatmulKernel {
public:
    __aicore__ inline MatmulKernel(){};
    __aicore__ inline void Init(
        __gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* bias, __gm__ uint8_t* c, const TCubeTiling& tiling);
    __aicore__ inline void Process();
    AscendC::Matmul<
        AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, AType, isTransA>,
        AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BType, isTransB>,
        AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, CType>,
        AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BiasType>, CFG_NORM>
        matmulObj;

private:
    __aicore__ inline void CalcOffset(
        int32_t blockIdx, int32_t& offsetA, int32_t& offsetB, int32_t& offsetC, int32_t& offsetBias);

    AscendC::GlobalTensor<AType> aGlobal;
    AscendC::GlobalTensor<BType> bGlobal;
    AscendC::GlobalTensor<CType> cGlobal;
    AscendC::GlobalTensor<BiasType> biasGlobal;
    TCubeTiling tiling;
};

template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulKernel<AType, BType, CType, BiasType>::Init(
    __gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* bias, __gm__ uint8_t* c, const TCubeTiling& tiling)
{
    this->tiling = tiling;

    aGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ AType*>(a), tiling.M * tiling.Ka);
    bGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ BType*>(b), tiling.Kb * tiling.N);
    cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ CType*>(c), tiling.M * tiling.N);
    biasGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ BiasType*>(bias), tiling.N);

    // clear gm
    Fill(cGlobal, tiling.M * tiling.N, (CType)0);

    int32_t offsetA = 0;
    int32_t offsetB = 0;
    int32_t offsetC = 0;
    int32_t offsetBias = 0;
    CalcOffset(AscendC::GetBlockIdx(), offsetA, offsetB, offsetC, offsetBias);
    aGlobal = aGlobal[offsetA];
    bGlobal = bGlobal[offsetB];
    cGlobal = cGlobal[offsetC];
    biasGlobal = biasGlobal[offsetBias];

    if (GetSysWorkSpacePtr() == nullptr) {
        return;
    }
}

template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulKernel<AType, BType, CType, BiasType>::Process()
{
    matmulObj.SetTensorA(aGlobal, isTransA);
    matmulObj.SetTensorB(bGlobal, isTransB);
    if (tiling.isBias) {
        matmulObj.SetBias(biasGlobal);
    }
    uint8_t enAtomic = 1; // set AtomicAdd
    matmulObj.IterateAll(cGlobal, enAtomic);
    matmulObj.End();
}

template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulKernel<AType, BType, CType, BiasType>::CalcOffset(
    int32_t blockIdx, int32_t& offsetA, int32_t& offsetB, int32_t& offsetC, int32_t& offsetBias)
{
    const TCubeTiling& tiling = this->tiling;
    auto temp0 = DivCeil(tiling.M, tiling.singleCoreM);
    auto temp1 = DivCeil(tiling.N, tiling.singleCoreN);
    auto temp2 = DivCeil(tiling.Ka, tiling.singleCoreK);

    auto divideKCoreNum = tiling.usedCoreNum / temp2;

    auto mCoreIndex = (blockIdx % divideKCoreNum) % temp0;
    auto nCoreIndex = (blockIdx % divideKCoreNum) / temp0;
    auto subKIndex = blockIdx / divideKCoreNum;

    offsetA = mCoreIndex * tiling.Ka * tiling.singleCoreM + subKIndex * tiling.singleCoreK;
    offsetB = subKIndex * tiling.singleCoreK * tiling.N + nCoreIndex * tiling.singleCoreN;
    offsetC = mCoreIndex * tiling.N * tiling.singleCoreM + nCoreIndex * tiling.singleCoreN;
    offsetBias = nCoreIndex * tiling.singleCoreN;
}

__global__ __cube__ void matmul_splitk_custom(
    __gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c, __gm__ uint8_t* workspace, __gm__ uint8_t* tilingGm)
{
    TCubeTiling tiling;
    CopyTiling(&tiling, tilingGm);

    MatmulKernel<half, half, float, float> matmulKernel;
    AscendC::TPipe pipe;
    REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), matmulKernel.matmulObj, &tiling);

    matmulKernel.Init(a, b, nullptr, c, tiling);
    matmulKernel.Process();
}

void GenerateTiling(platform_ascendc::PlatformAscendC* ascendcPlatform, uint8_t* tilingBuf)
{
    optiling::TCubeTiling tilingData;
    matmul_tiling::MultiCoreMatmulTiling tilingApi(*ascendcPlatform);

    tilingApi.SetDim(ascendcPlatform->GetCoreNumAic());
    tilingApi.SetAType(
        matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16, isTransA);
    tilingApi.SetBType(
        matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16, isTransB);
    tilingApi.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
    tilingApi.SetBiasType(
        matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);

    tilingApi.SetOrgShape(M, N, K);
    tilingApi.SetShape(M, N, K);
    tilingApi.EnableBias(isBias);
    tilingApi.SetBufferSpace(-1, -1, -1);
    // tiling enable split K
    tilingApi.EnableMultiCoreSplitK(true);

    int64_t res = tilingApi.GetTiling(tilingData); // Get matmul tiling data.
    if (res == -1) {
        std::cout << "gen tiling failed" << std::endl;
        return;
    }
    uint32_t tcubeTilingSize = tilingData.GetDataSize();
    tilingData.SaveToBuffer(tilingBuf, tcubeTilingSize);
}

int32_t main(int32_t argc, char* argv[])
{
#if defined(ASCENDC_CPU_DEBUG) && (__NPU_ARCH__ == 2201)
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance("Ascend910B1");
#elif defined(ASCENDC_CPU_DEBUG) && (__NPU_ARCH__ == 3510)
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance("Ascend950PR_9589");
#else
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
#endif

    size_t aFileSize = static_cast<size_t>(M * K) * sizeof(uint16_t); // uint16_t represent half
    size_t bFileSize = static_cast<size_t>(K * N) * sizeof(uint16_t); // uint16_t represent half
    size_t cFileSize = static_cast<size_t>(M * N) * sizeof(float);

    size_t userWorkspaceSize = 0;
    size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform->GetLibApiWorkSpaceSize());
    size_t workspaceSize = userWorkspaceSize + systemWorkspaceSize;

    // matmul TCubeTiling
    size_t tilingFileSize = sizeof(TCubeTiling);
    uint8_t* tilingBuf = (uint8_t*)malloc(tilingFileSize);
    GenerateTiling(ascendcPlatform, tilingBuf);

    uint32_t numBlocks = reinterpret_cast<TCubeTiling*>(tilingBuf)->usedCoreNum;

    int32_t deviceId = 0;
    aclrtStream stream = nullptr;
    aclrtContext context;

    aclInit(nullptr);
    aclrtSetDevice(deviceId);
    aclrtCreateContext(&context, deviceId);
    aclrtCreateStream(&stream);

    uint8_t* aHost;
    uint8_t* aDevice;
    aclrtMallocHost((void**)(&aHost), aFileSize);
    aclrtMalloc((void**)&aDevice, aFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
    ReadFile("./input/x1_gm.bin", aFileSize, aHost, aFileSize);
    aclrtMemcpy(aDevice, aFileSize, aHost, aFileSize, ACL_MEMCPY_HOST_TO_DEVICE);

    uint8_t* bHost;
    uint8_t* bDevice;
    aclrtMallocHost((void**)(&bHost), bFileSize);
    aclrtMalloc((void**)&bDevice, bFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
    ReadFile("./input/x2_gm.bin", bFileSize, bHost, bFileSize);
    aclrtMemcpy(bDevice, bFileSize, bHost, bFileSize, ACL_MEMCPY_HOST_TO_DEVICE);

    uint8_t* cHost;
    uint8_t* cDevice;
    aclrtMallocHost((void**)(&cHost), cFileSize);
    aclrtMalloc((void**)&cDevice, cFileSize, ACL_MEM_MALLOC_HUGE_FIRST);

    uint8_t* workspaceDevice;
    aclrtMalloc((void**)&workspaceDevice, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);

    uint8_t* tilingHost;
    uint8_t* tilingDevice;
    aclrtMallocHost((void**)(&tilingHost), tilingFileSize);
    aclrtMalloc((void**)&tilingDevice, tilingFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMemcpy(tilingHost, tilingFileSize, tilingBuf, tilingFileSize, ACL_MEMCPY_HOST_TO_HOST);
    aclrtMemcpy(tilingDevice, tilingFileSize, tilingHost, tilingFileSize, ACL_MEMCPY_HOST_TO_DEVICE);

    matmul_splitk_custom<<<numBlocks, 0, stream>>>(aDevice, bDevice, cDevice, workspaceDevice, tilingDevice);
    aclrtSynchronizeStream(stream);

    aclrtMemcpy(cHost, cFileSize, cDevice, cFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
    WriteFile("./output/output.bin", cHost, cFileSize);

    aclrtFree(aDevice);
    aclrtFreeHost(aHost);
    aclrtFree(bDevice);
    aclrtFreeHost(bHost);
    aclrtFree(workspaceDevice);
    aclrtFree(tilingDevice);
    aclrtFreeHost(tilingHost);
    aclrtFree(cDevice);
    aclrtFreeHost(cHost);

    aclrtDestroyStream(stream);
    aclrtDestroyContext(context);
    aclrtResetDevice(deviceId);
    aclFinalize();

    free(tilingBuf);
    return 0;
}