/**
 * UMMA QK GEMM Test — HD=64 (4 K-tiles, separate SMEM per K-tile)
 * 
 * Each K-tile gets its own (128, 16) SMEM region — no offset descriptors.
 * Source data stride handled correctly (SRC_HD=64, SMEM_HD=16).
 */

#include <cuda_runtime.h>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <cstring>

#include "dsv4/kernels/attention/fmha_common.cuh"
#include "dsv4/kernels/attention/fmha_umma_desc.cuh"

using namespace dsv4::kernels::attention;

static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }

constexpr int HD = 64;
constexpr int SK = 128;
constexpr int NKT = HD / MMA_K_BF16;  // 4
constexpr int BLOCK_MN = 128;
constexpr int TILE_SZ = BLOCK_MN * MMA_K_BF16;  // 128*16 = 2048 BF16 per K-tile

__global__ void __launch_bounds__(128)
test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
    float* __restrict__ s_out, float* __restrict__ s_scalar, float scale)
{
    const int tid = threadIdx.x;
    const int wid = tid / WARP, lane = tid % WARP;

    // SMEM: tmem_base(4) + pad(12) + Q tiles (4 × 2048 BF16) + K tiles (4 × 2048 BF16)
    extern __shared__ char sbuf[];
    uint32_t* sTmemBase = (uint32_t*)sbuf;
    bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
    bf16_t* sQ1 = sQ0 + TILE_SZ;
    bf16_t* sQ2 = sQ1 + TILE_SZ;
    bf16_t* sQ3 = sQ2 + TILE_SZ;
    bf16_t* sK0 = sQ3 + TILE_SZ;
    bf16_t* sK1 = sK0 + TILE_SZ;
    bf16_t* sK2 = sK1 + TILE_SZ;
    bf16_t* sK3 = sK2 + TILE_SZ;

    constexpr int CORES_MN = 16;  // 128/8

    // Load Q K-tiles: Q is (1, 64), each K-tile takes 16 dims
    // Zero all tiles
    for (int i = tid; i < NKT * TILE_SZ; i += 128) { sQ0[i] = 0; sK0[i] = 0; }
    __syncthreads();

    // Write Q row 0 to each K-tile's SMEM
    for (int kt = 0; kt < NKT; kt++) {
        bf16_t* sq = sQ0 + kt * TILE_SZ;
        for (int d = tid; d < MMA_K_BF16; d += 128) {
            int ck = d / 8, lc = d % 8;
            sq[ck * CORES_MN * 64 + lc] = q[kt * MMA_K_BF16 + d];
        }
        // Write K for this K-tile: K[r, 16*kt + d] for r=0..127, d=0..15
        bf16_t* sk = sK0 + kt * TILE_SZ;
        for (int r = 0; r < SK; r++) {
            for (int d = tid; d < MMA_K_BF16; d += 128) {
                int ck = d / 8, lc = d % 8;
                int tmn = r / 8, lr = r % 8;
                sk[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * MMA_K_BF16 + d];
            }
        }
    }
    __syncthreads();

    // TMEM alloc
    if (wid == 1) {
        tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128);
    }
    __syncthreads();
    uint32_t tb = *sTmemBase;

    // Multi-K-tile QK GEMM with separate SMEM per K-tile
    bf16_t* sQ_arr[NKT] = {sQ0, sQ1, sQ2, sQ3};
    bf16_t* sK_arr[NKT] = {sK0, sK1, sK2, sK3};
    uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN);

    for (int kt = 0; kt < NKT; kt++) {
        uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ_arr[kt]), BLOCK_MN);
        uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK_arr[kt]), BLOCK_MN);

        if (tid == 0) {
            umma_ss_f16(tb, dq, dk, idesc, kt > 0);
        }
        asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
        __syncthreads();
    }
    asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
    __syncthreads();

    // Read S from TMEM
    for (int n = 0; n < 128 / 8; n++) {
        const int row = wid * 32;
        const int col = n * 8;
        const int addr = tb + (row << 16) + col;
        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"(addr));
        asm volatile("tcgen05.wait::ld.sync.aligned;");

        int out_row = wid * 32 + lane;
        if (out_row < SK) {
            for (int c = 0; c < 8; c++) {
                int out_col = n * 8 + c;
                if (out_col < SK) {
                    s_out[out_row * SK + out_col] = tmp[c] * scale;
                }
            }
        }
    }
    __syncthreads();

    // Scalar reference
    if (tid == 0) {
        for (int j = 0; j < SK; j++) {
            float dot = 0.0f;
            for (int d = 0; d < HD; d++)
                dot += bf16_to_f32(q[d]) * bf16_to_f32(k[j * HD + d]);
            s_scalar[j] = dot * scale;
        }
    }

    if (wid == 0) tmem_dealloc(tb, 128);
}

