#include "qk_gemm_utils_mls_ds.h"
#include "static_switch.h"


template<bool PREFETCH_K, int kHeadDim, int kHeadDimV, int computeHeadDim, int TailTile16, int kBlockM, int kBlockN, int kBlockK, int WARP_M, int WARP_N, int STAGES, typename Element, typename ElementAccum, bool Is_even_MN>
__forceinline__ __device__ void pv_gemm_prefetch_k_mls_ds(
        vec4_uint v_ptr,
        vec4_uint k_ptr,
        Element* v_lds,
        Element* k_lds,
        union_vec2_f16x2<Element> p_reg[(WARP_M / 32) * (kBlockK / 32)][4],
        vec4_Accum<ElementAccum> pv_reg[(kHeadDimV / kBlockN) * (WARP_M / 32) * (kBlockN / 32)][4],
        int warp_id,
        int seqlen_k_stride,
        int seqlen_v_stride,
        int max_seq_kv_offset=0) {

    constexpr int WARP_NUM = kBlockM * kBlockN / (WARP_M * WARP_N);
    constexpr int WARP_K   = 32;
    constexpr int READ_ONCE_COUNT = 32 * 32;
    constexpr int kHeadDimV_OPT = (kHeadDimV == 192) ? 128 : kHeadDimV;
    constexpr int V_LDS_LOAD_NUM  = (kHeadDimV_OPT * WARP_K) / READ_ONCE_COUNT;
    constexpr int V_LOAD_REQUESTS = V_LDS_LOAD_NUM / WARP_NUM;
    constexpr int ELEMENT_BYTES   = sizeof(Element);
    static_assert (kBlockK >= 32, "Error: pv gemm kBlockK must be equal or greater than 32");
    static_assert (kBlockM >= WARP_M, "Error: pv gemm kBlockM must be equal or greater than WARP_M");
    static_assert (kBlockN == WARP_N, "Error: pv gemm kBlockN must be equal to WARP_N");
    static_assert (WARP_K == 32 and "Error: To simplify, only WARP_K = 32 is supported!");
    static_assert (WARP_M == 32 and "Error: To simplify, only WARP_M = 32 is supported!");
    static_assert (WARP_N == 32 and "Error: To simplify, only WARP_N = 32 is supported!");

    // 计算 V lds 起始偏移量
    int v_lds_base = reinterpret_cast<size_t>(v_lds);

    // 准备 V 寄存器
    union_vec4_f16x2<Element> v_reg[STAGES * (32 * WARP_N) / (32 * 32) * 2];

    // MLS
    vec4_uint v_srsrc;
    v_srsrc[0] = v_ptr[0];
    v_srsrc[1] = v_ptr[1];
    v_srsrc[2] = seqlen_v_stride; // stride
    v_srsrc[3] = 0;

    int lds_stage_id = 1;
    for (int n_loop = 1; n_loop < (kBlockK / WARP_K); ++n_loop) {
        if constexpr (kHeadDimV == 128) {
            // MLS
            if constexpr (Is_even_MN) { // Is_even_MN 简单场景下 nm_filter 场景简化, 非 BlockM = 128 场景未必全支持
                *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (n_loop * WARP_K * seqlen_v_stride + warp_id * 32) * ELEMENT_BYTES);
                if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
            } else {
                int nm_filter_max = n_loop * WARP_K + 32 - max_seq_kv_offset;
                int real_mls_loop = nm_filter_max >= 32 ? 0: n_loop; // 如果全越界了, 则只访问 n_loop = 0 的那波数据
                int nm_filter     = inline_min_max<0, 32>(real_mls_loop * WARP_K + 32 - max_seq_kv_offset); // 重新计算 nm_filter
                v_srsrc[3]        = nm_filter << 8;
                *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (real_mls_loop * WARP_K * seqlen_v_stride + warp_id * 32) * ELEMENT_BYTES);
                if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
            }
            int lds_offset = (lds_stage_id * WARP_K * kHeadDimV + warp_id * 32 * 32) * ELEMENT_BYTES;
            flash::wait_all_warp_arrived();
            union union_vec4_uint v_rsrc_bits;
            v_rsrc_bits.v32 = v_srsrc;
            size_t lds_addr_warp = reinterpret_cast<size_t>(v_lds) + lds_offset;
            matrix_load_b16_lds_builtin<32, 32, 1, 0>(lds_addr_warp, v_rsrc_bits.i32, 0);
        } else if constexpr (kHeadDimV == 192) {
            int warp_id_m = warp_id % 2;    // w0 w2
            int warp_id_n = warp_id / 2;    // w1 w3
            int n_load = 1;   
            int n_loop_ = n_loop - 1; // Add this to support headdimv>128, we have to deal with previous n_loop
            *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (n_loop_ * WARP_K * seqlen_v_stride + warp_id_m * 16 * seqlen_v_stride + warp_id_n * 32 + n_load * WARP_NUM * 32) * ELEMENT_BYTES);
            if constexpr (true) {
                int nm_filter = inline_min_max<0, 16>(n_loop_ * WARP_K + warp_id_m * 16 + 16 - max_seq_kv_offset); // 重新计算 nm_filter
                v_srsrc[3]    = max_seq_kv_offset % kBlockK == 0 ? 0: nm_filter << 8;
                if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
            }
            int lds_offset = (lds_stage_id * WARP_K * kHeadDimV_OPT + warp_id * 32 * 16) * ELEMENT_BYTES;
            flash::wait_all_warp_arrived();
            union union_vec4_uint v_rsrc_bits;
            v_rsrc_bits.v32 = v_srsrc;
            size_t lds_addr_warp = reinterpret_cast<size_t>(v_lds) + lds_offset;
            matrix_load_b16_lds_builtin<32, 16, 1, 0>(lds_addr_warp, v_rsrc_bits.i32, 0);
        }
        // DS
        lds_stage_id ^= 1;
        int stage_id = 0;
        flash::wait_buffer_data_arrived<true>(V_LOAD_REQUESTS);

        int lds_load_offset = v_lds_base + (0/*k_loop*/ * 32 * 32 + lds_stage_id * WARP_K * kHeadDimV_OPT) * ELEMENT_BYTES;
        if constexpr (TailTile16 == 2) {
            DS_READ_MATRIX_32X32_B16_ALT2(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
        } else {
            DS_READ_MATRIX_32X32_B16(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
        }

        stage_id ^= 1;
        for (int k_loop = 1; k_loop < (computeHeadDim / kBlockN); ++k_loop) {

            // Wait for special headdim
            if constexpr (kHeadDimV == 192) {
                if ((k_loop & 3) == 0x0) {
                    flash::wait_buffer_data_arrived<true>(0);   
                }
            }

            int lds_load_offset = v_lds_base + (k_loop * 32 * 32 + lds_stage_id * WARP_K * kHeadDimV_OPT) * ELEMENT_BYTES;
            if constexpr (TailTile16 == 2) {
                DS_READ_MATRIX_32X32_B16_ALT2(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
            } else {
                DS_READ_MATRIX_32X32_B16(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
            }
            flash::wait_lds_data_arrived<false>(3);

            // MMAC
            flash::raise_priority();
            stage_id ^= 1;
            {
                constexpr int min_tile_k = 0;
                #pragma unroll
                for(int min_tile_n = 0; min_tile_n < 2; ++min_tile_n) {
                    #pragma unroll
                    for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                        int pv_tile_id = (STAGES == 2) ? k_loop - 1: k_loop;
                        int v_tile_id = stage_id * 2 + min_tile_k;
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                            p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                            v_reg[v_tile_id].f16x4[min_tile_n],
                            pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                    }
                }
            }
            flash::wait_lds_data_arrived<false>(2);
            {
                constexpr int min_tile_k = 1;
                #pragma unroll
                for(int min_tile_n = 0; min_tile_n < 2; ++min_tile_n) {
                    #pragma unroll
                    for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                        int pv_tile_id = (STAGES == 2) ? k_loop - 1: k_loop;
                        int v_tile_id = stage_id * 2 + min_tile_k;
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                            p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                            v_reg[v_tile_id].f16x4[min_tile_n],
                            pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                    }
                }
            }
            flash::lower_priority();

            // MLS for special headdimV
            if constexpr (kHeadDimV == 192) { // Once finish 128x32, we need issue next MLS instruction, and get now we need use data via wait
                if ((k_loop & 3) == 0x0) {
                    int n_load = k_loop / 4;
                    int n_loop_ = ((kHeadDimV / kBlockN) - k_loop) < 4 ? (n_load = 0, n_loop): n_loop_;   // if finish kHeadDimV*WarpK prefetch, we prefetch next n_loop data
                    if constexpr (Is_even_MN) {
                        *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (n_loop_ * WARP_K * seqlen_v_stride + warp_id * 32 + n_load * 32 * WARP_NUM) * ELEMENT_BYTES);
                        if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
                    } else {
                        int nm_filter_max = n_loop_ * WARP_K + 32 - max_seq_kv_offset;
                        int real_mls_loop = nm_filter_max >= 32 ? 0: n_loop_; // 如果全越界了, 则只访问 n_loop = 0 的那波数据
                        int nm_filter     = inline_min_max<0, 32>(real_mls_loop * WARP_K + 32 - max_seq_kv_offset); // 重新计算 nm_filter
                        v_srsrc[3]        = nm_filter << 8;
                        if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
                        *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (real_mls_loop * WARP_K * seqlen_v_stride + warp_id * 32 + n_load * 32 * WARP_NUM) * ELEMENT_BYTES);
                    }
                    int lds_offset = (lds_stage_id * WARP_K * kHeadDimV_OPT + warp_id * 32 * 32) * ELEMENT_BYTES;
                    flash::wait_all_warp_arrived();
                    union union_vec4_uint v_rsrc_bits;
                    v_rsrc_bits.v32 = v_srsrc;
                    size_t lds_addr_warp = reinterpret_cast<size_t>(v_lds) + lds_offset;
                    matrix_load_b16_lds_builtin<32, 32, 1, 0>(lds_addr_warp, v_rsrc_bits.i32, 0);
                }
            }
        }
        stage_id ^= 1;
        // Wait DS
        flash::wait_lds_data_arrived<false>(1);
        // last mmac
        flash::raise_priority();
        {
            constexpr int min_tile_k = 0;
            #pragma unroll
            for(int min_tile_n = 0; min_tile_n < TailTile16; ++min_tile_n) {
                #pragma unroll
                for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                    int pv_tile_id = (computeHeadDim / kBlockN) - 1;
                    int v_tile_id = stage_id * 2 + min_tile_k;
                    pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                        p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                        v_reg[v_tile_id].f16x4[min_tile_n],
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                }
            }
        }
        flash::wait_lds_data_arrived<false>(0);
        {
            constexpr int min_tile_k = 1;
            #pragma unroll
            for(int min_tile_n = 0; min_tile_n < TailTile16; ++min_tile_n) {
                #pragma unroll
                for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                    int pv_tile_id = (computeHeadDim / kBlockN) - 1;
                    int v_tile_id = stage_id * 2 + min_tile_k;
                    pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                        p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                        v_reg[v_tile_id].f16x4[min_tile_n],
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                }
            }
        }
        flash::lower_priority();
    }

    // 128 提前发读取 K 的指令
    if constexpr (kHeadDimV == 128) {
        if constexpr (PREFETCH_K) {
            prefetch_k_to_lds_mls_ds<kHeadDim, kBlockK, kBlockN, WARP_NUM, WARP_N, Element, Is_even_MN>(k_ptr, k_lds, warp_id, seqlen_k_stride, max_seq_kv_offset - kBlockK);
        }
    }

    {
        constexpr int n_loop = kBlockK / WARP_K;
        // MLS for special headdimV
        if constexpr (kHeadDimV == 192) {
            int warp_id_m = warp_id % 2;    // w0 w2
            int warp_id_n = warp_id / 2;    // w1 w3
            constexpr int n_loop_ = n_loop - 1;
            int n_load = 1;
            lds_stage_id ^= 1;
            
            if constexpr (Is_even_MN) {
                *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (n_loop_ * WARP_K * seqlen_v_stride + warp_id_m * 16 * seqlen_v_stride + warp_id_n * 32 + n_load * WARP_NUM * 32) * ELEMENT_BYTES);
                if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
            } else {
                int nm_filter_max = n_loop_ * WARP_K + warp_id_m * 16 + 16 - max_seq_kv_offset;
                int real_mls_loop = nm_filter_max >= 16 ? 0: n_loop_; // 如果全越界了, 则只访问 n_loop = 0 的那波数据
                int nm_filter     = inline_min_max<0, 16>(real_mls_loop * WARP_K + warp_id_m * 16 + 16 - max_seq_kv_offset); // 重新计算 nm_filter
                v_srsrc[3]        = nm_filter << 8;
                if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
                *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (real_mls_loop * WARP_K * seqlen_v_stride + warp_id_m * 16 * seqlen_v_stride + warp_id_n * 32 + n_load * WARP_NUM * 32) * ELEMENT_BYTES);
            }
            int lds_offset = (lds_stage_id * WARP_K * kHeadDimV_OPT + warp_id * 32 * 16) * ELEMENT_BYTES;
            flash::wait_all_warp_arrived();
            union union_vec4_uint v_rsrc_bits;
            v_rsrc_bits.v32 = v_srsrc;
            size_t lds_addr_warp = reinterpret_cast<size_t>(v_lds) + lds_offset;
            matrix_load_b16_lds_builtin<32, 16, 1, 0>(lds_addr_warp, v_rsrc_bits.i32, 0);
        }

        lds_stage_id ^= 1;
        int stage_id = 0;

        if constexpr (kHeadDimV == 128) {
            if constexpr (PREFETCH_K) {
                flash::wait_buffer_data_arrived<true>(V_LOAD_REQUESTS);
            } else {
                flash::wait_buffer_data_arrived<true>(0);
            }
        } else if constexpr (kHeadDimV == 192) {
                flash::wait_buffer_data_arrived<true>(V_LOAD_REQUESTS); // prefetchK置后
        }
        __builtin_amdgcn_sched_barrier(0);

        // DS
        int lds_load_offset = v_lds_base + (0/*k_loop*/ * 32 * 32 + lds_stage_id * WARP_K * kHeadDimV_OPT) * ELEMENT_BYTES;
        if constexpr (TailTile16 == 2) {
            DS_READ_MATRIX_32X32_B16_ALT2(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
        } else {
            DS_READ_MATRIX_32X32_B16(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
        }

        stage_id ^= 1;
        for (int k_loop = 1; k_loop < (computeHeadDim / kBlockN); ++k_loop) {

            // Wait for special headdim
            if constexpr (kHeadDimV == 192) {
                if ((k_loop & 3) == 0x0) {
                    flash::wait_buffer_data_arrived<true>(0);   
                }
            }

            // DS
            int lds_load_offset = v_lds_base + (k_loop * 32 * 32 + lds_stage_id * WARP_K * kHeadDimV_OPT) * ELEMENT_BYTES;
            if constexpr (TailTile16 == 2) {
                DS_READ_MATRIX_32X32_B16_ALT2(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
            } else {
                DS_READ_MATRIX_32X32_B16(lds_load_offset, v_reg[stage_id * 2 + 0].f16, v_reg[stage_id * 2 + 1].f16, false/*transpose*/);
            }
            flash::wait_lds_data_arrived<false>(3);

            // MMAC
            flash::raise_priority();
            stage_id ^= 1;
            {
                constexpr int min_tile_k = 0;
                #pragma unroll
                for(int min_tile_n = 0; min_tile_n < 2; ++min_tile_n) {
                    #pragma unroll
                    for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                        int pv_tile_id = (STAGES == 2) ? k_loop - 1: k_loop;
                        int v_tile_id = stage_id * 2 + min_tile_k;
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                            p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                            v_reg[v_tile_id].f16x4[min_tile_n],
                            pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                    }
                }
            }
            flash::wait_lds_data_arrived<false>(2);
            {
                constexpr int min_tile_k = 1;
                #pragma unroll
                for(int min_tile_n = 0; min_tile_n < 2; ++min_tile_n) {
                    #pragma unroll
                    for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                        int pv_tile_id = (STAGES == 2) ? k_loop - 1: k_loop;
                        int v_tile_id = stage_id * 2 + min_tile_k;
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                            p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                            v_reg[v_tile_id].f16x4[min_tile_n],
                            pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                    }
                }
            }
            flash::lower_priority();
        }

        stage_id ^= 1;
        flash::wait_lds_data_arrived<false>(1);
        // last mmac
        flash::raise_priority();
        {
            constexpr int min_tile_k = 0;
            #pragma unroll
            for(int min_tile_n = 0; min_tile_n < TailTile16; ++min_tile_n) {
                #pragma unroll
                for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                    int pv_tile_id = (computeHeadDim / kBlockN) - 1;
                    int v_tile_id = stage_id * 2 + min_tile_k;
                    pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                        p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                        v_reg[v_tile_id].f16x4[min_tile_n],
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                }
            }
        }
        flash::wait_lds_data_arrived<false>(0);
        {
            constexpr int min_tile_k = 1;
            #pragma unroll
            for(int min_tile_n = 0; min_tile_n < TailTile16; ++min_tile_n) {
                #pragma unroll
                for(int min_tile_m = 0; min_tile_m < 2; ++min_tile_m) {
                    int pv_tile_id = (computeHeadDim / kBlockN) - 1;
                    int v_tile_id = stage_id * 2 + min_tile_k;
                    pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32 = mmac_4interleave<Element, ElementAccum>(
                        p_reg[n_loop - 1][2 * min_tile_k + min_tile_m].f16x4,
                        v_reg[v_tile_id].f16x4[min_tile_n],
                        pv_reg[pv_tile_id][min_tile_n * 2 + min_tile_m].f32);
                }
            }
        }
        flash::lower_priority();
    }

    if constexpr (kHeadDimV == 192) {
        if constexpr (PREFETCH_K) {
            prefetch_k_to_lds_mls_ds<kHeadDim, kBlockK, kBlockN, WARP_NUM, WARP_N, Element, Is_even_MN>(k_ptr, k_lds, warp_id, seqlen_k_stride, max_seq_kv_offset - kBlockK);
        }
    }
}
