"""GQA prefill kernel: composite + softmax_merge fused variant (3).

Identical mapping to ``_gqa_attention_prefill_short.py``. Difference vs
the GEMM-only composite: per-tile softmax is folded into
the P·V composite via the ``softmax_merge`` prologue recipe, eliminating
the GEMM/MATH engine bubble.

Multi-cube (A1/A2/A4): non-root PEs receive K/V via IPCQ ``tl.recv``.
Recv'd slots in on-chip memory are pinned (read in place as a composite
operand, not DMA-streamed from HBM).

Three-variant comparison:
  (1) without composite                    : ``_gqa_attention_prefill_short.py``
  (2) with composite (GEMM-only, no fuse)  : ``…_composite.py``
  (3) with composite + softmax_merge fuse  : this file

Shard addressing, layouts, and caller contract are identical to the
first-level kernel (ADR-0011 D-VA1).
"""
from __future__ import annotations


TILE_S_KV = 1024


def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
                     h_q: int, h_kv: int, S_kv: int) -> None:
    """Validate composite-prefill config — caller-side, sim-cost 0.

    Mirrors first-level prefill ``_validate_config``.
    """
    if kv_per_cube not in (1, 2, 4, 8):
        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
    if P != 8:
        raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
    if h_q % h_kv != 0:
        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
    if C != h_kv // kv_per_cube:
        raise ValueError(
            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
        )
    group_size = P // kv_per_cube
    if T_q < group_size:
        raise ValueError(
            f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
            f"has q-tile work for the broadcast chain"
        )
    if S_kv % TILE_S_KV != 0:
        raise ValueError(
            f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
        )


def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
    m_new = tl.maximum(m_local, m_other)
    scale_old = tl.exp(m_local - m_new)
    scale_new = tl.exp(m_other - m_new)
    l_new = l_local * scale_old + l_other * scale_new
    O_new = O_local * scale_old + O_other * scale_new
    return m_new, l_new, O_new


def gqa_attention_prefill_short_composite_fused_kernel(
    q_ptr: int,
    k_ptr: int,
    v_ptr: int,
    o_ptr: int,
    T_q: int,
    S_kv: int,
    h_q: int,
    h_kv: int,
    d_head: int,
    C: int,
    P: int,
    kv_per_cube: int,
    *,
    tl,
) -> None:
    """Composite-GEMM prefill — same mapping as first-level + tl.composite.

    Caller must invoke ``_validate_config(...)`` first.
    """
    group_size = P // kv_per_cube

    G = h_q // h_kv
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    pe_in_group = pe_id % group_size
    group_id_in_cube = pe_id // group_size

    q_start = (T_q * pe_in_group) // group_size
    q_end = (T_q * (pe_in_group + 1)) // group_size
    T_q_pe = q_end - q_start

    n_tiles = S_kv // TILE_S_KV

    Q_ROW_BYTES = G * d_head * 2
    KV_ROW_BYTES = d_head * 2
    K_TILE_BYTES = d_head * TILE_S_KV * 2
    K_HEAD_BYTES = n_tiles * K_TILE_BYTES
    V_HEAD_BYTES = S_kv * KV_ROW_BYTES

    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
    q_base = (q_ptr
              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
    k_head_shard_base = (k_ptr
                         + cube_id * kv_per_cube * K_HEAD_BYTES
                         + group_id_in_cube * K_HEAD_BYTES)
    v_head_shard_base = (v_ptr
                         + cube_id * kv_per_cube * V_HEAD_BYTES
                         + group_id_in_cube * V_HEAD_BYTES)

    Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")

    group_cols = min(4, group_size)
    group_rows = (group_size + group_cols - 1) // group_cols
    pe_col_in_group = pe_in_group % group_cols
    pe_row_in_group = pe_in_group // group_cols

    # ── Tile 0 — broadcast + persistent (m, ℓ, O) ──
    if pe_in_group == 0:
        K_T = tl.load(k_head_shard_base,
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.load(v_head_shard_base,
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if group_cols > 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)
        if group_rows > 1:
            tl.send(dir="intra_S", src=K_T)
            tl.send(dir="intra_S", src=V)
    elif pe_col_in_group == 0 and pe_row_in_group > 0:
        K_T = tl.recv(dir="intra_N",
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.recv(dir="intra_N",
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if group_cols > 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)
    else:
        K_T = tl.recv(dir="intra_W",
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.recv(dir="intra_W",
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if pe_col_in_group < group_cols - 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)

    # Tile 0 — primitives establish (m, ℓ, O); recipe fusion enters tile 1+.
    scores = tl.dot(Q, K_T)
    m_local = tl.max(scores, axis=-1)
    centered = scores - m_local
    exp_scores = tl.exp(centered)
    l_local = tl.sum(exp_scores, axis=-1)
    O_local = tl.dot(exp_scores, V)

    # ── Tiles 1..n_tiles-1 ──
    for tile_idx in range(1, n_tiles):
        with tl.scratch_scope():
            if pe_in_group == 0:
                K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if group_cols > 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
                if group_rows > 1:
                    tl.send(dir="intra_S", src=K_T_t)
                    tl.send(dir="intra_S", src=V_t)
            elif pe_col_in_group == 0 and pe_row_in_group > 0:
                K_T_t = tl.recv(dir="intra_N",
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.recv(dir="intra_N",
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if group_cols > 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
            else:
                K_T_t = tl.recv(dir="intra_W",
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.recv(dir="intra_W",
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if pe_col_in_group < group_cols - 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
            # Two-composite fusion: Q·Kᵀ composite → softmax_merge prologue
            # binds P (pinned primary-out) to the P·V composite, folding
            # the new tile's contribution into O_local.
            scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
            tl.composite(
                prologue=[{"op": "softmax_merge", "s": scores_t,
                           "m": m_local, "l": l_local, "O": O_local}],
                op="gemm", b=V_t, out=O_local,
                epilogue=[{"op": "add", "other": O_local}],
            )

    O_final = O_local / l_local
    o_base = (o_ptr
              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
    tl.store(o_base, O_final)
