kernbench2/tests/test_tl_load_lazy.py

"""Phase 1 spec test for lazy ``tl.load`` (ADR-0062).

ADR-0062 §D1/§D2: ``tl.load`` is non-blocking. It issues a ``DmaReadCmd``
and returns immediately with a ``TensorHandle`` carrying a pending
``LoadFuture``. The runtime auto-inserts a wait on that future at the
first consuming op (``tl.dot``, MATH ops, ``tl.store``, ``tl.send``,
``tl.copy_to``, ``tl.composite`` operands). Symmetric with the existing
``recv_async`` / ``RecvFuture`` pattern, generalised to HBM loads.

Phase 1 (this commit): tests only — production code lands in Phase 2.
All tests currently fail because:
  - ``TensorHandle`` has no ``pending`` field
  - ``LoadFuture`` class doesn't exist
"""
from __future__ import annotations

from kernbench.common.pe_commands import DmaReadCmd
from kernbench.triton_emu.tl_context import TLContext


def _ctx() -> TLContext:
    """TLContext with a real scratch base so _make_compute_out allocates."""
    return TLContext(
        pe_id=0, num_programs=1, dispatch_cycles=0,
        scratch_base=1 << 24, scratch_size=1 << 20,
    )


# ── 1. tl.load handle carries a `pending` field ──────────────────────


def test_tl_load_handle_carries_pending_field():
    """ADR-0062 §D2: tl.load handle has a ``pending`` attribute that is
    not None — it references the LoadFuture the consumer will await."""
    ctx = _ctx()
    handle = ctx.load(0x1000, shape=(8, 64), dtype="f16")
    assert hasattr(handle, "pending"), (
        "TensorHandle must have a `pending` field per ADR-0062 §D2"
    )
    assert handle.pending is not None, (
        "tl.load handle must carry a non-None pending (LoadFuture)"
    )


# ── 2. Constant / math-output handles have pending=None ──────────────


def test_constant_handle_pending_is_none():
    """Non-loaded handles (tl.zeros, tl.full, _make_compute_out, tl.arange)
    have ``pending=None`` — there is no DMA to wait on. Consumer ops can
    safely skip auto-wait for these inputs."""
    ctx = _ctx()
    z = ctx.zeros((8, 64), dtype="f16")
    assert getattr(z, "pending", "MISSING") is None, (
        "tl.zeros handle must have pending=None"
    )
    f = ctx.full((8, 64), value=0.0, dtype="f16")
    assert getattr(f, "pending", "MISSING") is None, (
        "tl.full handle must have pending=None"
    )
    a = ctx.arange(0, 8, dtype="i32")
    assert getattr(a, "pending", "MISSING") is None, (
        "tl.arange handle must have pending=None"
    )


def test_compute_out_pending_is_none():
    """Scratch-allocated output handles (via _make_compute_out) have
    pending=None. Math ops produce them; no DMA → no auto-wait needed
    when they're consumed downstream."""
    ctx = _ctx()
    out = ctx._make_compute_out(shape=(8, 64), dtype="f16")
    assert getattr(out, "pending", "MISSING") is None, (
        "compute-out handle must have pending=None"
    )


# ── 3. LoadFuture class exists ───────────────────────────────────────


def test_loadfuture_class_exists():
    """ADR-0062 §D2: LoadFuture class wraps the DMA done event and a
    resolved-flag, analogous to RecvFuture for IPCQ. Lives in
    tl_context.py to keep all greenlet bridge types in one module."""
    from kernbench.triton_emu.tl_context import LoadFuture  # added Phase 2
    assert LoadFuture is not None


# ── 4. Distinct loads produce distinct LoadFutures ───────────────────


def test_distinct_loads_have_distinct_pending():
    """ADR-0062 §D2: two tl.load calls (different addresses) produce two
    independent LoadFuture objects. Each consumer must wait on the
    right one — sharing the same future would over-serialise."""
    ctx = _ctx()
    h1 = ctx.load(0x1000, shape=(8,), dtype="f16")
    h2 = ctx.load(0x2000, shape=(8,), dtype="f16")
    assert h1.pending is not None and h2.pending is not None
    assert h1.pending is not h2.pending, (
        "distinct loads must have distinct LoadFuture objects"
    )


# ── 5. LoadFuture starts unresolved ──────────────────────────────────


def test_loadfuture_starts_unresolved():
    """ADR-0062 §D2: a fresh LoadFuture has resolved=False — the consumer
    op must await it before reading. Once awaited (in greenlet mode by
    auto-wait at first use), the SimPy event triggers and a subsequent
    consumer can skip the yield."""
    ctx = _ctx()
    h = ctx.load(0x1000, shape=(8,), dtype="f16")
    assert h.pending.resolved is False, (
        "LoadFuture must start unresolved; got resolved=True at issue time"
    )


# ── 6. op_log: dma_read_count unchanged from blocking version ────────


def test_op_log_dma_read_count_unchanged():
    """ADR-0062 §D2 (last paragraph) / Test Req 4: each tl.load still
    emits exactly one DmaReadCmd. The asynchrony is a scheduling
    property, not a new op kind — dma_read_count and existing op_log
    consumers see no change."""
    ctx = _ctx()
    ctx.load(0x1000, shape=(8, 64), dtype="f16")
    ctx.load(0x2000, shape=(8, 64), dtype="f16")
    ctx.load(0x3000, shape=(8, 64), dtype="f16")
    dma_cmds = [c for c in ctx.commands if isinstance(c, DmaReadCmd)]
    assert len(dma_cmds) == 3, (
        f"lazy tl.load must still emit one DmaReadCmd per call; "
        f"got {len(dma_cmds)}"
    )


# ── 7. Handle metadata fields preserved across the lazy change ───────


def test_tl_load_handle_metadata_unchanged():
    """ADR-0062 §D1: the tl surface is unchanged. The handle's addr,
    shape, dtype, nbytes, space, pinned fields are identical to the
    blocking version — only `pending` is new."""
    ctx = _ctx()
    h = ctx.load(0x1000, shape=(8, 64), dtype="f16")
    assert h.addr == 0x1000
    assert h.shape == (8, 64)
    assert h.dtype == "f16"
    assert h.nbytes == 8 * 64 * 2
    assert h.space == "hbm"
    assert h.pinned is True


# ── 8. LoadFuture carries the DmaReadCmd it wraps ────────────────────


def test_loadfuture_carries_cmd():
    """ADR-0062 §D2: LoadFuture references the DmaReadCmd it wraps so
    the runtime can resolve handle → command → completion event at
    auto-wait time. Mirrors RecvFuture.cmd."""
    ctx = _ctx()
    h = ctx.load(0x1000, shape=(8, 64), dtype="f16")
    assert hasattr(h.pending, "cmd"), "LoadFuture.cmd must exist"
    assert isinstance(h.pending.cmd, DmaReadCmd)
    assert h.pending.cmd.handle is h