# SPDX-License-Identifier: Apache-2.0
"""C20: attention on the p150 through ``ttnn.transformer.scaled_dot_product_attention`` (PLAN.md 1.2 C20, probe P7).

SDPA has three silent traps (common/PROBES.md P7; ``probes/g4-attention-gridsample-msda.md``):

- ``is_causal`` defaults to True: with Sq == Sk it silently applies a causal mask, with Sq != Sk it is rejected;
- the default ``scale`` is ``1/sqrt(padded head_dim)``: after a 16 -> 32 head padding it is wrong (PCC 0.979);
- the default program config (32/32 chunks) is 2.5-7x slower everywhere and inaccurate over >= 32k keys.

:func:`sdpa` is the only way the ports call the op: ``is_causal=False`` always, ``scale`` is a required keyword,
and a ``program_config`` is always passed (chunk sizes from :func:`chunk_config`, the P7 table). Without a mask,
pass the logical, unaligned Sq / Sk: the op pads to its chunks and masks its own padding element by element. A user
mask is an additive bf16 TILE tensor **in DRAM**, ``[1|B, 1|H, Sq, Sk]``; ``-inf`` is safe, also when a whole first
key chunk is masked (P7).

**A user mask needs a tile-aligned Sk** (found by the Diffusion Planner port, 2026-10-07): with a mask, the reader
fills only the key *tiles* past ``ceil(Sk / 32)`` with ``-inf`` and reads the last partial tile from the mask
itself, whose tile padding is 0 after a host tilize (``reader_interleaved.cpp``, "Mask read"). The up to 31 padded
keys then join the softmax with score ``q . k_pad`` and value ``v_pad``: rel-L2 0.18 instead of 0.025 at 321 x 321
with 89 valid keys, 0.06 at P7's masked 564 x 564. :func:`sdpa` therefore refuses a mask with ``Sk % 32 != 0``:
give K / V a tile-aligned key count (:func:`aligned_keys`; the extra keys are masked with ``-inf`` in the mask,
``key_bias(padded_validity, sq)``), which costs nothing because the tiles are padded anyway.

Head layout. The op wants ``[B, H, S, D]`` bf16 TILE interleaved tensors with ``D`` a multiple of 32 and no padding on
B / H / D. :func:`split_heads` / :func:`split_qkv` / :func:`split_q_kv` turn projection outputs ``[B, 1, S, H*D]``
into heads (``ttnn.experimental.nlp_create_qkv_heads``, one program, logical S kept), and ``sdpa(...,
concat_heads=True)`` (the op's ``output_concat_heads``) or :func:`merge_heads` go back to ``[B, 1, S, H*D]``.
A head dim below 32 (16 in TransFusion / BEVFusion / PTv3) is zero-padded *in the weights*
(:func:`pad_head_columns` on the Q / K / V projections, :func:`pad_head_rows` on the output projection): exact
(padded output columns are 0, P7) and free at run time; pass ``scale=1/sqrt(logical head dim)``.

Masks. :func:`key_bias` builds ``[1, 1, Sq, Sk]`` (0 for valid keys, ``-inf`` otherwise) on the host; inside a
trace :func:`expand_key_bias` repeats a persistent ``[1, 1, 1, Sk]`` bias row (:func:`key_bias_row`) to
``[1, 1, Sq, Sk]`` in DRAM (one program, the upload shrinks from Sq * Sk to Sk values).

Fallback. :func:`attention_matmul` is plain ``softmax(scale Q K^T + bias) V`` with matmuls (for dtypes or layouts the
op rejects, e.g. fp32 attention); it needs a tile-aligned Sk (pad the keys and mask them with ``-inf``).
:func:`attention_reference` is the float64 numpy oracle of both.

Importing this module never imports ttnn (ttaw rule); ttnn is imported inside the functions that need it.
"""
from __future__ import annotations

import math
from typing import Any, Dict, Optional, Sequence, Tuple

import numpy as np

