/**
 * Copyright (c) 2026 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 debug.asc
 * \brief 向量加法算子调试样例，演示 msdebug 工具的断点调试与变量查看功能
 */

#include <cstdint>
#include <vector>
#include <iostream>
#include "acl/acl.h"
#include "kernel_operator.h"

#define CHECK_ACL(call, msg) do { \
    aclError ret = (call); \
    if (ret != ACL_SUCCESS) { \
        std::cerr << "[ERROR] " << msg << " failed, ret=" << ret << std::endl; \
        return std::vector<float>(); \
    } \
} while(0)

template <uint32_t tileSize>
__vector__ __global__ void vec_add_demo(__gm__ float* a, __gm__ float* b, __gm__ float* c)
{
    AscendC::InitSocState();

    AscendC::GlobalTensor<float> gmA, gmB, gmC;
    gmA.SetGlobalBuffer(a + block_idx * tileSize, tileSize);
    gmB.SetGlobalBuffer(b + block_idx * tileSize, tileSize);
    gmC.SetGlobalBuffer(c + block_idx * tileSize, tileSize);

    AscendC::LocalMemAllocator<AscendC::Hardware::UB> ubAlloc;
    AscendC::LocalTensor<float> bufA = ubAlloc.Alloc<float, tileSize>();
    AscendC::LocalTensor<float> bufB = ubAlloc.Alloc<float, tileSize>();
    AscendC::LocalTensor<float> bufC = ubAlloc.Alloc<float, tileSize>();

    AscendC::DataCopy(bufA, gmA, tileSize);
    AscendC::DataCopy(bufB, gmB, tileSize);
    AscendC::PipeBarrier<PIPE_ALL>();

    AscendC::Add(bufC, bufA, bufB, tileSize);
    AscendC::PipeBarrier<PIPE_ALL>();

    AscendC::DataCopy(gmC, bufC, tileSize);
    AscendC::PipeBarrier<PIPE_ALL>();
}

std::vector<float> launch_add(std::vector<float>& a, std::vector<float>& b)
{
    constexpr uint32_t numBlocks = 8;
    constexpr uint32_t tileSize = 2048;
    uint32_t totalLen = a.size();
    size_t totalBytes = totalLen * sizeof(float);
    float* devA = nullptr;
    float* devB = nullptr;
    float* devC = nullptr;
    uint8_t* hostC = nullptr;

    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId), "aclrtSetDevice");
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream), "aclrtCreateStream");

    CHECK_ACL(aclrtMalloc((void**)&devA, totalBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc devA");
    CHECK_ACL(aclrtMalloc((void**)&devB, totalBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc devB");
    CHECK_ACL(aclrtMalloc((void**)&devC, totalBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc devC");
    CHECK_ACL(aclrtMallocHost((void**)&hostC, totalBytes), "aclrtMallocHost hostC");

    CHECK_ACL(aclrtMemcpy(devA, totalBytes, a.data(), totalBytes, ACL_MEMCPY_HOST_TO_DEVICE), "aclrtMemcpy H2D devA");
    CHECK_ACL(aclrtMemcpy(devB, totalBytes, b.data(), totalBytes, ACL_MEMCPY_HOST_TO_DEVICE), "aclrtMemcpy H2D devB");

    vec_add_demo<tileSize><<<numBlocks, nullptr, stream>>>(devA, devB, devC);
    CHECK_ACL(aclrtSynchronizeStream(stream), "aclrtSynchronizeStream");

    CHECK_ACL(aclrtMemcpy(hostC, totalBytes, devC, totalBytes, ACL_MEMCPY_DEVICE_TO_HOST), "aclrtMemcpy D2H");
    std::vector<float> c((float*)hostC, (float*)(hostC + totalBytes));

    aclrtFree(devA);
    aclrtFree(devB);
    aclrtFree(devC);
    aclrtFreeHost(hostC);
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);

    return c;
}

int32_t main(int32_t argc, char* argv[])
{
    aclInit(nullptr);

    constexpr uint32_t totalLen = 8 * 2048;
    std::vector<float> a(totalLen);
    std::vector<float> b(totalLen);

    for (uint32_t i = 0; i < totalLen; ++i) {
        a[i] = i * 0.1f;
        b[i] = i * 0.2f;
    }

    std::vector<float> result = launch_add(a, b);
    if (result.empty()) {
        std::cerr << "[ERROR] launch_add failed" << std::endl;
        aclFinalize();
        return 1;
    }

    std::vector<float> golden(totalLen);
    for (uint32_t i = 0; i < totalLen; ++i) {
        golden[i] = a[i] + b[i];
    }

    bool pass = std::equal(golden.begin(), golden.end(), result.begin());
    std::cout << (pass ? "test pass!" : "test failed!") << std::endl;

    aclFinalize();
    return pass ? 0 : 1;
}