/**
 * 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 fixpipe_perf.asc
 * \brief Fixpipe（L0C 搬出）性能测试示例
 *
 * 测试矩阵计算的搬出路径：把 Cube 计算结果从 L0C Buffer 搬出到 L1 Buffer 或 Unified Buffer。

 *
 * 测试场景编号说明：
 * 场景 1-9：  适用于 Atlas A3/A2 训练/推理平台（dav-2201 架构）
 * 场景 11-19：适用于 Ascend 950PR/950DT 平台（dav-3510 架构）
 *
 * 场景 1： L0C->L1  DataCopy float -> half   （随路 F322F16）
 * 场景 2： L0C->L1  DataCopy float -> int8_t （随路 QF322B8_PRE）
 *
 * 场景 11：L0C->L1  DataCopy float -> half   （随路 F322F16）
 * 场景 12：L0C->L1  DataCopy float -> int8_t （随路 QF322B8_PRE）
 * 场景 13：L0C->UB  Fixpipe  float -> float 非双目标模式（dualDstCtl=0b00，仅 dav-3510）
 * 场景 14：L0C->UB  Fixpipe  float -> float 双目标模式按 M 拆分（dualDstCtl=0b01，仅 dav-3510）
 */
#ifdef ASCENDC_CPU_DEBUG
#include "cpu_debug_launch.h"
#endif
#include "acl/acl.h"
#include "kernel_operator.h"
#include <cstdio>
#include <cstdlib>

constexpr uint32_t BLOCK_CUBE = 16;
constexpr uint32_t C0_SIZE = 16;
constexpr uint64_t ADDR_0 = 0;
// 搬出目的矩阵在 L1 上的声明地址
constexpr uint64_t DST_L1_ADDR = 128 * 1024;

// L0C -> L1 搬出性能测试：源矩阵 C 为 float（L0C），目的类型 T 为 half / int8_t（L1）。
template <class T>
class KernelFixpipeL0c2L1 {
public:
    __aicore__ inline KernelFixpipeL0c2L1(uint32_t m, uint32_t k, uint32_t n)
    {
        mSize = m;
        kSize = k;
        nSize = n;

        aSizeAlignL0 = CeilAlign(mSize, BLOCK_CUBE) * CeilAlign(kSize, C0_SIZE);
        bSizeAlignL0 = CeilAlign(kSize, C0_SIZE) * CeilAlign(nSize, BLOCK_CUBE);
        cSizeAlignL0 = CeilAlign(mSize, BLOCK_CUBE) * CeilAlign(nSize, BLOCK_CUBE);
        dstSizeAlignL1 = CeilAlign(mSize, BLOCK_CUBE) * CeilAlign(nSize, BLOCK_CUBE);
    }

    __aicore__ inline void Process()
    {
        AscendC::LocalTensor<half> a2Local(AscendC::TPosition::A2, ADDR_0, aSizeAlignL0);
        AscendC::LocalTensor<half> b2Local(AscendC::TPosition::B2, ADDR_0, bSizeAlignL0);
        AscendC::LocalTensor<float> co1Local(AscendC::TPosition::CO1, ADDR_0, cSizeAlignL0);
        AscendC::LocalTensor<T> c1Local(AscendC::TPosition::C1, DST_L1_ADDR, dstSizeAlignL1);

        AscendC::PipeBarrier<PIPE_ALL>();

        AscendC::MmadParams mmadParams;
        mmadParams.m = mSize;
        mmadParams.n = nSize;
        mmadParams.k = kSize;
        mmadParams.cmatrixInitVal = true;
        AscendC::Mmad(co1Local, a2Local, b2Local, mmadParams);

        AscendC::PipeBarrier<PIPE_ALL>();

        // 搬出：L0C -> L1，随路量化/格式转换
        AscendC::DataCopyCO12DstParams intriParams;
        intriParams.nSize = nSize;
        intriParams.mSize = mSize;
        intriParams.srcStride = CeilAlign(mSize, BLOCK_CUBE);
        intriParams.dstStride = CeilAlign(mSize, BLOCK_CUBE);
        intriParams.reluPre = 0;
        intriParams.channelSplit = false;
        intriParams.nz2ndEn = false;
        if constexpr (AscendC::IsSameType<T, half>::value) {
            // float -> half：非量化 Cast
            intriParams.quantPre = QuantMode_t::F322F16;
        } else {
            // float -> int8_t：scalar 量化，需设置量化参数
            intriParams.quantPre = QuantMode_t::QF322B8_PRE;
            AscendC::SetFixpipePreQuantFlag(1);
        }
        AscendC::DataCopy(c1Local, co1Local, intriParams);

        AscendC::PipeBarrier<PIPE_ALL>();
    }

private:
    // 向上整除
    __aicore__ inline uint32_t CeilDiv(uint32_t numerator, uint32_t denominator)
    {
        return (numerator + denominator - 1) / denominator;
    }
    // 向上对齐
    __aicore__ inline uint32_t CeilAlign(uint32_t numerator, uint32_t denominator)
    {
        return (numerator + denominator - 1) / denominator * denominator;
    }

private:
    uint32_t mSize, kSize, nSize;
    uint32_t aSizeAlignL0, bSizeAlignL0, cSizeAlignL0, dstSizeAlignL1;
};

