// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include <miopen/config.h>
#include <miopen/convolution.hpp>
#include <miopen/db.hpp>
#include <miopen/env.hpp>
#include <miopen/gcn_asm_utils.hpp>
#include <miopen/mlo_internal.hpp>
#include <miopen/conv/solvers.hpp>
#include <miopen/readonlyramdb.hpp>
#include <miopen/datatype.hpp>
#include <miopen/version.h>
#include <miopen/stringutils.hpp>
#include <miopen/hip_build_utils.hpp>
#include <miopen/any_solver.hpp>

#include <cmath>
#include <cstring>
#include <iomanip>
#include <memory>
#include <sstream>
#include <unordered_map>

#if MIOPEN_ENABLE_SQLITE
#include <miopen/sqlite_db.hpp>
#endif

miopen::PerformanceDb miopen::GetDb(const miopen::ExecutionContext& ctx)
{
    return {DbKinds::PerfDb, ctx.GetPerfDbPath(), ctx.GetUserPerfDbPath()};
}

static auto GetGemmSolvers()
{
    return miopen::solver::SolverContainer<miopen::solver::conv::GemmFwd1x1_0_1,
                                           miopen::solver::conv::GemmFwd1x1_0_1_int8,
                                           miopen::solver::conv::GemmFwd1x1_0_2,
                                           miopen::solver::conv::GemmFwdRest,

                                           miopen::solver::conv::GemmBwd1x1_stride1,
                                           miopen::solver::conv::GemmBwd1x1_stride2,
                                           miopen::solver::conv::GemmBwdRest,

                                           miopen::solver::conv::GemmWrw1x1_stride1,
                                           miopen::solver::conv::GemmWrwUniversal>{};
}

static auto GetDirectSolvers()
{
    return miopen::solver::SolverContainer<miopen::solver::conv::ConvAsm3x3U,
                                           miopen::solver::conv::ConvAsm1x1U,
                                           miopen::solver::conv::ConvAsm1x1UV2,
                                           miopen::solver::conv::ConvAsm5x10u2v2f1,
                                           miopen::solver::conv::ConvAsm7x7c3h224w224k64u2v2p3q3f1,
                                           miopen::solver::conv::ConvAsm5x10u2v2b1,
                                           miopen::solver::conv::ConvHipDirectFwd11x11,
                                           miopen::solver::conv::ConvHipDirectFwd,
                                           miopen::solver::conv::ConvDirectNaiveConvFwd,
                                           miopen::solver::conv::ConvDirectNaiveConvBwd,
                                           miopen::solver::conv::ConvDirectNaiveConvWrw,
                                           miopen::solver::conv::ConvDepthwiseFwd2D,
                                           miopen::solver::conv::ConvDepthwiseFwd3D>{};
}

static auto GetImplicitGemmSolvers()
{
    return miopen::solver::SolverContainer<
        miopen::solver::conv::ConvHipImplicitGemmForwardV4R5Xdlops,
        miopen::solver::conv::ConvHipImplicitGemmForwardV4R4Xdlops,
        miopen::solver::conv::ConvHipImplicitGemmForwardV4R4Xdlops_Padded_Gemm,
        miopen::solver::conv::ConvHipImplicitGemmBwdDataV4R1Xdlops,
        miopen::solver::conv::ConvHipImplicitGemmBwdDataV1R1Xdlops,
        miopen::solver::conv::ConvHipImplicitGemmV4R1Fwd,
        miopen::solver::conv::ConvHipImplicitGemmV4R4Fwd,
        miopen::solver::conv::ConvMlirIgemmFwdXdlops,
        miopen::solver::conv::ConvMlirIgemmFwd,
        miopen::solver::conv::ConvMlirIgemmBwdXdlops,
        miopen::solver::conv::ConvMlirIgemmBwd,
        miopen::solver::conv::ConvHipImplicitGemmBwdDataV1R1,
        miopen::solver::conv::ConvHipImplicitGemmBwdDataV4R1,
        miopen::solver::conv::ConvAsmImplicitGemmV4R1DynamicFwd_1x1,
        miopen::solver::conv::ConvAsmImplicitGemmV4R1DynamicFwd,
        miopen::solver::conv::ConvAsmImplicitGemmV4R1DynamicBwd,
        miopen::solver::conv::ConvAsmImplicitGemmGTCDynamicFwdXdlops,
        miopen::solver::conv::ConvAsmImplicitGemmGTCDynamicBwdXdlops,
        miopen::solver::conv::ConvAsmImplicitGemmGTCDynamicFwdXdlopsNHWC,
        miopen::solver::conv::ConvAsmImplicitGemmGTCDynamicBwdXdlopsNHWC,
        miopen::solver::conv::ConvCkIgemmFwdV6r1DlopsNchw,
#if MIOPEN_BACKEND_HIP
        miopen::solver::conv::ConvHipImplicitGemmGroupFwdXdlops,
        miopen::solver::conv::ConvHipImplicitGemmGroupBwdXdlops,
        miopen::solver::conv::ConvHipImplicitGemm3DGroupFwdXdlops,
        miopen::solver::conv::ConvHipImplicitGemm3DGroupBwdXdlops,
#endif // MIOPEN_BACKEND_HIP
        miopen::solver::conv::ConvAsmImplicitGemmGTCDynamicFwdDlopsNCHWC>{};
}

