"""Mesh-native bidirectional AllReduce-mlo attention — decode (ADR-0059 Proposed).

Every rank holds the full Q (replicated, small at ``S_q=1``) and 1/n_ranks
of KV (sequence-sharded). Each rank computes its partial attention
against own KV in ONE shot, then runs a bidirectional fan-out of the
``(m, ℓ, o)`` triplet: the triplet originating at rank i reaches rank j at
step ``|i - j|``. Every rank merges every other rank's triplet exactly
once over ``n_ranks - 1`` steps, ending with the final answer replicated
on every rank.

Supersedes ADR-0056's closed-ring ``_attention_allreduce_mlo.py``. Both
modules stay on disk during the transition; this one runs on the
hardware's actual open-mesh wiring (no closed-ring SFR install required).

Imported by ``milestone_gqa_llama70b`` (after the bench's Phase 2 switches
its imports) and invoked through ``torch.launch(...)`` — not through
``dist.all_reduce(...)``. See ADR-0056 Context for why this kernel is not
backend-dispatched via ADR-0050's algorithm-module contract.
"""
from __future__ import annotations

from kernbench.common.pe_commands import TensorHandle


def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
    """Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
    return TensorHandle(
        id=handle.id,
        addr=handle.addr,
        shape=new_shape,
        dtype=handle.dtype,
        nbytes=handle.nbytes,
        data=handle.data,
        space=handle.space,
        pinned=handle.pinned,
    )


def attention_mesh_mlo_kernel(
    q_ptr: int,
    k_ptr: int,
    v_ptr: int,
    o_ptr: int,
    S_q: int,
    S_kv_per_rank: int,
    h_q: int,
    h_kv: int,
    d_head: int,
    n_ranks: int,
    rank_axis: int = 0,
    *,
    tl,
) -> None:
    """Mesh-native bidirectional AllReduce-mlo — see module docstring.

    ``rank_axis`` selects which program-id dimension carries the ring rank,
    matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction
    (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`):

      0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE
          id in cube). KV is split @ PEs **intra-cube**; ring runs over
          the 8 PEs of one cube (NOC-only). At Llama-70B headline scale
          this kernel launches once per cube; 64 such cubes run in
          parallel for one user (1 Q-head per cube × 8 cubes per KV-group
          × 8 KV-groups). The PE-level ring inside each cube is
          independent of the others.

      1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1)
          (cube id). KV is split @ cubes **inter-cube**; ring runs over
          the cubes of one KV-group. The kernel gates ``pe_id != 0`` to
          return early — a v1 simplification: at headline scale (B=8) the
          study's "Batch on batch" pattern would have all 8 PEs each handle
          one user's batch element instead of staying silent. Validation
          shipped with B=1 to focus on the inter-cube ring's correctness;
          adding the per-cube batch dimension is sub-cycle 4c headline work.
    """
    # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
    if rank_axis != 0 and tl.program_id(axis=0) != 0:
        return
    rank = tl.program_id(axis=rank_axis)
    has_E = rank < n_ranks - 1
    has_W = rank > 0

    # Q is replicated on every rank — loaded once.
    Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")

    # Local KV chunk. KV is sequence-sharded and stays put on this rank for
    # the entire fan-out — distinguishing decode from prefill (ADR-0059 D3)
    # where KV circulates.
    K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
    V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")

    # ── One-shot local partial attention ──────────────────────────
    K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
    V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
    scores = tl.dot(Q, K_2d_T)
    m = tl.max(scores, axis=-1)
    P = tl.softmax(scores, axis=-1)
    scores_centered = scores - m
    exp_scores = tl.exp(scores_centered)
    ell = tl.sum(exp_scores, axis=-1)
    o = tl.dot(P, V_2d)

    # Seed bidirectional waves with own triplet (step-1 send).
    to_send_east_m: TensorHandle | None = m
    to_send_east_ell: TensorHandle | None = ell
    to_send_east_o: TensorHandle | None = o
    to_send_west_m: TensorHandle | None = m
    to_send_west_ell: TensorHandle | None = ell
    to_send_west_o: TensorHandle | None = o

    # Bidirectional fan-out of (m, ℓ, o) triplets — n_ranks - 1 steps.
    for step in range(1, n_ranks):
        # Send eastbound triplet (own at step 1; forwarded at later steps).
        if has_E and to_send_east_m is not None:
            tl.send(dir="E", src=to_send_east_m)
            tl.send(dir="E", src=to_send_east_ell)
            tl.send(dir="E", src=to_send_east_o)
        # Send westbound triplet.
        if has_W and to_send_west_m is not None:
            tl.send(dir="W", src=to_send_west_m)
            tl.send(dir="W", src=to_send_west_ell)
            tl.send(dir="W", src=to_send_west_o)

        # Receive eastbound triplet from W (originated at rank - step).
        m_from_W: TensorHandle | None = None
        ell_from_W: TensorHandle | None = None
        o_from_W: TensorHandle | None = None
        if has_W and (rank - step) >= 0:
            m_from_W = tl.recv(dir="W", shape=m.shape, dtype="f16")
            ell_from_W = tl.recv(dir="W", shape=ell.shape, dtype="f16")
            o_from_W = tl.recv(dir="W", shape=o.shape, dtype="f16")
            m_combined = tl.maximum(m, m_from_W)
            scale_old = tl.exp(m - m_combined)
            scale_new = tl.exp(m_from_W - m_combined)
            ell = ell * scale_old + ell_from_W * scale_new
            o = o * scale_old + o_from_W * scale_new
            m = m_combined

        # Receive westbound triplet from E (originated at rank + step).
        m_from_E: TensorHandle | None = None
        ell_from_E: TensorHandle | None = None
        o_from_E: TensorHandle | None = None
        if has_E and (rank + step) < n_ranks:
            m_from_E = tl.recv(dir="E", shape=m.shape, dtype="f16")
            ell_from_E = tl.recv(dir="E", shape=ell.shape, dtype="f16")
            o_from_E = tl.recv(dir="E", shape=o.shape, dtype="f16")
            m_combined = tl.maximum(m, m_from_E)
            scale_old = tl.exp(m - m_combined)
            scale_new = tl.exp(m_from_E - m_combined)
            ell = ell * scale_old + ell_from_E * scale_new
            o = o * scale_old + o_from_E * scale_new
            m = m_combined

        # Forward the original received triplet (not the merged running state)
        # so neighbors get the original wave. ``None`` propagates if nothing
        # arrived this step.
        to_send_east_m = m_from_W
        to_send_east_ell = ell_from_W
        to_send_east_o = o_from_W
        to_send_west_m = m_from_E
        to_send_west_ell = ell_from_E
        to_send_west_o = o_from_E

    # Final normalize: O := o / ℓ.
    O_final = o / ell
    tl.store(o_ptr, O_final)