#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)
// L0C -> UB 搬出性能测试（仅 dav-3510）：源矩阵 C 为 float（L0C），输出 float 到 UB（双目标模式仅支持 float）。
// scenario==13：非双目标模式（dualDstCtl=0b00）；scenario==14：双目标按 M 拆分（dualDstCtl=0b01）。
// 需 __mix__ 混合编程：AIC 侧 Mmad + Fixpipe，AIV 侧将 UB 数据搬出到 GM。
constexpr AscendC::FixpipeConfig CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true};

class KernelFixpipeL0c2UB {
public:
    __aicore__ inline KernelFixpipeL0c2UB(uint32_t scenario, uint32_t m, uint32_t k, uint32_t n)
    {
        scenarioNum = scenario;
        mSize = m;
        kSize = k;
        nSize = n;

        aSizeAlignL0 = CeilAlign(mSize, BLOCK_CUBE) * CeilAlign(kSize, C0_SIZE);
        bSizeAlignL0 = CeilAlign(kSize, C0_SIZE) * CeilAlign(nSize, BLOCK_CUBE);
        cSizeAlignL0 = CeilAlign(mSize, BLOCK_CUBE) * CeilAlign(nSize, BLOCK_CUBE);
        cUBSize = mSize * nSize;
    }

    __aicore__ inline void Init() { cUB = AscendC::LocalTensor<float>(AscendC::TPosition::VECOUT, ADDR_0, cUBSize); }

    __aicore__ inline void Process()
    {
        if ASCEND_IS_AIC {
            AscendC::LocalTensor<half> a2Local(AscendC::TPosition::A2, ADDR_0, aSizeAlignL0);
            AscendC::LocalTensor<half> b2Local(AscendC::TPosition::B2, ADDR_0, bSizeAlignL0);
            AscendC::LocalTensor<float> co1Local(AscendC::TPosition::CO1, ADDR_0, cSizeAlignL0);

            AscendC::PipeBarrier<PIPE_ALL>();

            AscendC::MmadParams mmadParams;
            mmadParams.m = mSize;
            mmadParams.n = nSize;
            mmadParams.k = kSize;
            mmadParams.cmatrixInitVal = true;
            AscendC::Mmad(co1Local, a2Local, b2Local, mmadParams);

            AscendC::PipeBarrier<PIPE_ALL>();

            // 搬出：L0C Buffer -> UB（float -> float，双目标模式仅支持 float 输出，不支持随路转换）
            AscendC::FixpipeParamsArch3510<AscendC::CO2Layout::ROW_MAJOR> fixpipeParams;
            fixpipeParams.srcStride = CeilAlign(mSize, BLOCK_CUBE);
            if (scenarioNum == 14) {
                // 双目标模式按 M 拆分，M 须为 2 的倍数
                fixpipeParams.mSize = CeilAlign(mSize, 2);
                fixpipeParams.nSize = nSize;
                fixpipeParams.dstStride = nSize;
                fixpipeParams.dualDstCtl = 0b01;
            } else {
                // 非双目标模式
                fixpipeParams.mSize = mSize;
                fixpipeParams.nSize = nSize;
                fixpipeParams.dstStride = nSize;
                fixpipeParams.dualDstCtl = 0b00;
            }
            AscendC::Fixpipe<float, float, CFG_ROW_MAJOR_UB>(cUB, co1Local, fixpipeParams);

            AscendC::PipeBarrier<PIPE_ALL>();
        }
    }

private:
    __aicore__ inline uint32_t CeilAlign(uint32_t numerator, uint32_t denominator)
    {
        return (numerator + denominator - 1) / denominator * denominator;
    }

private:
    uint32_t scenarioNum;
    uint32_t mSize, kSize, nSize;
    uint32_t aSizeAlignL0, bSizeAlignL0, cSizeAlignL0, cUBSize;
    AscendC::LocalTensor<float> cUB;
};
#endif