static auto GetWindogradSolvers()
{
    return miopen::solver::SolverContainer<
        miopen::solver::conv::ConvBinWinograd3x3U,
        miopen::solver::conv::ConvBinWinoRxS<3, 2>,
        miopen::solver::conv::ConvBinWinoRxS<2, 3>,
        miopen::solver::conv::ConvBinWinogradRxSf2x3g1,
        miopen::solver::conv::ConvBinWinogradRxS,
        miopen::solver::conv::ConvMPBidirectWinograd<2, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd<3, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd<4, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd<5, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd<6, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd_xdlops<2, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd_xdlops<3, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd_xdlops<4, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd_xdlops<5, 3>,
        miopen::solver::conv::ConvMPBidirectWinograd_xdlops<6, 3>,
        miopen::solver::conv::ConvWinoFuryRxS<2, 3>,
        miopen::solver::conv::ConvWinoRageRxS<2, 3>,
        // Transposed Winograd solvers for NHWC layout support
        miopen::solver::conv::TransposedConvBinWinograd3x3U,
        miopen::solver::conv::TransposedConvBinWinogradRxS,
        miopen::solver::conv::TransposedConvBinWinogradRxSf2x3g1,
        miopen::solver::conv::TransposedConvWinoFuryRxS<2, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd<2, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd<3, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd<4, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd<5, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd<6, 3>,
        miopen::solver::conv::TransposedConvBinWinoRxS<2, 3>,
        miopen::solver::conv::TransposedConvBinWinoRxS<3, 2>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd_xdlops<2, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd_xdlops<3, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd_xdlops<4, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd_xdlops<5, 3>,
        miopen::solver::conv::TransposedConvMPBidirectWinograd_xdlops<6, 3>>{};
}

static auto GetImplicitGemmWrWSolvers()
{
    return miopen::solver::SolverContainer<
        miopen::solver::conv::ConvHipImplicitGemmWrwV4R4Xdlops,
        miopen::solver::conv::ConvHipImplicitGemmWrwV4R4Xdlops_Padded_Gemm,
        miopen::solver::conv::ConvHipImplicitGemmV4R1WrW,
        miopen::solver::conv::ConvHipImplicitGemmV4R4WrW,
        miopen::solver::conv::ConvAsmImplicitGemmV4R1DynamicWrw,
        miopen::solver::conv::ConvMlirIgemmWrWXdlops,
        miopen::solver::conv::ConvMlirIgemmWrW,
        miopen::solver::conv::ConvAsmImplicitGemmGTCDynamicWrwXdlops,
#if MIOPEN_BACKEND_HIP
        miopen::solver::conv::ConvHipImplicitGemmGroupWrwXdlops,
        miopen::solver::conv::ConvHipImplicitGemm3DGroupWrwXdlops,
#endif // MIOPEN_BACKEND_HIP
        miopen::solver::conv::ConvAsmImplicitGemmGTCDynamicWrwXdlopsNHWC>{};
}

