/**
* 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 cos.asc
 * \brief
 */

#include "acl/acl.h"
#include "data_utils.h"
#include "kernel_operator.h"

struct CosCustomTilingData {
    uint32_t totalLength;
    uint32_t tileNum;
};

void GenerateTilingData(uint8_t* tilingBuf, const uint32_t totalLength, const uint32_t coreNum)
{
    CosCustomTilingData* tiling = reinterpret_cast<CosCustomTilingData*>(tilingBuf);

    tiling->totalLength = totalLength;
    tiling->tileNum = coreNum;
}

namespace MyCustomKernel {
constexpr int32_t BUFFER_NUM = 2;

template <typename T>
class KernelCos {
public:
    __aicore__ inline KernelCos() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, CosCustomTilingData tilingData, AscendC::TPipe* pipeIn)
    {
        ASCENDC_ASSERT(AscendC::GetBlockNum() != 0, { KERNEL_LOG(KERNEL_ERROR, "block dim can not be zero!"); });
        this->tileNum = tilingData.tileNum;
        this->blockLength = tilingData.totalLength / AscendC::GetBlockNum();
        this->tileLength = this->blockLength / tileNum / BUFFER_NUM;

        xGm.SetGlobalBuffer((__gm__ T*)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        yGm.SetGlobalBuffer((__gm__ T*)y + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);

        pipe = pipeIn;
        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        AscendC::LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
        AscendC::Cos<T, false>(yLocal, xLocal, this->tileLength);
        outQueueY.EnQue<T>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<T> yLocal = outQueueY.DeQue<T>();
        AscendC::DataCopy(yGm[progress * this->tileLength], yLocal, this->tileLength);
        outQueueY.FreeTensor(yLocal);
    }

private:
    AscendC::TPipe* pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueY;
    AscendC::GlobalTensor<T> xGm;
    AscendC::GlobalTensor<T> yGm;

    uint32_t tileNum = 0;
    uint32_t blockLength = 0;
    uint32_t tileLength = 0;
};
} // namespace MyCustomKernel

__global__ __aicore__ void cos_custom(GM_ADDR x, GM_ADDR y, CosCustomTilingData tiling)
{
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    AscendC::TPipe pipe;
    MyCustomKernel::KernelCos<float> op;
    op.Init(x, y, tiling, &pipe);
    op.Process();
}

namespace {
constexpr uint32_t USED_CORE_NUM = 8;
constexpr uint32_t TILINGDATA_SIZE = 2;
constexpr uint32_t TOTAL_LENGTH = 8 * 2048;
} // namespace

static bool CompareResult(const void* outputData, uint32_t outSize)
{
    void* goldenData;
    aclrtMallocHost((void**)(&goldenData), outSize);
    size_t goldenSize = outSize;
    bool ret = ReadFile("./output/golden.bin", goldenSize, goldenData, goldenSize);
    if (ret) {
        printf("ReadFile golden.bin success!\n");
    } else {
        printf("test failed!\n");
        return false;
    }
    constexpr float EPS = 1e-4;
    int64_t wrongNum = 0;

    for (size_t i = 0; i < outSize / sizeof(float); i++) {
        float a = (reinterpret_cast<const float*>(outputData))[i];
        float b = (reinterpret_cast<const float*>(goldenData))[i];
        float ae = std::abs(a - b);
        float re = ae / std::abs(b);
        if (ae > EPS && re > EPS) {
            printf("CompareResult golden.bin failed output is %lf, golden is %lf\n", a, b);
            wrongNum++;
        }
    }
    aclrtFreeHost(goldenData);
    if (wrongNum != 0) {
        return false;
    } else {
        printf("CompareResult golden.bin success!\n");
        return true;
    }
}

int32_t main(int32_t argc, char* argv[])
{
    size_t tilingSize = TILINGDATA_SIZE * sizeof(uint32_t);
    size_t inputSize = TOTAL_LENGTH * sizeof(uint32_t);
    size_t outputSize = inputSize;

    aclInit(nullptr);
    int32_t deviceId = 0;
    aclrtSetDevice(deviceId);
    aclrtStream stream = nullptr;
    aclrtCreateStream(&stream);

    uint8_t *xHost, *yHost, *tiling;
    uint8_t *xDevice, *yDevice;

    aclrtMallocHost((void**)(&xHost), inputSize);
    aclrtMallocHost((void**)(&yHost), outputSize);
    aclrtMallocHost((void**)(&tiling), tilingSize);

    aclrtMalloc((void**)&xDevice, inputSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void**)&yDevice, outputSize, ACL_MEM_MALLOC_HUGE_FIRST);

    ReadFile("./input/input_x.bin", inputSize, xHost, inputSize);

    GenerateTilingData(tiling, TOTAL_LENGTH, USED_CORE_NUM);

    // Copy host memory to device memory
    aclrtMemcpy(xDevice, inputSize, xHost, inputSize, ACL_MEMCPY_HOST_TO_DEVICE);

    // Execute the kernel
    cos_custom<<<USED_CORE_NUM, nullptr, stream>>>(xDevice, yDevice, *reinterpret_cast<CosCustomTilingData*>(tiling));

    // Wait for the stop event to complete
    aclrtSynchronizeStream(stream);

    // Copy result to host memory and write to output file
    aclrtMemcpy(yHost, outputSize, yDevice, outputSize, ACL_MEMCPY_DEVICE_TO_HOST);
    WriteFile("./output/output.bin", yHost, outputSize);

    // Compare the result with the golden result
    bool goldenResult = true;
    goldenResult = CompareResult(yHost, outputSize);

    // Clean up memory
    aclrtFree(xDevice);
    aclrtFree(yDevice);
    aclrtFreeHost(xHost);
    aclrtFreeHost(yHost);
    aclrtFreeHost(tiling);

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

    if (goldenResult) {
        printf("test pass!\n");
    } else {
        printf("test failed!\n");
    }
    return 0;
}