"""GQA decode kernel — Case 5 (Cube-TP on d_head × PE-SP on S_kv).

Per the per-PE memory/comm summary chart  (paper_plot_gqa_4cases_summary.py):
this is the second d_head-TP variant — KV is sharded 64-way (cube=
column_wise on d_head, pe=row_wise on S_kv). Per-PE KV memory matches
Cases 4 and 6 ★ (640 MB / PE at 1M context across 80 layers), but the
attention algorithm pays the partial-score AR over INTER-cube fabric
(UCIe, slower) instead of Case 4's intra-cube fabric (NoC).

Symmetric mirror of Case 4 — intra and inter-cube AR roles SWAPPED:

  Step             Case 4 (Cube-SP × PE-TP/d_head)   Case 5 (THIS — Cube-TP/d_head × PE-SP)
  ───────────────  ────────────────────────────────  ──────────────────────────────────────
  partial Q·K^T    AR-SUM intra-cube (NoC)           AR-SUM inter-cube (UCIe)
  online-softmax   inter-cube lrab on (m, ℓ, O)      intra-cube merge on (m, ℓ, O)
                   (every PE participates; each       (online-softmax reduce-to-PE0)
                    PE owns its d_head/P slice)
  writer           root cube's P PEs (each writes    each cube's PE 0 writes its
                   its d_head/P slice)                d_head/C slice (C writers total)

Algorithm:
  1. Each PE loads its (d_head/C × S_kv/P) slice of K and V; full-d_head
     Q from HBM but only its d_head/C slice via column_wise DPPolicy.
  2. Per S_kv tile (per PE — already a strict subset of the cube's
     S_kv), compute partial_score = Q_shard · K_shard^T.  Shape:
     (G·T_q, S_kv/P).  AR-SUM these PARTIAL scores INTER-CUBE across
     the C cubes (same PE rank, different d_head slices) so every PE
     in cube c, rank p has the full Q · K^T_p scores for THIS PE's
     S_kv/P token range.
  3. Local softmax + P · V_shard → partial output (G·T_q, d_head/C)
     for this PE's S_kv/P slice and this cube's d_head/C slice.
  4. INTRA-CUBE online-softmax merge of (m, ℓ, O) reduce-to-PE0
     across the cube's P PEs (different S_kv slices, same d_head slice).
  5. PE 0 of each cube normalises and writes the cube's d_head/C slice.

Tensor layout:
  Q : (T_q, h_q · d_head)  cube=column_wise on d_head, pe=replicate.
      Per PE shape (T_q, h_q · d_head/C); reshape to (G·T_q, d_head/C).
  K : (S_kv, h_kv · d_head) cube=column_wise on d_head, pe=row_wise on S_kv.
      Per PE shape (S_kv/P, d_head/C); reshape to (d_head/C, S_kv/P) for K^T.
  V : same layout as K.
  O : (T_q, h_q · d_head) cube=column_wise on d_head, pe=replicate.
      Per-cube writer (PE 0) owns its d_head/C slice.

Topology / SFR:
  - Requires ``configure_sfr_intercube_multisip`` (disjoint intra_*
    and E/W/N/S namespaces).
  - INTRA-CUBE PEs arranged 2×4 (used for online-softmax merge).
  - INTER-CUBE 4×2 cube sub-mesh (used for partial-score AR sum).

ADR-0065 P5 scope note: op_log / latency measurement only; numeric
data-mode parity is a separate follow-up (matches the milestone-1h-gqa
sweep mode).
"""
from __future__ import annotations


TILE_S_KV = 1024  # per-tile S_kv width within a PE's own S_kv/P shard.

# Cube sub-mesh geometry (4×2). For Case 5 we pick cube 0 as the AR-sum
# root (simpler than lrab; partial-score AR doesn't need geometric centering
# because the reduction is element-wise sum, not online-softmax merge).
_SUB_W = 4
_SUB_H = 2
_AR_ROOT_CUBE = 0


