#pragma once // prepare for prefetch V in qk gemm
#include "intrinsic_mls_ds.h"


template<int kHeadDim, int kBlockM, int kBlockN, int kBlockK, int WARP_M, int WARP_N, int TailTile16, typename Element, bool Is_even_MN>
__forceinline__ __device__ void prefetch_v_to_lds_mls_ds(
        vec4_uint v_ptr,
        Element* v_lds,
        int warp_id,
        int seqlen_v_stride,
        int max_seq_kv_offset=0) {

    constexpr int ELEMENT_BYTES = sizeof(Element);
    constexpr int WARP_NUM = kBlockM * kBlockN / (WARP_M * WARP_N);
    constexpr int WARP_K   = 32;
    constexpr int kHeadDim_OPT = (kHeadDim == 192) ? 128 : kHeadDim;

    // MLS
    // int n_loop = 0;
    int n_load = 0;
    vec4_uint v_srsrc;
    if constexpr (kHeadDim == 128) {
        *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (/*n_loop * WARP_K * seqlen_v_stride*/ + warp_id * 32) * ELEMENT_BYTES);
    } else if constexpr (kHeadDim == 192) {
        *(uint64_t*)&v_srsrc = VA_LIMIT_BITS(*(uint64_t*)&v_ptr + (n_load * WARP_NUM * 32 + warp_id * 32) * ELEMENT_BYTES);
    }
    v_srsrc[2] = seqlen_v_stride; // stride
    if constexpr (true) {
        int nm_filter = inline_min_max<0, 32>(0 * WARP_K + 32 - max_seq_kv_offset);
        v_srsrc[3] = max_seq_kv_offset % kBlockK == 0 ? 0: nm_filter << 8;
        if constexpr (TailTile16 == 2) { v_srsrc[3] += 0x20000; }
    }
    int lds_stage_id = 0;
    int lds_offset = (lds_stage_id * WARP_K * kHeadDim_OPT + warp_id * 32 * 32) * ELEMENT_BYTES;

    flash::wait_all_warp_arrived(); // 防止写 v lds 和读 k lds 冲突, qk 可能有的 warp 没结束
    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);
    __builtin_amdgcn_sched_barrier(0);
}