__all__ = [
    "aligned_keys",
    "HEAD_DIM_MULTIPLE",
    "PROGRAM_CONFIG_MIN_SK",
    "CHUNK_TABLE",
    "padded_head_dim",
    "chunk_config",
    "sdpa_program_config",
    "sdpa_compute_config",
    "sdpa",
    "split_heads",
    "split_qkv",
    "split_q_kv",
    "merge_heads",
    "key_bias",
    "key_bias_row",
    "expand_key_bias",
    "pad_head_dim",
    "pad_head_columns",
    "pad_head_rows",
    "attention_reference",
    "attention_matmul",
]

HEAD_DIM_MULTIPLE = 32  # SDPA rejects a padded head dim ("Padding is not supported on the head_dim dimension")
TILE = 32
# P7: from about 1,000 keys on, the default 32/32 chunks are 2.5-7x slower; this wrapper passes a config always.
PROGRAM_CONFIG_MIN_SK = 1000

# Exact (Sq, Sk) shapes measured on the p150b (ETH, 12x10): (q_chunk, k_chunk) and the evidence. Shapes not listed
# use the Sk-range rules of :func:`chunk_config`. Ports add their measured shapes here (with the log).
CHUNK_TABLE: Dict[Tuple[int, int], Tuple[int, int, str]] = {
    (500, 500): (128, 128, "P7 tf_self 0.029 ms (default 0.047)"),
    (512, 512): (512, 512, "P7 pt_patch_s0/s4 0.172-0.222 ms (q128 k128 0.413-0.571)"),
    (564, 564): (64, 64, "P7 dp_fusion_masked 0.058 ms (default 0.094, q128 k128 0.158)"),
    (321, 564): (64, 64, "P7 dp_cross 0.029 ms (default 0.030)"),
    (576, 576): (64, 64, "DP fusion, masked: 0.061-0.092 ms (q32 k64 0.117, q32 k32 0.129, q128 k128 0.275)"),
    (352, 352): (64, 64, "DP DiT self, masked: 0.063 ms (q64 k128 0.059, q32 k32 0.080, q128 k128 0.071)"),
    (352, 564): (64, 64, "DP DiT cross: 0.033-0.043 ms (q32 k32 0.035, q128 k128 0.052, q64 k128 0.077)"),
    (900, 1668): (64, 128, "P7 sp_self 0.052 ms (default 0.139)"),
    (900, 6000): (64, 512, "P7 sp_cross 0.106 ms (default 0.469)"),
}


