// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.

#pragma once

#include "ck_tile/core.hpp"
#include "ck_tile/ops/common/tensor_layout.hpp"
#include "ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general_policy.hpp"

namespace ck_tile {

/*
This pipeline deal with a gemm(actually 2 gemm) with one very small(token), one very big(weight)
we need to design the pipeline such that all waves along gemm-N dim (gemm-m only 1 wave)

    <----- gemm-N ------>
    +----+----+----+----+
    | w0 | w1 | w2 | w3 | gemm-m
    +----+----+----+----+
*/
template <typename Problem_, typename Policy_ = FusedMoeGemmPipelineGeneralPolicy>
struct FusedMoeGemmPipeline_General
{
    using Problem = remove_cvref_t<Problem_>;
    using Policy  = remove_cvref_t<Policy_>;

    using BlockShape = typename Problem::BlockShape; // this is FusedMoeGemmShape

    using ADataType            = typename Problem::ADataType;
    using GDataType            = typename Problem::GDataType;
    using DDataType            = typename Problem::DDataType;
    using AccDataType          = typename Problem::AccDataType;
    using ODataType            = typename Problem::ODataType;
    using AScaleDataType       = typename Problem::AScaleDataType;
    using GScaleDataType       = typename Problem::GScaleDataType;
    using DScaleDataType       = typename Problem::DScaleDataType;
    using YSmoothScaleDataType = typename Problem::YSmoothScaleDataType;
    using TopkWeightDataType   = typename Problem::TopkWeightDataType;
    using IndexDataType        = typename Problem::IndexDataType;
    using YDataType            = typename Problem::YDataType;

    using Traits = typename Problem::Traits;

    static constexpr bool IsGateOnly          = Traits::IsGateOnly;
    static constexpr bool UseSmoothQuant      = Traits::UseSmoothQuant;
    static constexpr bool PadHiddenSize       = Traits::PadHiddenSize;
    static constexpr bool PadIntermediateSize = Traits::PadIntermediateSize;

    static constexpr index_t kAlignmentA = Policy::template GetAlignment_A<Problem>();
    static constexpr index_t kAlignmentG = Policy::template GetAlignment_G<Problem>();
    static constexpr index_t kAlignmentD = Policy::template GetAlignment_D<Problem>();
    static constexpr index_t kAlignmentO = Policy::template GetAlignment_O<Problem>();