int main() {
    printf("=== UMMA QK GEMM HD=64 (separate SMEM per K-tile) ===\n");
    const float SCALE = 1.0f / sqrtf((float)HD);

    bf16_t* h_q = (bf16_t*)malloc(HD * sizeof(bf16_t));
    bf16_t* h_k = (bf16_t*)malloc(SK * HD * sizeof(bf16_t));
    float* h_s_out = (float*)calloc(SK * SK, sizeof(float));
    float* h_s_scalar = (float*)calloc(SK, sizeof(float));

    srand(42);
    for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
    for (int i = 0; i < SK*HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);

    bf16_t *d_q, *d_k; float *d_s_out, *d_s_scalar;
    cudaMalloc(&d_q, HD*sizeof(bf16_t)); cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));
    cudaMalloc(&d_s_out, SK*SK*sizeof(float)); cudaMalloc(&d_s_scalar, SK*sizeof(float));
    cudaMemcpy(d_q, h_q, HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);

    // SMEM: 4 + 12(pad) + 8 * 2048*2
    int smem = (4 + 16 + NKT * 2 * TILE_SZ * sizeof(bf16_t) + 256 + 127) & ~127;
    printf("SMEM: %d bytes (%d KB)\n", smem, smem / 1024);

    test_umma_hd64<<<1, 128, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE);

    cudaError_t err = cudaDeviceSynchronize();
    if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }

    cudaMemcpy(h_s_out, d_s_out, SK*SK*sizeof(float), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost);

    printf("S[0,0..7] MMA:  "); for(int c=0;c<8;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n");
    printf("S[0,0..7] ref:  "); for(int c=0;c<8;c++) printf("%.6f ",h_s_scalar[c]); printf("\n");
    printf("S[0,64..71] MMA: "); for(int c=64;c<72;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n");
    printf("S[0,64..71] ref: "); for(int c=64;c<72;c++) printf("%.6f ",h_s_scalar[c]); printf("\n");

    float max_diff = 0.0f, max_val = 0.0f;
    for (int c = 0; c < SK; c++) {
        float diff = fabsf(h_s_out[0*SK+c] - h_s_scalar[c]);
        max_diff = fmaxf(max_diff, diff);
        max_val = fmaxf(max_val, fabsf(h_s_scalar[c]));
    }
    float rel_err = max_val > 0 ? max_diff / max_val : max_diff;
    printf("Row 0 rel err (128 cols): %.8f\n", rel_err);

    float max_nonzero = 0.0f;
    for (int r = 1; r < SK; r++)
        for (int c = 0; c < SK; c++)
            max_nonzero = fmaxf(max_nonzero, fabsf(h_s_out[r*SK+c]));
    printf("Rows 1-127 max abs: %.8f\n", max_nonzero);

    bool row0_ok = rel_err < 0.01f;
    bool rows_zero = max_nonzero < 1e-4f;
    printf("Row 0: %s | Rows 1-127 zero: %s\n",
        row0_ok ? "PASS" : "FAIL", rows_zero ? "PASS" : "FAIL");
    printf("Overall: %s\n", (row0_ok && rows_zero) ? "PASSED" : "FAILED");

    cudaFree(d_q); cudaFree(d_k); cudaFree(d_s_out); cudaFree(d_s_scalar);
    free(h_q); free(h_k); free(h_s_out); free(h_s_scalar);
    return (row0_ok && rows_zero) ? 0 : 1;
}