def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
    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 _inter_cube_allreduce_sum(tensor, *, tl, cube_id: int,
                              sub_w: int = _SUB_W, sub_h: int = _SUB_H):
    """Sum-AllReduce a tensor INTER-CUBE across the 4×2 cube sub-mesh.

    Reduce-to-cube-0 (row chain west + col bridge north) followed by
    broadcast-from-cube-0 (col bridge south + row chain east). The
    payload is the partial-score tile — each cube has a different
    d_head slice's contribution; element-wise sum recovers the full
    Q·K^T scores for this PE rank's S_kv/P tokens.

    All P PEs of every cube participate independently (per-PE-rank
    inter-cube fabric lanes).
    """
    row = cube_id // sub_w
    col = cube_id % sub_w

    # ── Phase 1a: row reduce (west) — accumulate at col=0 ────────────
    if sub_w > 1:
        if col == sub_w - 1:
            tl.send(dir="W", src=tensor)
        elif 0 < col < sub_w - 1:
            with tl.scratch_scope():
                other = tl.recv(dir="E", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, tensor + other)
            tl.send(dir="W", src=tensor)
        elif col == 0:
            with tl.scratch_scope():
                other = tl.recv(dir="E", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, tensor + other)

    # ── Phase 1b: col reduce on col=0 (north) — accumulate at row=0 ──
    if col == 0 and sub_h > 1:
        if row == sub_h - 1:
            tl.send(dir="N", src=tensor)
        elif row == 0:
            with tl.scratch_scope():
                other = tl.recv(dir="S", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, tensor + other)

    # ── Phase 2a: col broadcast (south) — row=0 → row=1 on col=0 ─────
    if col == 0 and sub_h > 1:
        if row == 0:
            tl.send(dir="S", src=tensor)
        elif row > 0:
            with tl.scratch_scope():
                received = tl.recv(dir="N", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, received)

    # ── Phase 2b: row broadcast (east) — col=0 → all on each row ─────
    if sub_w > 1:
        if col == 0:
            tl.send(dir="E", src=tensor)
        elif 0 < col < sub_w - 1:
            with tl.scratch_scope():
                received = tl.recv(dir="W", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, received)
            tl.send(dir="E", src=tensor)
        elif col == sub_w - 1:
            with tl.scratch_scope():
                received = tl.recv(dir="W", shape=tensor.shape, dtype="f16")
                tl.copy_to(tensor, received)


def gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_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,
    *,
    tl,
) -> None:
    """Case-5 decode: Cube-TP-on-d_head × PE-SP-on-S_kv."""
    G = h_q // h_kv
    if d_head % C != 0:
        raise ValueError(
            f"Case 5 requires d_head={d_head} divisible by C={C}"
        )
    if S_kv % P != 0:
        raise ValueError(
            f"Case 5 requires S_kv={S_kv} divisible by P={P}"
        )

    d_head_local = d_head // C       # 128 / 8 = 16 in the headline setup
    S_local_pe = S_kv // P           # each PE owns S_kv/P tokens
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)

    PE_GRID_COLS = 4
    pe_col = pe_id % PE_GRID_COLS
    pe_row = pe_id // PE_GRID_COLS
    pe_cols_used = min(PE_GRID_COLS, P)
    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS

    KV_ROW_BYTES = d_head_local * 2  # bytes per K/V row at this PE's d_head slice

    # ── Load this PE's d_head slice of Q ────────────────────────────
    # Q in HBM is (T_q, h_q · d_head); column_wise cube shard gives each
    # cube (T_q, h_q · d_head/C); pe=replicate within cube → every PE
    # in the cube has the same data. Reshape to (G·T_q, d_head/C).
    Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16")

    # ── Local attention (tile sweep on this PE's S_kv/P shard) ──────
    n_tiles = (S_local_pe + TILE_S_KV - 1) // TILE_S_KV

    # Tile 0: bootstrap (m, ℓ, O_shard). The partial Q·K^T is across
    # d_head — must AR inter-cube before softmax.
    tile_s0 = min(TILE_S_KV, S_local_pe)
    K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16")
    V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16")
    partial_scores0 = tl.dot(Q, K_T)
    _inter_cube_allreduce_sum(partial_scores0, tl=tl, cube_id=cube_id)
    m_local = tl.max(partial_scores0, axis=-1)
    exp_scores = tl.exp(partial_scores0 - m_local)
    l_local = tl.sum(exp_scores, axis=-1)
    # Partial output (G·T_q, d_head/C) — sharded on d_head only; PE
    # still holds the full d_head/C slice for its S_kv/P tokens.
    O_local = tl.dot(exp_scores, V)

    for tile_idx in range(1, n_tiles):
        tile_start = tile_idx * TILE_S_KV
        tile_s = min(TILE_S_KV, S_local_pe - tile_start)
        with tl.scratch_scope():
            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
                            shape=(d_head_local, tile_s), dtype="f16")
            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
                          shape=(tile_s, d_head_local), dtype="f16")
            partial_scores_t = tl.dot(Q, K_T_t)
            _inter_cube_allreduce_sum(partial_scores_t, tl=tl,
                                      cube_id=cube_id)
            m_tile = tl.max(partial_scores_t, axis=-1)
            exp_scores_t = tl.exp(partial_scores_t - m_tile)
            l_tile = tl.sum(exp_scores_t, axis=-1)
            O_tile = tl.dot(exp_scores_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)

    # ── INTRA-CUBE online-softmax reduce-to-PE0 (different S_kv slices) ─
    # Each cube ends with (m, ℓ, O) for its d_head/C slice at PE 0.
    if pe_cols_used > 1:
        if pe_col < pe_cols_used - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_col > 0:
            tl.send(dir="intra_W", src=m_local)
            tl.send(dir="intra_W", src=l_local)
            tl.send(dir="intra_W", src=O_local)

    if pe_col == 0 and pe_rows_used > 1:
        if pe_row < pe_rows_used - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_row > 0:
            tl.send(dir="intra_N", src=m_local)
            tl.send(dir="intra_N", src=l_local)
            tl.send(dir="intra_N", src=O_local)

    # ── Final normalise + store (one writer per cube: PE 0; each
    #     cube owns a disjoint d_head/C slice of the output) ─────────
    if pe_id == 0:
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