def padded_head_dim(head_dim: int, multiple: int = HEAD_DIM_MULTIPLE) -> int:
    """The head dim SDPA sees: ``head_dim`` rounded up to a multiple of 32."""
    if head_dim <= 0:
        raise ValueError(f"head_dim must be positive, got {head_dim}")
    return -(-int(head_dim) // multiple) * multiple


def aligned_keys(n: int, multiple: int = TILE) -> int:
    """The tile-aligned key count a masked attention over ``n`` real keys runs on (the extra keys are masked)."""
    if n <= 0:
        raise ValueError(f"key count must be positive, got {n}")
    return -(-int(n) // multiple) * multiple


def chunk_config(sq: int, sk: int) -> Tuple[int, int]:
    """``(q_chunk, k_chunk)`` for logical ``Sq`` x ``Sk`` (multiples of 32, as the op requires).

    Exact shapes come from :data:`CHUNK_TABLE`. Otherwise, by key length (P7: chunk sizes are what matters for
    long keys; q64 keeps ``B * H * Sq / 64`` work units near the 120 cores at Sq ~ 500-900):
    Sk >= 16,384 -> q64 k1024 (P7 tf / bf / pt cross 0.45-0.88 ms; q64 k512 + ``fp32_acc`` halves the max error at
    ~1.4x the time); Sk >= 4,096 -> q64 k512; Sk >= 1,000 -> q64 k128; shorter keys -> q64 k64 (32 when the
    sequence fits one tile).
    """
    sq, sk = int(sq), int(sk)
    if sq <= 0 or sk <= 0:
        raise ValueError(f"sequence lengths must be positive, got Sq={sq}, Sk={sk}")
    if (sq, sk) in CHUNK_TABLE:
        q, k, _ = CHUNK_TABLE[(sq, sk)]
        return q, k
    q = 32 if sq <= TILE else 64
    if sk >= 16384:
        return q, 1024
    if sk >= 4096:
        return q, 512
    if sk >= PROGRAM_CONFIG_MIN_SK:
        return q, 128
    return q, (32 if sk <= TILE else 64)


def sdpa_program_config(device: Any, sq: int, sk: int, *, q_chunk: Optional[int] = None, k_chunk: Optional[int] = None,
                   exp_approx_mode: Optional[bool] = None):
    """``ttnn.SDPAProgramConfig`` over the device's whole compute grid (read from the device, never hard-coded)
    with the :func:`chunk_config` chunks unless ``q_chunk`` / ``k_chunk`` are given."""
    import ttnn

    q_default, k_default = chunk_config(sq, sk)
    q_chunk = int(q_chunk or q_default)
    k_chunk = int(k_chunk or k_default)
    for name, value in (("q_chunk", q_chunk), ("k_chunk", k_chunk)):
        if value <= 0 or value % TILE:
            raise ValueError(f"{name}={value}: must be a positive multiple of {TILE}")
    return ttnn.SDPAProgramConfig(compute_with_storage_grid_size=device.compute_with_storage_grid_size(),
                                  q_chunk_size=q_chunk, k_chunk_size=k_chunk, exp_approx_mode=exp_approx_mode)


def sdpa_compute_config(*, fp32_acc: bool = False, fidelity: str = "HiFi2", approx: bool = True):
    """Explicit SDPA compute config. The default equals the op's own (HiFi2, bf16 DST, approximate exp): P7 found
    neither HiFi4 nor exact exp more accurate; ``fp32_acc=True`` (the op's standard instead of its streaming kernel)
    halves the max error (0.11 -> 0.056 over 36k keys) at 1.35-1.45x the time."""
    from ..precision import compute_kernel_config

    return compute_kernel_config(fidelity, fp32_acc=fp32_acc, approx=approx)


def _seq(t: Any) -> int:
    return int(t.shape[-2])


def sdpa(q: Any, k: Any, v: Any, *, scale: float, attn_mask: Any = None, program_config: Any = "auto",
         compute_kernel_config: Any = None, fp32_acc: bool = False, concat_heads: bool = False,
         memory_config: Any = None):
    """Non-causal scaled dot-product attention ``softmax(scale * Q K^T + attn_mask) V``.

    ``q`` ``[B, H, Sq, D]``, ``k`` / ``v`` ``[B, Hkv, Sk, D]`` (bf16 / bfp8 / bfp4, TILE, interleaved, D a multiple of
    32); ``scale`` is required: ``1/sqrt(logical head dim)`` (never the op's padded-head default). ``attn_mask``:
    additive ``[1|B, 1|H, Sq, Sk]`` bf16 TILE tensor in DRAM, or None; with a mask ``Sk`` must be a multiple of 32
    (module docstring). ``program_config``: ``"auto"`` (the
    :func:`chunk_config` chunks over the whole grid), an ``SDPAProgramConfig``, or None (op default; refused once
    Sk >= :data:`PROGRAM_CONFIG_MIN_SK`, P7). ``compute_kernel_config`` overrides :func:`sdpa_compute_config`
    (``fp32_acc``). ``concat_heads=True`` returns ``[B, 1, Sq, H*D]`` instead of ``[B, H, Sq, D]`` (no separate
    concat program). Returns the device tensor.
    """
    import ttnn

    scale = float(scale)
    if not math.isfinite(scale) or scale <= 0.0:
        raise ValueError(f"scale must be a positive finite float, got {scale!r}")
    if len(q.shape) != 4 or len(k.shape) != 4 or len(v.shape) != 4:
        raise ValueError(f"sdpa expects [B, H, S, D] tensors, got q{tuple(q.shape)} k{tuple(k.shape)} "
                         f"v{tuple(v.shape)}")
    d = int(q.shape[-1])
    if d % HEAD_DIM_MULTIPLE:
        raise ValueError(f"head dim {d} is not a multiple of {HEAD_DIM_MULTIPLE}: zero-pad it in the projection "
                         "weights (pad_head_columns / pad_head_rows) and pass scale=1/sqrt(logical head dim)")
    sq, sk = _seq(q), _seq(k)
    if attn_mask is not None and sk % TILE:
        raise ValueError(f"a user mask needs a tile-aligned key count (Sk={sk}): the op masks the padded keys of the "
                         "last partial tile only through the mask's own tile padding (0 after a host tilize), so "
                         f"they would join the softmax. Run on aligned_keys({sk}) = {aligned_keys(sk)} keys and mask "
                         "the extra ones with -inf (key_bias on the padded validity)")
    if isinstance(program_config, str):
        if program_config != "auto":
            raise ValueError(f"program_config={program_config!r}: expected 'auto', a config object or None")
        program_config = sdpa_program_config(q.device(), sq, sk)
    elif program_config is None and sk >= PROGRAM_CONFIG_MIN_SK:
        raise ValueError(f"Sk={sk} >= {PROGRAM_CONFIG_MIN_SK}: the op's default 32/32 chunks are slow and inaccurate "
                         "for long keys (probe P7); pass program_config='auto' or an SDPAProgramConfig")
    if compute_kernel_config is None:
        compute_kernel_config = sdpa_compute_config(fp32_acc=fp32_acc)
    kwargs: Dict[str, Any] = {}
    if concat_heads:
        kwargs["output_concat_heads"] = True
    return ttnn.transformer.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_mask, is_causal=False, scale=scale, memory_config=memory_config,
        program_config=program_config, compute_kernel_config=compute_kernel_config, **kwargs)


# ---------------------------------------------------------------------------------------------------- heads

def _check_4d_heads(x: Any, what: str) -> None:
    if len(x.shape) != 4 or int(x.shape[1]) != 1:
        raise ValueError(f"{what} must be [B, 1, S, width] (a 4-D projection output), got {tuple(x.shape)}")


def split_heads(x: Any, num_heads: int, *, memory_config: Any = None):
    """``[B, 1, S, H*D]`` -> ``[B, H, S, D]`` (one ``nlp_create_qkv_heads`` program with no K / V heads; the logical
    S is kept, so unaligned sequences stay unaligned for :func:`sdpa`)."""
    import ttnn

    _check_4d_heads(x, "split_heads input")
    width = int(x.shape[-1])
    if width % num_heads or (width // num_heads) % HEAD_DIM_MULTIPLE:
        raise ValueError(f"width {width} does not split into {num_heads} heads of a multiple of {HEAD_DIM_MULTIPLE}")
    out, _, _ = ttnn.experimental.nlp_create_qkv_heads(x, num_heads=num_heads, num_kv_heads=0,
                                                       transpose_k_heads=False, memory_config=memory_config)
    return out


def split_qkv(qkv: Any, num_heads: int, *, num_kv_heads: Optional[int] = None, memory_config: Any = None):
    """Fused ``[B, 1, S, (H + 2 Hkv) * D]`` (Q | K | V column blocks, heads contiguous) -> ``(q, k, v)``
    ``[B, H|Hkv, S, D]`` in one program (K not transposed)."""
    import ttnn

    _check_4d_heads(qkv, "split_qkv input")
    kv_heads = num_heads if num_kv_heads is None else int(num_kv_heads)
    width = int(qkv.shape[-1])
    if width % (num_heads + 2 * kv_heads) or (width // (num_heads + 2 * kv_heads)) % HEAD_DIM_MULTIPLE:
        raise ValueError(f"width {width} does not split into {num_heads} + 2 x {kv_heads} heads of a multiple of "
                         f"{HEAD_DIM_MULTIPLE}")
    return ttnn.experimental.nlp_create_qkv_heads(qkv, num_heads=num_heads, num_kv_heads=kv_heads,
                                                  transpose_k_heads=False, memory_config=memory_config)


def split_q_kv(q: Any, kv: Any, num_heads: int, *, num_kv_heads: Optional[int] = None, memory_config: Any = None):
    """Separate ``q`` ``[B, 1, S, H*D]`` and ``kv`` ``[B, 1, S, 2*Hkv*D]`` (K | V column blocks) with the SAME S
    (self-attention with Q and K / V from different inputs, e.g. a pre-LN Q next to an un-normalised K / V)
    -> ``(q, k, v)`` heads in one program. Cross-attention with another key length: :func:`split_heads` each."""
    import ttnn

    _check_4d_heads(q, "split_q_kv q")
    _check_4d_heads(kv, "split_q_kv kv")
    if int(q.shape[-2]) != int(kv.shape[-2]):
        raise ValueError(f"split_q_kv needs the same sequence length (Q {q.shape[-2]}, K/V {kv.shape[-2]})")
    kv_heads = num_heads if num_kv_heads is None else int(num_kv_heads)
    return ttnn.experimental.nlp_create_qkv_heads(q, kv, num_heads=num_heads, num_kv_heads=kv_heads,
                                                  transpose_k_heads=False, memory_config=memory_config)


def merge_heads(x: Any, *, memory_config: Any = None):
    """``[B, H, S, D]`` -> ``[B, 1, S, H*D]`` (``nlp_concat_heads``). Prefer ``sdpa(..., concat_heads=True)``."""
    import ttnn

    if len(x.shape) != 4:
        raise ValueError(f"merge_heads expects [B, H, S, D], got {tuple(x.shape)}")
    return ttnn.experimental.nlp_concat_heads(x, memory_config=memory_config)


# ---------------------------------------------------------------------------------------------------- masks

def key_bias_row(valid: Sequence[bool], *, dtype: Any = np.float32) -> np.ndarray:
    """``[1, 1, 1, Sk]`` additive key bias: 0 for valid keys, ``-inf`` for masked ones (bf16-exact values)."""
    v = np.asarray(valid, dtype=bool).reshape(-1)
    if v.size == 0:
        raise ValueError("key_bias_row: empty validity vector")
    if not v.any():
        raise ValueError("key_bias_row: every key is masked (the softmax of a fully masked row is undefined)")
    return np.where(v, 0.0, -np.inf).astype(dtype).reshape(1, 1, 1, -1)


def key_bias(valid: Sequence[bool], sq: int, *, dtype: Any = np.float32) -> np.ndarray:
    """``[1, 1, Sq, Sk]`` additive key mask for :func:`sdpa` (identical rows), built on the host. Upload it as a
    bf16 TILE tensor in DRAM (the op's requirement), or keep a :func:`key_bias_row` input and expand it on the
    device with :func:`expand_key_bias`."""
    row = key_bias_row(valid, dtype=dtype)
    return np.ascontiguousarray(np.broadcast_to(row, (1, 1, int(sq), row.shape[-1])))


def expand_key_bias(row: Any, sq: int, *, memory_config: Any = None):
    """Device-side :func:`key_bias`: repeat a ``[1, 1, 1, Sk]`` bf16 TILE bias row to ``[1, 1, Sq, Sk]`` in DRAM
    (one program; traceable)."""
    import ttnn

    if len(row.shape) != 4 or int(row.shape[-2]) != 1:
        raise ValueError(f"expand_key_bias expects a [1, 1, 1, Sk] row, got {tuple(row.shape)}")
    return ttnn.repeat(row, [1, 1, int(sq), 1],
                       memory_config=ttnn.DRAM_MEMORY_CONFIG if memory_config is None else memory_config)


# ---------------------------------------------------------------------------------------------------- head padding

def pad_head_dim(x: Any, padded: Optional[int] = None) -> np.ndarray:
    """Zero-pad the last (head) dim of a host array ``[..., D]`` to ``padded`` (default :func:`padded_head_dim`)."""
    a = np.asarray(x)
    d = a.shape[-1]
    p = padded_head_dim(d) if padded is None else int(padded)
    if p < d:
        raise ValueError(f"padded head dim {p} < head dim {d}")
    if p == d:
        return a
    return np.pad(a, [(0, 0)] * (a.ndim - 1) + [(0, p - d)])


def pad_head_columns(w: Any, b: Any, num_heads: int, head_dim: int,
                     padded: Optional[int] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]:
    """Projection ``[in, H*D]`` (``y = x @ w + b``) -> ``[in, H*Dp]`` with zero columns after each head's D (and the
    bias likewise), so the projection itself emits padded heads (exact: the padded columns are 0)."""
    w = np.asarray(w)
    p = padded_head_dim(head_dim) if padded is None else int(padded)
    if w.shape[-1] != num_heads * head_dim:
        raise ValueError(f"weight has {w.shape[-1]} columns, expected {num_heads} x {head_dim}")
    wh = w.reshape(w.shape[:-1] + (num_heads, head_dim))
    wp = np.zeros(w.shape[:-1] + (num_heads, p), w.dtype)
    wp[..., :head_dim] = wh
    bp = None
    if b is not None:
        b = np.asarray(b).reshape(num_heads, head_dim)
        bp = np.zeros((num_heads, p), b.dtype)
        bp[:, :head_dim] = b
        bp = bp.reshape(-1)
    return wp.reshape(w.shape[:-1] + (num_heads * p,)), bp


def pad_head_rows(w: Any, num_heads: int, head_dim: int, padded: Optional[int] = None) -> np.ndarray:
    """Output projection ``[H*D, out]`` -> ``[H*Dp, out]`` with zero rows at each head's padded positions, matching
    the ``[.., H*Dp]`` concatenated-heads output of a padded attention."""
    w = np.asarray(w)
    p = padded_head_dim(head_dim) if padded is None else int(padded)
    if w.shape[0] != num_heads * head_dim:
        raise ValueError(f"weight has {w.shape[0]} rows, expected {num_heads} x {head_dim}")
    wp = np.zeros((num_heads, p) + w.shape[1:], w.dtype)
    wp[:, :head_dim] = w.reshape((num_heads, head_dim) + w.shape[1:])
    return wp.reshape((num_heads * p,) + w.shape[1:])


# ---------------------------------------------------------------------------------------------------- oracle / fallback

def attention_reference(q: Any, k: Any, v: Any, *, scale: float, bias: Any = None) -> np.ndarray:
    """float64 numpy ``softmax(scale Q K^T + bias) V`` over ``[B, H, S, D]`` arrays (K / V heads broadcast when
    ``Hkv`` divides ``H``); ``bias`` broadcasts to ``[B, H, Sq, Sk]`` (``-inf`` allowed)."""
    q = np.asarray(q, np.float64)
    k = np.asarray(k, np.float64)
    v = np.asarray(v, np.float64)
    if k.shape[1] != q.shape[1]:
        rep = q.shape[1] // k.shape[1]
        k, v = np.repeat(k, rep, axis=1), np.repeat(v, rep, axis=1)
    s = np.einsum("bhqd,bhkd->bhqk", q, k) * float(scale)
    if bias is not None:
        s = s + np.asarray(bias, np.float64)
    s = s - s.max(axis=-1, keepdims=True)
    p = np.exp(s)
    p /= p.sum(axis=-1, keepdims=True)
    return np.einsum("bhqk,bhkd->bhqd", p, v)


def attention_matmul(q: Any, k: Any, v: Any, *, scale: float, attn_mask: Any = None,
                     compute_kernel_config: Any = None):
    """Small-MHA fallback for what :func:`sdpa` rejects (e.g. fp32 Q / K / V): ``softmax(scale Q K^T + mask) V``
    with two matmuls. ``q`` ``[B, H, Sq, D]``, ``k`` / ``v`` ``[B, H, Sk, D]`` with a tile-aligned ``Sk`` (pad the
    keys and mask the padding with ``-inf``: a softmax over tile padding would include it). Materialises the
    ``[B, H, Sq, Sk]`` scores, so keep it for small shapes."""
    import ttnn

    from ..precision import compute_kernel_config as ckc

    sk = int(k.shape[-2])
    if sk % TILE:
        raise ValueError(f"attention_matmul needs a tile-aligned Sk (got {sk}): pad the keys and mask them")
    cfg = compute_kernel_config if compute_kernel_config is not None else ckc("HiFi4", fp32_acc=True)
    scores = ttnn.matmul(q, k, transpose_b=True, compute_kernel_config=cfg)
    scores = ttnn.multiply(scores, float(scale))
    if attn_mask is not None:
        scores = ttnn.add(scores, attn_mask)
    probs = ttnn.softmax(scores, dim=-1, numeric_stable=True, compute_kernel_config=cfg)
    return ttnn.matmul(probs, v, compute_kernel_config=cfg)