# SPDX-License-Identifier: Apache-2.0
"""Attention score scale + additive mask as one ``ttnn.generic_op`` (``ATTN_SMASK``, OPT round 2 item 4b).

``scale_mask(s, scale, mask)``: fp32 TILE scores ``[1, H, Sq, Sk]`` and a ``[1, 1, Sq, Sk]`` mask (bf16 or fp32,
broadcast over the heads) -> ``s * scale + mask`` fp32, the two stock ``binary_ng`` programs of
``tt/attention.py`` (``ttnn.multiply(s, scale)``, ``ttnn.add(., mask)``) in one pass over the scores
(``kernels/smask_*.cpp``): bit-identical, half the DRAM traffic and one program less per attention. A bf16 mask
sends the stock fp32 + bf16 add down binary_ng's FPU path, which truncates the scores to TF32 (measured); the
kernel reproduces that truncation (the mask holds 0 / -inf only, so the add itself is exact).
The mask-plane tiles are split in contiguous ranges over the grid; each core streams the H score tiles of its mask
tiles.
"""
from __future__ import annotations

import os
import struct
from typing import Any

__all__ = ["scale_mask", "supported"]

_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
TB = 4096
N_RT = 2


def supported(s: Any, mask: Any) -> bool:
    import ttnn

    try:
        ss, ms = list(s.padded_shape), list(mask.padded_shape)
        return (hasattr(ttnn, "generic_op") and s.dtype == ttnn.float32 and s.layout == ttnn.TILE_LAYOUT
                and mask.layout == ttnn.TILE_LAYOUT and mask.dtype in (ttnn.bfloat16, ttnn.float32)
                and not s.is_sharded() and not mask.is_sharded() and len(ss) == 4 and len(ms) == 4
                and ss[0] == 1 and ss[1] % 2 == 0 and ms[0] == 1 and ms[1] == 1 and ss[2:] == ms[2:])
    except Exception:  # noqa: BLE001 - the host fake ttnn
        return False


def scale_mask(s: Any, scale: float, mask: Any):
    import ttnn

    dev = s.device()
    _, H, Sq, Sk = list(s.padded_shape)
    plane = (Sq // 32) * (Sk // 32)
    out = ttnn.allocate_tensor_on_device(s.shape, ttnn.float32, ttnn.TILE_LAYOUT, dev, ttnn.DRAM_MEMORY_CONFIG)
    g = dev.compute_with_storage_grid_size()
    n = min(plane, g.x * g.y)
    cs = [(i % g.x, i // g.x) for i in range(n)]
    full, rem = divmod(n, g.x)
    rs = []
    if full:
        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(g.x - 1, full - 1)))
    if rem:
        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, full), ttnn.CoreCoord(rem - 1, full)))
    crs = ttnn.CoreRangeSet(set(rs))
    base, extra = divmod(plane, n)
    rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
    m0 = 0
    for i, (cx, cy) in enumerate(cs):
        k = base + (1 if i < extra else 0)
        rd[cx][cy] = [m0, k]
        wr[cx][cy] = [m0, k]
        cp[cx][cy] = [k, 0]
        m0 += k
    mbf = mask.dtype == ttnn.bfloat16
    tbm = 2048 if mbf else 4096

    def cb(idx, pages, dt, tb):
        return ttnn.CBDescriptor(total_size=pages * tb, core_ranges=crs, format_descriptors=[
            ttnn.CBFormatDescriptor(buffer_index=idx, data_format=dt, page_size=tb)])

    cbs = [cb(0, 2 * H, ttnn.float32, TB), cb(1, 2, mask.dtype, tbm), cb(2, 1, ttnn.float32, TB),
           cb(16, 4, ttnn.float32, TB)]
    um = [ttnn.UnpackToDestMode.Default] * 64
    um[0] = um[2] = ttnn.UnpackToDestMode.UnpackToDestFp32
    if not mbf:
        um[1] = ttnn.UnpackToDestMode.UnpackToDestFp32
    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
                                        math_approx_mode=False)
    ccfg.unpack_to_dest_mode = um
    bits = int.from_bytes(struct.pack("<f", float(scale)), "little")
    SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH

    def acc(t):
        return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())

    def kd(name, ct, rt, common, config):
        return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
                                     compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
                                     config=config)

    ks = [kd("smask_reader.cpp", [H, plane, bits, N_RT, tbm] + acc(s) + acc(mask), rd,
             [s.buffer_address(), mask.buffer_address()], ttnn.ReaderConfigDescriptor()),
          kd("smask_writer.cpp", [H, plane, N_RT] + acc(out), wr, [out.buffer_address()],
             ttnn.WriterConfigDescriptor()),
          kd("smask_compute.cpp", [H, N_RT, int(mbf)], cp, [], ccfg)]
    ttnn.generic_op([s, mask, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
    return out