static auto GetWindogradWrWSolvers()
{
    return miopen::solver::SolverContainer<
        miopen::solver::conv::ConvBinWinogradRxS,
        miopen::solver::conv::ConvBinWinoRxS<3, 2>,
        miopen::solver::conv::ConvBinWinoRxS<2, 3>,
        miopen::solver::conv::ConvBinWinogradRxSf2x3g1,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<3, 2>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<3, 3>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<3, 4>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<3, 5>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<3, 6>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<7, 2>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<7, 3>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<7, 3, 1, 1>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<7, 2, 1, 1>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<1, 1, 7, 2>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<1, 1, 7, 3>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<5, 3>,
        miopen::solver::conv::ConvWinograd3x3MultipassWrW<5, 4>,
        miopen::solver::conv::ConvWinoFuryRxS<2, 3>,
        miopen::solver::conv::ConvWinoRageRxS<2, 3>,
        // Transposed Winograd solvers for NHWC layout support
        miopen::solver::conv::TransposedConvBinWinograd3x3U,
        miopen::solver::conv::TransposedConvBinWinogradRxS,
        miopen::solver::conv::TransposedConvBinWinogradRxSf2x3g1,
        miopen::solver::conv::TransposedConvWinoFuryRxS<2, 3>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<3, 2>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<3, 3>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<3, 4>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<3, 5>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<3, 6>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<7, 2>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<7, 3>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<1, 1, 7, 2>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<1, 1, 7, 3>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<7, 2, 1, 1>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<7, 3, 1, 1>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<5, 3>,
        miopen::solver::conv::TransposedConvWinograd3x3MultipassWrW<5, 4>,
        miopen::solver::conv::TransposedConvBinWinoRxS<2, 3>,
        miopen::solver::conv::TransposedConvBinWinoRxS<3, 2>>{};
}

static auto GetBwdWrW2DSolvers()
{
    return miopen::solver::SolverContainer<miopen::solver::conv::ConvAsmBwdWrW1x1,
                                           miopen::solver::conv::ConvAsmBwdWrW3x3,
                                           miopen::solver::conv::ConvHipBwdWrW53,
                                           miopen::solver::conv::ConvDirectNaiveConvFwd,
                                           miopen::solver::conv::ConvDirectNaiveConvBwd,
                                           miopen::solver::conv::ConvDirectNaiveConvWrw>{};
}

static auto GetFFTSolvers() { return miopen::solver::SolverContainer<miopen::solver::conv::fft>{}; }

miopen::DbGetter::DbGetter(std::function<PerformanceDb()>&& init_) : init(init_) {}

auto miopen::DbGetter::operator()() -> PerformanceDb&
{
    if(!db)
        db.emplace(init());

    return *db;
}

auto miopen::MakeConvDbGetter(const ExecutionContext& ctx) -> DbGetter
{
    return DbGetter{[&]() { return GetDb(ctx); }};
}

