/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.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_intf.h
 * \brief
 */
#ifndef __LIB_MATMUL_INTF_H__
#define __LIB_MATMUL_INTF_H__
#if __CCE_AICORE__ == 220
#include "lib/matmul/kernel_kfc.h"
#else
#include "lib/matmul/matmul.h"
#endif

#define REGIST_MATMUL_OBJ_STATIC REGIST_MATMUL_OBJ
#define REGIST_MATMUL_OBJ_STATIC_REMOTE REGIST_MATMUL_OBJ_REMOTE
namespace matmul {
#ifdef ASCENDC_CPU_DEBUG
#if __CCE_AICORE__ == 220
#ifdef ASCENDC_CUBE_ONLY
template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
    class MM_CB = MatmulCallBackFunc<nullptr, nullptr, nullptr>>
using Matmul = matmul::MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;
#else
template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
    class MM_CB = MatmulCallBackFunc<nullptr, nullptr, nullptr>>
using Matmul = matmul::MatmulClient<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;
#endif

#ifdef ASCENDC_CUBE_ONLY
template <class T, class... Args> __aicore__ static T* GetCurTiling(T* t, Args&&... b)
{
    return t;
}

template <class T, class... Args>
__aicore__ inline void InitCurObjSkip(TPipe* tpipe, T* a,
    Args&&... b)
{
    InitCurObj(tpipe, b...);
}

template <class T, class... Args>
__aicore__ inline void InitCurObj(TPipe* tpipe, T& a, Args&&... b)
{
    ASSERT(tpipe != nullptr && "tpipe cannot be nullptr");
    if constexpr (sizeof...(b) == 0) {
        SetTPipe(a, tpipe);
    } else {
        auto tiling = GetCurTiling(b...);
        a.SetSubBlockIdx(0);
        a.Init(tiling, tpipe);
        if constexpr (sizeof...(b) > 1) {
            InitCurObjSkip(tpipe, b...);
        }
    }
}

#define REGIST_MATMUL_OBJ(tpipe, workspace, ...) \
    InitCurObj(tpipe, __VA_ARGS__)

#else
#define REGIST_MATMUL_OBJ(tpipe, workspace, ...)              \
    if ASCEND_IS_AIC {                                        \
        AscendC::KfcServer server;                            \
        server.Init(workspace);                               \
        server.InitObj(tpipe, __VA_ARGS__);                   \
        while (server.isRun()) {                              \
            server.Run(__VA_ARGS__);                          \
        };                                                    \
        server.Quit();                                        \
        return;                                               \
    }                                                         \
    KfcCommClient __kfcClient__(workspace, GetSubBlockIdx()); \
    g_kfcClient = &__kfcClient__;                             \
    AscendC::SetMatrixKfc(tpipe, &__kfcClient__, 0, workspace, __VA_ARGS__)
#endif

#else

template <class T, class... Args> __aicore__ static T* GetCurTiling(T* t, Args&&... b)
{
    return t;
}

template <class T, class... Args>
__aicore__ inline void InitCurObjSkip(TPipe* tpipe, T* a,
    Args&&... b)
{
    InitCurObj(tpipe, b...);
}

template <class T, class... Args>
__aicore__ inline void InitCurObj(TPipe* tpipe, T& a, Args&&... b)
{
    ASSERT(tpipe != nullptr && "tpipe cannot be nullptr");
    if constexpr (sizeof...(b) == 0) {
        SetTPipe(a, tpipe);
    } else {
        auto tiling = GetCurTiling(b...);
        a.Init(tiling, tpipe);
        if constexpr (sizeof...(b) > 1) {
            InitCurObjSkip(tpipe, b...);
        }
    }
}

#define REGIST_MATMUL_OBJ(tpipe, workspace, ...) \
    InitCurObj(tpipe, __VA_ARGS__)

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
    class MM_CB = MatmulCallBackFunc<nullptr, nullptr, nullptr>>
using Matmul = matmul::MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;

#endif

#else

__aicore__ inline void clearWorkspace(__gm__ uint8_t* workspace)
{
    SetAtomicNone();
#if __CCE_AICORE__ == 220
    if ASCEND_IS_AIC {
        SetMaskNorm();
        SetLoadDataBoundary((uint64_t)0);
        SetLoadDataPaddingValue((uint64_t)0);
    } else {
        AscendCUtils::SetMask<uint64_t>((uint64_t)-1, (uint64_t)-1);
        SetMaskNorm();
    }
#endif

#ifdef __DAV_C220_CUBE__
    ClearWorkspaceImpl(workspace);
    uint16_t eventID = 3;
    NotifyEvent<PIPE_MTE3>(eventID);
#endif
}

#ifdef __DAV_C220_CUBE__
#ifdef ASCENDC_CUBE_ONLY
template <class T, class... Args> __aicore__ static T* GetCurTiling(T* t, Args&&... b)
{
    return t;
}

template <class T, class... Args>
__aicore__ inline void InitCurObjSkip(TPipe* tpipe, T* a,
    Args&&... b)
{
    InitCurObj(tpipe, b...);
}

template <class T, class... Args>
__aicore__ inline void InitCurObj(TPipe* tpipe, T& a, Args&&... b)
{
    ASSERT(tpipe != nullptr && "tpipe cannot be nullptr");
    if constexpr (sizeof...(b) == 0) {
        SetTPipe(a, tpipe);
    } else {
        auto tiling = GetCurTiling(b...);
        a.SetSubBlockIdx(0);
        a.Init(tiling, tpipe);
        if constexpr (sizeof...(b) > 1) {
            InitCurObjSkip(tpipe, b...);
        }
    }
}

#define REGIST_MATMUL_OBJ(tpipe, workspace, ...) \
    InitCurObj(tpipe, __VA_ARGS__)

#define REGIST_MATMUL_OBJ_REMOTE(tpipe, workspace, ...)

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
    class MM_CB = MatmulCallBackFunc<nullptr, nullptr, nullptr>>
using Matmul = matmul::MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;
#else
#define REGIST_MATMUL_OBJ(tpipe, workspace, ...) \
    if ASCEND_IS_AIC {           \
        AscendC::KfcServer server;               \
        server.Init(workspace);                  \
        server.InitObj(tpipe, __VA_ARGS__);      \
        while (server.isRun()) {                 \
            server.Run(__VA_ARGS__);             \
        };                                       \
        server.Quit();                           \
        return;                                  \
    }

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
    class MM_CB = MatmulCallBackFunc<nullptr>>
using Matmul = MatmulServiceAux<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;
#endif

#elif defined(__DAV_C220_VEC__)
#ifdef ASCENDC_CUBE_ONLY
#define REGIST_MATMUL_OBJ(tpipe, workspace, ...) \
    return
#else
#define REGIST_MATMUL_OBJ(tpipe, workspace, ...)                             \
    AscendC::KfcCommClient __kfcClient__(workspace, GetSubBlockIdx());       \
    g_kfcClient = &__kfcClient__;                                            \
    AscendC::SetMatrixKfc(tpipe, &__kfcClient__, 0, workspace, __VA_ARGS__); \
    WaitEvent(0x3)
#endif

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
    class MM_CB = MatmulCallBackFunc<nullptr, nullptr, nullptr>>
using Matmul = matmul::MatmulClient<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;
#else

template <class T, class... Args> __aicore__ static T* GetCurTiling(T* t, Args&&... b)
{
    return t;
}

template <class T, class... Args>
__aicore__ inline void InitCurObjSkip(TPipe* tpipe, T* a,
    Args&&... b)
{
    InitCurObj(tpipe, b...);
}

template <class T, class... Args>
__aicore__ inline void InitCurObj(TPipe* tpipe, T& a, Args&&... b)
{
    ASSERT(tpipe != nullptr && "tpipe cannot be nullptr");
    if constexpr (sizeof...(b) == 0) {
        SetTPipe(a, tpipe);
    } else {
        auto tiling = GetCurTiling(b...);
        a.Init(tiling, tpipe);
        if constexpr (sizeof...(b) > 1) {
            InitCurObjSkip(tpipe, b...);
        }
    }
}

#define REGIST_MATMUL_OBJ(tpipe, workspace, ...) \
    InitCurObj(tpipe, __VA_ARGS__)

#define REGIST_MATMUL_OBJ_REMOTE(tpipe, workspace, ...)

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
    class MM_CB = MatmulCallBackFunc<nullptr, nullptr, nullptr>>
using Matmul = matmul::MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;

#endif
#endif
}; // namespace matmul

#endif