"""Compute-bound prefill attention — 3 command-form variants (single-rank).

Companion to the memory-bound decode study
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp*``). Prefill processes a
block of T_q query positions at once, so the score / context GEMMs have a
large M = G·T_q and high arithmetic intensity (~M flops/byte) — the
workload is **compute-bound** (above the roofline ridge), unlike T_q=1
decode (M=8, memory-bound). This is the regime where the composite
command's value shows: it streams DMA↔compute per HW tile to keep the MAC
array fed, while the primitive kernel serializes load→dot and starves it.

Single-rank (C=P=1): no cross-device reduce — the focus is the per-PE
GEMM-issue mechanism. FlashAttention 2-D tiling (Q-block × S_kv-tile,
online softmax) bounds the TCM scratch.

Three forms, differing only in how each Q-block's local attention is
issued (placement/softmax identical):
  primitive           tl.dot per S_kv-tile + primitive online merge.
  composite           one coarse composite GEMM per Q-block over the whole
                      S_kv (K, V as HBM refs → PE_SCHEDULER tiles/streams).
  composite_extended  Q·Kᵀ composite + softmax_merge recipe composite.
"""
from __future__ import annotations

from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import _merge_running

M_BLOCK = 128       # query rows per FlashAttention Q-block (tile-filling).
TILE_S_KV = 256     # primitive per-tile S_kv width.
_SEED_S_KV = 64     # composite_extended recipe seed slice.


def _qblock_bounds(T_q: int, h_q: int, h_kv: int):
    G = h_q // h_kv
    M_total = G * T_q
    n_qblocks = (M_total + M_BLOCK - 1) // M_BLOCK
    return G, M_total, n_qblocks


# ── primitive: hand-tiled monolithic dots + online merge ─────────────


def gqa_prefill_primitive_kernel(
    q_ptr, k_ptr, v_ptr, o_ptr,
    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
) -> None:
    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
    ROW = d_head * 2  # f16
    n_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV

    for qb in range(n_qblocks):
        # Per-Q-block scratch_scope: this block's running (m,ℓ,O) plus its
        # tile-0 transients are freed before the next block, so scratch
        # stays O(one block) rather than accumulating across all blocks.
        with tl.scratch_scope():
            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
            q_off = qb * M_BLOCK * ROW
            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")

            tile_s0 = min(TILE_S_KV, S_kv)
            K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
            V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
            scores = tl.dot(Q, K_T)
            m_local = tl.max(scores, axis=-1)
            exp_s = tl.exp(scores - m_local)
            l_local = tl.sum(exp_s, axis=-1)
            O_local = tl.dot(exp_s, V)

            for ti in range(1, n_tiles):
                start = ti * TILE_S_KV
                tile_s = min(TILE_S_KV, S_kv - start)
                with tl.scratch_scope():
                    K_T_t = tl.load(k_ptr + start * ROW,
                                    shape=(d_head, tile_s), dtype="f16")
                    V_t = tl.load(v_ptr + start * ROW,
                                  shape=(tile_s, d_head), dtype="f16")
                    scores_t = tl.dot(Q, K_T_t)
                    m_tile = tl.max(scores_t, axis=-1)
                    exp_t = tl.exp(scores_t - m_tile)
                    l_tile = tl.sum(exp_t, axis=-1)
                    O_tile = tl.dot(exp_t, V_t)
                    m_new, l_new, O_new = _merge_running(
                        m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
                    )
                    tl.copy_to(m_local, m_new)
                    tl.copy_to(l_local, l_new)
                    tl.copy_to(O_local, O_new)

            O_final = O_local / l_local
            tl.store(o_ptr + q_off, O_final)


# ── composite: one coarse composite GEMM per Q-block ─────────────────


def gqa_prefill_composite_kernel(
    q_ptr, k_ptr, v_ptr, o_ptr,
    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
) -> None:
    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
    ROW = d_head * 2

    for qb in range(n_qblocks):
        with tl.scratch_scope():
            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
            q_off = qb * M_BLOCK * ROW
            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
            K_T = tl.ref(k_ptr, shape=(d_head, S_kv), dtype="f16")
            V = tl.ref(v_ptr, shape=(S_kv, d_head), dtype="f16")

            scores = tl.composite(op="gemm", a=Q, b=K_T)      # Q·Kᵀ, one command
            m_local = tl.max(scores, axis=-1)
            exp_s = tl.exp(scores - m_local)
            l_local = tl.sum(exp_s, axis=-1)
            O_local = tl.zeros((m_blk, d_head), dtype="f16")
            tl.composite(op="gemm", a=exp_s, b=V, out=O_local)  # P·V, one command

            O_final = O_local / l_local
            tl.store(o_ptr + q_off, O_final)


# ── composite_extended: Q·Kᵀ composite + softmax_merge recipe ────────


def gqa_prefill_composite_ext_kernel(
    q_ptr, k_ptr, v_ptr, o_ptr,
    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
) -> None:
    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
    ROW = d_head * 2

    for qb in range(n_qblocks):
        with tl.scratch_scope():
            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
            q_off = qb * M_BLOCK * ROW
            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")

            seed = min(_SEED_S_KV, S_kv)
            K_T0 = tl.load(k_ptr, shape=(d_head, seed), dtype="f16")
            V0 = tl.load(v_ptr, shape=(seed, d_head), dtype="f16")
            scores0 = tl.dot(Q, K_T0)
            m_local = tl.max(scores0, axis=-1)
            exp0 = tl.exp(scores0 - m_local)
            l_local = tl.sum(exp0, axis=-1)
            O_local = tl.dot(exp0, V0)

            rest = S_kv - seed
            if rest > 0:
                K_T1 = tl.ref(k_ptr + seed * ROW,
                              shape=(d_head, rest), dtype="f16")
                V1 = tl.ref(v_ptr + seed * ROW,
                            shape=(rest, d_head), dtype="f16")
                scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
                tl.composite(
                    prologue=[{"op": "softmax_merge", "s": scores1,
                               "m": m_local, "l": l_local, "O": O_local}],
                    op="gemm", b=V1, out=O_local,
                    epilogue=[{"op": "add", "other": O_local}],
                )

            O_final = O_local / l_local
            tl.store(o_ptr + q_off, O_final)
