/**
 * DSV4 FMHA — 6-warp specialized kernel, multi-row softmax, TMA async loads.
 *
 * ==================================================================
 * DESIGN
 * ==================================================================
 *
 * Same 6-warp design as fmha_6warp_multirow.cuh, but replaces scalar
 * GMEM reads in the load warp with TMA async bulk copies.
 *
 * 6-warp CTA: warps 0-3 = softmax, warp 4 = MMA, warp 5 = TMA load.
 * Grid: (1, n_h, batch) — each CTA processes one head of one batch item.
 *
 * TMA PIPELINE (single-stage, no overlap yet):
 *   For each K sub-tile (kt):
 *     1. TMA warp issues cp.async.bulk.tensor.2d for Q sub-tile and K sub-tile
 *     2. mbarrier wait for TMA completion
 *     3. Load warp transposes row-major SMEM → canonical K-major SMEM
 *     4. MMA warp runs tcgen05.mma as before
 *
 * KEY DECISION: Q is loaded per K-sub-tile, not once at the start.
 * This means Q gets TMA-loaded NKT_QK times (HD/16 times). The
 * alternative is to TMA-load Q once, but that requires the GMEM
 * allocation to be padded to (128, HD) even when T < 128.
 *
 * Per-sub-tile loading is cleaner because:
 * 1. TMA tiles are always (128, 16) BF16 = 4KB — same for Q, K, V
 * 2. No padding needed in GMEM
 * 3. Enables pipeline overlap (TMA of kt+1 overlaps MMA of kt)
 * 4. TMA bandwidth is 2.5 TB/s on B200 — re-reading Q 4-16x is cheap
 *
 * SMEM LAYOUT (simpler than the original design):
 *   sMbar, sTmemBase, sRowMax, sRowSum, sQ_tma, sK_tma, sQ, sK, sPk, sV_tma, sV
 *   All TMA staging buffers are (128, 16) row-major = 4KB each.
 * ==================================================================
 */

#pragma once

#include "fmha_common.cuh"
#include "fmha_umma_desc.cuh"
#include "fmha_tma.cuh"

namespace dsv4::kernels::attention {

struct FmhaMultiRowTmaParams {
    const bf16_t* __restrict__ q;
    const bf16_t* __restrict__ k;
    const bf16_t* __restrict__ v;
    bf16_t* __restrict__ o;
    float* __restrict__ lse;
    int s_k, T;
    float scale;
    int head_dim;
    int q_head_stride, q_batch_stride;
    int k_head_stride, k_batch_stride;
    int v_head_stride, v_batch_stride;
    int o_head_stride, o_batch_stride;
    int lse_head_stride, lse_batch_stride;
    // TMA descriptors (device pointers to CUtensorMap in GMEM)
    CUtensorMap* __restrict__ tma_q;  // Q: (T, HD)
    CUtensorMap* __restrict__ tma_k;  // K: (s_k, HD)
    CUtensorMap* __restrict__ tma_v;  // V: (HD, s_k)
};

template<int HD, int SK_TILE = 128>
__global__ void __launch_bounds__(192)
fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
    static constexpr int NKT_QK = HD / MMA_K_BF16;
    static constexpr int NKT_PV = SK_TILE / MMA_K_BF16;
    static constexpr int N_NSUB = HD / 16;
    static constexpr int TILE_SZ = 128 * MMA_K_BF16;  // 128*16 BF16
    static constexpr int V_SUB_SZ = 16 * MMA_K_BF16;
    static constexpr int TMEM_N = (HD <= 128) ? 128 : 256;
    static constexpr int MAX_ROWS = 128;
    static constexpr int CORES_MN = 128 / 8;
    static constexpr int NUM_READS = SK_TILE / 8;
    // TMA tile size in BF16: (128, 16) = 2048 BF16 = 4096 bytes
    static constexpr int TMA_TILE_BF16 = 128 * MMA_K_BF16;