__global__ __cube__ void fixpipe_perf_custom(
    __gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c, uint32_t scenario, uint32_t m, uint32_t k, uint32_t n)
{
    AscendC::InitSocState();
    // L0C -> L1 场景：half（场景 1/11）、int8_t（场景 2/12）
    if (scenario == 1 || scenario == 11) {
        KernelFixpipeL0c2L1<half> op(m, k, n);
        op.Process();
    } else if (scenario == 2 || scenario == 12) {
        KernelFixpipeL0c2L1<int8_t> op(m, k, n);
        op.Process();
    }
    AscendC::PipeBarrier<PIPE_ALL>();
}

#if defined(BUILD_DAV_3510)
// L0C -> UB 核函数（仅 dav-3510）：1 个 AIC + 2 个 AIV。
// 场景 13：非双目标模式；场景 14：双目标按 M 拆分。
__global__ __mix__(1, 2) void fixpipe_perf_ub_custom(
    __gm__ uint8_t* c, uint32_t scenario, uint32_t m, uint32_t k, uint32_t n)
{
    AscendC::InitSocState();
    KernelFixpipeL0c2UB op(scenario, m, k, n);
    op.Init();
    op.Process();
    AscendC::PipeBarrier<PIPE_ALL>();
}
#endif

int32_t main(int32_t argc, char* argv[])
{
    if (argc != 5) {
        std::printf("Usage: %s SCENARIO_NUM M K N\n", argv[0]);
        return 1;
    }

    uint32_t scenario = static_cast<uint32_t>(std::strtoul(argv[1], nullptr, 10));
    uint32_t m = static_cast<uint32_t>(std::strtoul(argv[2], nullptr, 10));
    uint32_t k = static_cast<uint32_t>(std::strtoul(argv[3], nullptr, 10));
    uint32_t n = static_cast<uint32_t>(std::strtoul(argv[4], nullptr, 10));

    if (m == 0 || k == 0 || n == 0) {
        std::printf("M, K and N must be positive integers.\n");
        return 1;
    }

    size_t aFileSize = static_cast<size_t>(m) * k * sizeof(half);
    size_t bFileSize = static_cast<size_t>(k) * n * sizeof(half);
    size_t cFileSize = static_cast<size_t>(m) * n * sizeof(float);

    uint32_t numBlocks = 1;

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

    uint8_t* aDevice = nullptr;
    uint8_t* bDevice = nullptr;
    uint8_t* cDevice = nullptr;
    aclrtMalloc((void**)&aDevice, aFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void**)&bDevice, bFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void**)&cDevice, cFileSize, ACL_MEM_MALLOC_HUGE_FIRST);

    if (scenario == 13 || scenario == 14) {
        // L0C -> UB 场景，仅 dav-3510，使用 __mix__ 核函数
#if defined(BUILD_DAV_3510)
        fixpipe_perf_ub_custom<<<numBlocks, nullptr, stream>>>(cDevice, scenario, m, k, n);
#endif
    } else {
        // L0C -> L1 场景
        fixpipe_perf_custom<<<numBlocks, nullptr, stream>>>(aDevice, bDevice, cDevice, scenario, m, k, n);
    }
    aclrtSynchronizeStream(stream);

    aclrtFree(aDevice);
    aclrtFree(bDevice);
    aclrtFree(cDevice);

    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}