    static constexpr index_t SLD_A = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::SLD_A);
    static constexpr index_t GLD_A = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::GLD_A);
    static constexpr index_t GLD_B = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::GLD_B);
    static constexpr index_t GST_O = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::GST_O);

    static constexpr index_t kBlockPerCu = []() {
        if constexpr(Problem::kBlockPerCu != -1)
            return Problem::kBlockPerCu;
        else
        {
            // minimize occupancy
            return 2;
        }
    }();

    static constexpr const char* name = "flatmm_gl";

    CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
    {
        // matrix a or tokens smem
        constexpr index_t smem_mat_a =
            BlockShape::Block_M0 * BlockShape::Block_K0 * sizeof(ADataType);
        // shuffle C matrix
        constexpr index_t smem_bridge =
            BlockShape::Block_M0 * BlockShape::Block_N0 * sizeof(YDataType);

        return max(smem_mat_a, smem_bridge);
        // return Policy::template GetSmemSize<Problem>();
    }

    // this is the thread-offset along row/col
    CK_TILE_HOST_DEVICE static auto GetACoord()
    {
        constexpr auto a_dist = Policy::template MakeGlobalTileDistribution_A<Problem>();
        const auto a_coord    = a_dist.calculate_index();
        return a_coord;
    }

    template <typename T>
    CK_TILE_HOST_DEVICE static void PrintMem(T& tensor)
    {
        constexpr auto spans = T::get_distributed_spans();
        int counter          = 0;
        sweep_tile_span(spans[number<0>{}], [&](auto idxn) {
            sweep_tile_span(spans[number<1>{}], [&](auto idxk) {
                constexpr auto i_j_idx = make_tuple(idxn, idxk);
                const auto tile_idx =
                    get_x_indices_from_distributed_indices(tensor.get_tile_distribution(), i_j_idx);
                if(threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0)
                {
                    const auto row = tile_idx.at(number<0>{});
                    const auto col = tile_idx.at(number<1>{});
                    printf("in G row is %d , col is %d, counter is %d, value is: %f"
                           " \n",
                           row,
                           col,
                           counter,
                           ck_tile::type_convert<float>(tensor(i_j_idx)));
                    counter = counter + 1;
                }
            });
        });
    }
    template <typename AWindow, typename GWindow, typename DWindow, typename OWindow>
    CK_TILE_DEVICE auto operator()(const AWindow& a_window_,
                                   const GWindow& g_window_,
                                   const DWindow& d_window_,
                                   OWindow& o_window_,
                                   TopkWeightDataType /*topk_weight*/,
                                   CK_TILE_LDS_ADDR void* smem,
                                   index_t hidden_size,
                                   index_t intermediate_size)
    {
        ignore = d_window_;
        ignore = hidden_size;
        ignore = intermediate_size;
        CK_TILE_LDS_ADDR ADataType* smem_0 = reinterpret_cast<CK_TILE_LDS_ADDR ADataType*>(smem);
        auto a_lds_view                    = make_tensor_view<address_space_enum::lds>(
            smem_0, Policy::template MakeLdsBlockDesc_A<Problem>());
        auto a_lds_win = make_tile_window(
            a_lds_view,
            make_tuple(number<BlockShape::Block_M0>{}, number<BlockShape::Block_K0>{}),
            {0, 0});

        auto a_global_to_dram_window = make_tile_window(
            a_window_.get_bottom_tensor_view(),
            make_tuple(number<BlockShape::Block_M0>{}, number<BlockShape::Block_K0>{}),
            a_window_.get_window_origin(),
            Policy::template MakeGlobalTileDistribution_A<Problem>());

        // load g to register
        auto g_global_to_dram_window = make_tile_window(
            g_window_.get_bottom_tensor_view(),
            make_tuple(number<BlockShape::Block_N0>{}, number<BlockShape::Block_K0>{}),
            g_window_.get_window_origin(),
            Policy::template MakeGlobalTileDistribution_G<Problem>());

        // gemm gate
        constexpr auto gemm_0   = Policy::template GetBlockGemm0<Problem>();
        using SaccBlockTileType = decltype(gemm_0.MakeCBlockTile());
        auto s_acc              = SaccBlockTileType{};

        // save tokens to lds
        auto a_dram_block = load_tile(a_global_to_dram_window);
        store_tile(a_lds_win, a_dram_block);
#if 0
        PrintMem(a_dram_block);
#endif

        auto g_dram_block = load_tile(g_global_to_dram_window);

#if 0
        PrintMem(g_dram_block);
#endif

        clear_tile(s_acc); // initialize C
        constexpr index_t kK0  = BlockShape::Block_K0;
        const index_t k0_loops = ck_tile::integer_divide_ceil(hidden_size, kK0);
        index_t iCounter0      = k0_loops - 1;
        while(iCounter0 > 0)
        {
            block_sync_lds();

            gemm_0(s_acc, a_lds_win, g_dram_block);

            block_sync_lds();
            move_tile_window(a_global_to_dram_window, {0, kK0});
            move_tile_window(g_global_to_dram_window, {0, kK0});

            a_dram_block = load_tile(a_global_to_dram_window);
            g_dram_block = load_tile(g_global_to_dram_window);

            store_tile(a_lds_win, a_dram_block);
            iCounter0--;
        }
        // tail
        {
            block_sync_lds();
            gemm_0(s_acc, a_lds_win, g_dram_block);
        }
        // relu
        const auto activation =  ck_tile::element_wise::Gelu{};
        tile_elementwise_inout(activation, s_acc, s_acc);
#if 0
        PrintMem(s_acc);
#endif
        // move sacc to LDS
        auto bridge_lds_view = make_tensor_view<address_space_enum::lds>(
            smem_0, Policy::template MakeBridgeLdsBlockDesc<Problem>());
        auto bridge_slds_win =
            make_tile_window(bridge_lds_view,
                             Policy::template MakeBridgeLdsBlockDesc<Problem>().get_lengths(),
                             {0, 0});
        // cast data to YDataType
        auto y_pre = cast_tile<YDataType>(s_acc);
#if 0
        PrintMem(y_pre);
#endif
        // save to lds
        store_tile(bridge_slds_win, y_pre);
        block_sync_lds();

        // gemm down
        constexpr auto gemm_1   = Policy::template GetBlockGemm1<Problem>();
        using OaccBlockTileType = decltype(gemm_1.MakeCBlockTile());
        auto o_acc              = OaccBlockTileType{};
        // y data
        auto bridge_llds_win =
            make_tile_window(bridge_lds_view,
                             Policy::template MakeBridgeLdsBlockDesc<Problem>().get_lengths(),
                             {0, 0},
                             Policy::template MakeYTileDistribution<Problem>());
        auto y = load_tile(bridge_llds_win);

        // d data
        auto d_global_to_dram_window = make_tile_window(
            d_window_.get_bottom_tensor_view(),
            make_tuple(number<BlockShape::Block_N0>{}, number<BlockShape::Block_K0>{}),
            d_window_.get_window_origin(),
            Policy::template MakeGlobalTileDistribution_D<Problem>());
        auto d = load_tile(d_global_to_dram_window);

        constexpr index_t kN1  = BlockShape::Block_N1;
        const index_t n1_loops = ck_tile::integer_divide_ceil(hidden_size, kN1);
        index_t iCounter1      = n1_loops - 1;
        while(iCounter1 > 0)
        {
            clear_tile(o_acc); 
            block_sync_lds();
            gemm_1(o_acc, y, d);
            block_sync_lds();
            move_tile_window(d_global_to_dram_window, {kN1, 0});
            d = load_tile(d_global_to_dram_window);
            // move out window and save data
            auto o = cast_tile<ODataType>(o_acc);
            store_tile(o_window_, o);
            move_tile_window(o_window_, {kN1, 0});

            iCounter1--;
        }
        // tail
        {
            clear_tile(o_acc); 
            block_sync_lds();
            gemm_1(o_acc, y, d);

            auto o = cast_tile<ODataType>(o_acc);
            store_tile(o_window_, o);
        }
#if 0
        PrintMem(o_acc);
#endif
        // store_tile(o_window_, a_dram_block);
    }
};

} // namespace ck_tile