    const int head_idx = blockIdx.y;
    const int batch_idx = blockIdx.z;
    const int tid = threadIdx.x;
    const int wid = tid / 32;
    const int lane = tid % 32;
    const bool is_softmax_warp = (wid < 4);
    const bool is_mma_warp = (wid == 4);
    const bool is_load_warp = (wid == 5);
    const int T = params.T;
    const int s_k = params.s_k;
    const float scale = params.scale;

    // Output pointers (still use GMEM strides)
    bf16_t* __restrict__ o_head = params.o + head_idx * params.o_head_stride + batch_idx * params.o_batch_stride;
    float* __restrict__ lse_head = params.lse ? params.lse + head_idx * params.lse_head_stride + batch_idx * params.lse_batch_stride : nullptr;

    // TMA descriptor pointers
    CUtensorMap* __restrict__ tma_q = params.tma_q;
    CUtensorMap* __restrict__ tma_k = params.tma_k;
    CUtensorMap* __restrict__ tma_v = params.tma_v;

    // ==================================================================
    // SMEM allocation
    // ==================================================================
    extern __shared__ char sbuf[];
    size_t off = 0;

    // sTmemBase (4B)
    uint32_t* sTmemBase = (uint32_t*)sbuf; off = 4;

    // sMbar (128B-aligned)
    off = (off + 127) & ~(size_t)127;
    uint64_t* sMbar = (uint64_t*)(sbuf + off); off += 8;