std::vector<miopen::solver::ConvSolution>
FindAllGemmSolutions(const miopen::ExecutionContext& ctx,
                     const miopen::conv::ProblemDescription& problem,
                     const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetGemmSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<std::pair<std::string, size_t>>
AllGemmWorkspaceSize(const miopen::ExecutionContext& ctx,
                     const miopen::conv::ProblemDescription& problem)
{
    return GetGemmSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<miopen::solver::ConvSolution>
FindAllDirectSolutions(const miopen::ExecutionContext& ctx,
                       const miopen::conv::ProblemDescription& problem,
                       const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetDirectSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<std::pair<std::string, size_t>>
AllDirectForwardBackwardDataWorkspaceSize(const miopen::ExecutionContext& ctx,
                                          const miopen::conv::ProblemDescription& problem)
{
    return GetDirectSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<std::pair<std::string, size_t>>
FindAllWinogradWorkspaceSizes(const miopen::ExecutionContext& ctx,
                              const miopen::conv::ProblemDescription& problem)
{
    return GetWindogradSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<std::pair<std::string, size_t>>
FindWinogradWrWWorkspaceSizes(const miopen::ExecutionContext& ctx,
                              const miopen::conv::ProblemDescription& problem)
{
    return GetWindogradWrWSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<std::pair<std::string, size_t>>
FindAllImplicitGemmWorkspaceSizes(const miopen::ExecutionContext& ctx,
                                  const miopen::conv::ProblemDescription& problem)
{
    return GetImplicitGemmSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<miopen::solver::ConvSolution>
FindAllImplicitGemmSolutions(const miopen::ExecutionContext& ctx,
                             const miopen::conv::ProblemDescription& problem,
                             const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetImplicitGemmSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<miopen::solver::ConvSolution>
FindAllWinogradSolutions(const miopen::ExecutionContext& ctx,
                         const miopen::conv::ProblemDescription& problem,
                         const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetWindogradSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<miopen::solver::ConvSolution>
FindWinogradWrWAllSolutions(const miopen::ExecutionContext& ctx,
                            const miopen::conv::ProblemDescription& problem,
                            const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetWindogradWrWSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<std::pair<std::string, size_t>>
AllDirectBwdWrW2DWorkspaceSize(const miopen::ExecutionContext& ctx,
                               const miopen::conv::ProblemDescription& problem)
{
    return GetBwdWrW2DSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<std::pair<std::string, size_t>>
FindImplicitGemmWrWWorkspaceSizes(const miopen::ExecutionContext& ctx,
                                  const miopen::conv::ProblemDescription& problem)
{
    return GetImplicitGemmWrWSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<miopen::solver::ConvSolution>
FindImplicitGemmWrWAllSolutions(const miopen::ExecutionContext& ctx,
                                const miopen::conv::ProblemDescription& problem,
                                const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetImplicitGemmWrWSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<miopen::solver::ConvSolution>
FindAllBwdWrW2DSolutions(const miopen::ExecutionContext& ctx,
                         const miopen::conv::ProblemDescription& problem,
                         const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetBwdWrW2DSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<miopen::solver::ConvSolution>
FindAllFFTSolutions(const miopen::ExecutionContext& ctx,
                    const miopen::conv::ProblemDescription& problem,
                    const miopen::AnyInvokeParams& invoke_ctx)
{
    return GetFFTSolvers().SearchForAllSolutions(
        ctx, problem, miopen::MakeConvDbGetter(ctx), invoke_ctx);
}

std::vector<std::pair<std::string, size_t>>
AllFFTForwardBackwardDataWorkspaceSize(const miopen::ExecutionContext& ctx,
                                       const miopen::conv::ProblemDescription& problem)
{
    return GetFFTSolvers().GetWorkspaceSizes(ctx, problem);
}

std::vector<std::string> GetAllFindSolverDbIds()
{
    std::vector<std::string> ids;
    const auto collect = [&](auto container) {
        container.Foreach([&](auto solver) { ids.push_back(solver.SolverDbId()); });
    };
    collect(GetGemmSolvers());
    collect(GetDirectSolvers());
    collect(GetImplicitGemmSolvers());
    collect(GetWindogradSolvers());
    collect(GetImplicitGemmWrWSolvers());
    collect(GetWindogradWrWSolvers());
    collect(GetBwdWrW2DSolvers());
    collect(GetFFTSolvers());
    return ids;
}

void mlo_construct_activ_lrn_pooling_common::setupFloats()
{
    if(_problem.GetInDataType() == miopenFloat && _problem.GetOutDataType() == miopenFloat)
    {
        _ctx.general_compile_options += " -DMIOPEN_USE_FP32=1 -DMIOPEN_USE_FP16=0";
    }
    else if(_problem.GetInDataType() == miopenHalf && _problem.GetOutDataType() == miopenHalf)
    {
        _ctx.general_compile_options += " -DMIOPEN_USE_FP32=0 -DMIOPEN_USE_FP16=1";
    }
    else
    {
        MIOPEN_LOG_W("Unsupported data types configuration: "
                     << miopen::GetDataTypeName(_problem.GetInDataType()) << "x"
                     << miopen::GetDataTypeName(_problem.GetOutDataType()));
    }
}
