"""Engine-fix spec test — composite GEMM with an on-chip (TCM) operand.

Root cause: ``_make_compute_out`` (tl_context) builds compute-result
handles with ``space="tcm"`` but omits ``pinned=True``. The PE_SCHEDULER
tiling gate (tiling.py:306) skips the operand DMA_READ only when
``pinned`` is set, so a compute result fed as a composite head-GEMM
operand ``a`` (e.g. the attention probabilities into P·V) is wrongly
scheduled as an HBM ``DMA_READ`` of its bit-61 scratch address →
``PhysAddrError`` in data mode.

Every other TCM-resident handle in tl_context is already marked
``pinned=True`` for exactly this reason (composite auto-output, recipe
scratch / primary-out). ``_make_compute_out`` is the one site that
forgot. ``.pinned`` is read ONLY by the tiling DMA gate, so setting it on
compute outputs is side-effect-free and byte-equal for existing benches
(which never feed an unpinned compute output into a composite operand).

Phase 1: this test FAILS until the fix (a DMA_READ for operand A is
emitted for the on-chip ``a``).
"""
from __future__ import annotations

from kernbench.common.pe_commands import CompositeCmd
from kernbench.components.builtin.pe_types import StageType
from kernbench.components.builtin.tiling import generate_plan_from_ops
from kernbench.triton_emu.tl_context import TLContext, run_kernel


def _pv_composite_kernel(*, tl) -> None:
    """P·V composite whose ``a`` is a real on-chip compute result."""
    x = tl.load(0x1000, shape=(8, 1024), dtype="f16")
    a = tl.exp(x)                                   # compute output → TCM
    V = tl.ref(0x3000, shape=(1024, 128), dtype="f16")   # HBM, streamed
    O = tl.zeros((8, 128), dtype="f16")
    tl.composite(op="gemm", a=a, b=V, out=O)


def _pv_plan():
    tl = TLContext(pe_id=0, num_programs=1,
                   scratch_base=0x200000, scratch_size=1 << 20)
    run_kernel(_pv_composite_kernel, tl)
    comp = next(c for c in tl.commands if isinstance(c, CompositeCmd))
    return generate_plan_from_ops(
        ops=comp.ops, tile_m=32, tile_k=64, tile_n=32,
        bytes_per_element=2, pe_prefix="sip0.cube0.pe0",
    )


def test_onchip_operand_a_not_dma_read():
    """The on-chip ``a`` (a compute result) must NOT be streamed via a
    DMA_READ; only the HBM ref ``b`` (V) is."""
    plan = _pv_plan()
    read_operands = [
        s.params.get("operand")
        for t in plan.tiles for s in t.stages
        if s.stage_type == StageType.DMA_READ
    ]
    assert "A" not in read_operands, read_operands
    assert "B" in read_operands, read_operands