    // sRowMax, sRowSum
    float* sRowMax = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float);
    float* sRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float);

    // sQ_tma: row-major (128, 16) — TMA destination for Q sub-tile
    off = (off + 127) & ~(size_t)127;
    bf16_t* sQ_tma = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t);

    // sK_tma: row-major (128, 16) — TMA destination for K sub-tile
    off = (off + 127) & ~(size_t)127;
    bf16_t* sK_tma = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t);

    // sQ: canonical K-major (128, 16) for MMA
    off = (off + 127) & ~(size_t)127;
    bf16_t* sQ = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);

    // sK: canonical K-major (128, 16) for MMA
    off = (off + 127) & ~(size_t)127;
    bf16_t* sK = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);

    // sPk: canonical K-major (128, 16) for P staging in PV
    off = (off + 127) & ~(size_t)127;
    bf16_t* sPk = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);

    // sV_tma: row-major (16, 128) — TMA destination for V sub-tile
    off = (off + 127) & ~(size_t)127;
    bf16_t* sV_tma = (bf16_t*)(sbuf + off); off += 16 * 128 * sizeof(bf16_t);

    // sV: canonical K-major — PV B operand (128, 16) transposed from (16, 128)
    off = (off + 127) & ~(size_t)127;
    bf16_t* sV = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);

    // ==================================================================
    // Initialize mbarrier
    // ==================================================================
    if (tid == 0) {
        uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
        tma_mbarrier_init(mbar_addr, 128 * 16 * 2);
    }

    // TMEM alloc
    if (is_mma_warp) {
        uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
        tmem_alloc(smem_ptr, TMEM_N);
    }
    __syncthreads();
    uint32_t tb = *sTmemBase;

    // Row assignment
    const bool my_warp_active = (T <= 32) ? (wid == 0) : is_softmax_warp;
    const int my_row = my_warp_active ? (wid * 32 + lane) : 0;
    const bool my_row_active = my_warp_active && (my_row < T);

    // ==================================================================
    // QK GEMM → S in TMEM (loop over K sub-tiles)
    // ==================================================================
    for (int kt = 0; kt < NKT_QK; kt++) {
        // --- TMA load Q sub-tile ---
        // Q is (T, HD). Load columns [kt*16, kt*16+16).
        // TMA coord: (coord_x = kt*16, coord_y = 0) — column offset in innermost dim
        // The TMA tile is (16, 128) for the (cols, rows) innermost-first convention.
        // This loads Q[0..127, kt*16..kt*16+15] into sQ_tma.

        // Re-init mbarrier
        if (tid == 0) {
            uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
            tma_mbarrier_init(mbar_addr, 128 * 16 * 2);
        }
        __syncthreads();

        if (is_load_warp && lane == 0) {
            uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sQ_tma);
            uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
            tma_load_3d(smem_dst, (uint64_t)tma_q, mbar_addr, kt * MMA_K_BF16, 0);
        }

        // Wait for Q TMA completion
        if (is_load_warp && lane == 0) {
            uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
            tma_mbarrier_wait(mbar_addr);
        }
        __syncthreads();

        // Transpose sQ_tma (row-major 128×16) → sQ (canonical 128×16)
        if (is_load_warp) {
            write_smem_canonical<128, MMA_K_BF16>(sQ, sQ_tma);
        }
        __syncthreads();

        // --- TMA load K sub-tile ---
        // K is (s_k, HD). Load columns [kt*16, kt*16+16).
        // TMA coord: (coord_x = kt*16, coord_y = 0)

        // Re-init mbarrier
        if (tid == 0) {
            uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
            tma_mbarrier_init(mbar_addr, 128 * 16 * 2);
        }
        __syncthreads();

        if (is_load_warp && lane == 0) {
            uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sK_tma);
            uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
            tma_load_3d(smem_dst, (uint64_t)tma_k, mbar_addr, kt * MMA_K_BF16, 0);
        }

        if (is_load_warp && lane == 0) {
            uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
            tma_mbarrier_wait(mbar_addr);
        }
        __syncthreads();

        // Transpose sK_tma → sK (canonical)
        if (is_load_warp) {
            write_smem_canonical<128, MMA_K_BF16>(sK, sK_tma);
        }
        __syncthreads();

        // MMA: sQ × sK → TMEM
        if (is_mma_warp) {
            uint32_t idesc = make_idesc(128, 128);
            uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ), 128);
            uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK), 128);
            if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0);
            asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
        }
        __syncthreads();
    }

    // TMEM visibility fence
    asm volatile("fence.sc.gpu;" ::: "memory");
    __syncthreads();

    // ==================================================================
    // SOFTMAX — compute P in registers (TWO passes over TMEM)
    // (Identical to non-TMA multirow kernel)
    // ==================================================================
    float my_row_max = -INFINITY;
    if (my_warp_active) {
        for (int n = 0; n < NUM_READS; n++) {
            float tmp[8];
            asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
                : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
                  "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
                : "r"(tb + n * 8));
            asm volatile("tcgen05.wait::ld.sync.aligned;");
            if (my_row_active) {
                for (int c = 0; c < 8; c++) {
                    int col = n * 8 + c;
                    if (col < s_k) my_row_max = fmaxf(my_row_max, tmp[c] * scale);
                }
            }
        }
    }

    if (my_row_active) sRowMax[my_row] = my_row_max;
    __syncthreads();

    // Pass 2: compute P values
    float my_p_vals[SK_TILE];
    float my_row_sum = 0.0f;

    if (my_warp_active) {
        float rm = my_row_active ? sRowMax[my_row] : 0.0f;
        for (int n = 0; n < NUM_READS; n++) {
            float tmp[8];
            asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
                : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
                  "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
                : "r"(tb + n * 8));
            asm volatile("tcgen05.wait::ld.sync.aligned;");
            if (my_row_active) {
                for (int c = 0; c < 8; c++) {
                    int col = n * 8 + c;
                    if (col < s_k) {
                        float p = expf(tmp[c] * scale - rm);
                        my_p_vals[col] = p;
                        my_row_sum += p;
                    }
                }
            }
        }
    }

    if (my_row_active) sRowSum[my_row] = my_row_sum;
    __syncthreads();

    // ==================================================================
    // PV GEMM — write P to sPk per K-tile, TMA load V, accumulate O in TMEM
    // ==================================================================
    for (int n_sub = 0; n_sub < N_NSUB; n_sub++) {
        int d_base = n_sub * 16;
        for (int pv_kt = 0; pv_kt < NKT_PV; pv_kt++) {
            const int col_start = pv_kt * MMA_K_BF16;

            // Zero sPk
            if (is_load_warp) {
                for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0;
            }
            __syncthreads();

            // Softmax warps: write P to sPk (same as non-TMA kernel)
            if (my_row_active) {
                for (int c = 0; c < MMA_K_BF16; c++) {
                    int gc = col_start + c;
                    int ck = c/8, lc = c%8;
                    int core_mn = my_row/8, local_r = my_row%8;
                    sPk[ck*CORES_MN*64 + core_mn*64 + local_r*8 + lc] = f32_to_bf16(my_p_vals[gc]);
                }
            }
            __syncthreads();

            // --- TMA load V sub-tile ---
            // V is (HD, s_k). Load rows [d_base, d_base+16), cols [col_start, col_start+128).
            // TMA coord: (coord_x = col_start, coord_y = d_base)
            // The TMA tile is (128, 16) — (cols, rows) innermost-first.
            // This loads V[d_base..d_base+15, col_start..col_start+127] into sV_tma.

            // Re-init mbarrier
            if (tid == 0) {
                uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
                tma_mbarrier_init(mbar_addr, 128 * 16 * 2);
            }
            __syncthreads();

            if (is_load_warp && lane == 0) {
                uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sV_tma);
                uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
                tma_load_3d(smem_dst, (uint64_t)tma_v, mbar_addr, col_start, d_base);
            }

            if (is_load_warp && lane == 0) {
                uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
                tma_mbarrier_wait(mbar_addr);
            }
            __syncthreads();

            // Transpose sV_tma (16, 128) row-major → sV (128, 16) canonical
            // B[r, d] where r = seq position, d = head dim offset
            // B[r, d] = V[d, r] = sV_tma[d * 128 + r]
            if (is_load_warp) {
                constexpr int SV_CORES_MN = 128 / 8;  // 16
                for (int i = lane; i < TILE_SZ; i += 32) sV[i] = 0;
                for (int i = lane; i < 16 * 128; i += 32) {
                    int d = i / 128;  // row in V_tma = head dim offset
                    int r = i % 128;  // col in V_tma = sequence position
                    int core_mn = r / 8;
                    int local_r = r % 8;
                    int core_k = d / 8;
                    int local_c = d % 8;
                    int dst_idx = core_k * SV_CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c;
                    sV[dst_idx] = sV_tma[i];
                }
            }
            __syncthreads();

            // MMA: sPk × sV → TMEM
            if (is_mma_warp) {
                uint32_t idesc_pv = make_idesc(128, 16);
                uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), 128);
                uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
                if (tid == 128) umma_ss_f16(tb + n_sub*16, dp, dv, idesc_pv, pv_kt > 0);
                asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
            }
            __syncthreads();
        }
    }

    // Ensure PV output is visible
    asm volatile("fence.sc.gpu;" ::: "memory");
    __syncthreads();

    // ==================================================================
    // EPILOGUE: TMEM → regs → normalize → BF16 → GMEM + LSE output
    // ==================================================================
    if (my_warp_active) {
        float rm = my_row_active ? sRowMax[my_row] : 0.0f;
        float rs = my_row_active ? sRowSum[my_row] : 0.0f;
        float inv_rs = my_row_active ? (1.0f / rs) : 0.0f;

        for (int n = 0; n < N_NSUB * 2; n++) {
            float tmp[8];
            asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
                : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
                  "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
                : "r"(tb + n * 8));
            asm volatile("tcgen05.wait::ld.sync.aligned;");

            if (my_row_active) {
                for (int c = 0; c < 8; c++) {
                    int d = n * 8 + c;
                    if (d < HD) o_head[my_row * HD + d] = f32_to_bf16(tmp[c] * inv_rs);
                }
            }
        }

        if (my_row_active && lse_head) lse_head[my_row] = logf(rs) + rm;
    }
    __syncthreads();

    if (is_mma_warp) tmem_dealloc(tb, TMEM_N);
}

} // namespace dsv4::kernels::attention
