kernbench2/tests/test_tl_scratch_scope.py

"""Phase 1 spec test for ``tl.scratch_scope`` (ADR-0063, P3a).

ADR-0063 D1: ``tl.scratch_scope()`` is a context manager whose ``__enter__``
snapshots ``_scratch_cursor`` and whose ``__exit__`` restores it, freeing
every handle allocated inside the ``with``-block. This pins the bump
allocator behaviour without changing any command emission.

Phase 1 (this commit): tests only — production code lands in Phase 2.
"""
from __future__ import annotations

from kernbench.triton_emu.tl_context import TLContext


# Configured TLContext with a real (non-zero) scratch base — required for
# _make_compute_out to allocate, see tl_context.py:_scratch_alloc.
def _ctx() -> TLContext:
    return TLContext(
        pe_id=0, num_programs=1,
        dispatch_cycles=0,
        scratch_base=1 << 24,         # 16 MiB base
        scratch_size=1 << 20,         # 1 MiB pool
    )


# ── 1. Method exists and returns a usable ctx-mgr ─────────────────────


def test_scratch_scope_method_exists():
    ctx = _ctx()
    assert hasattr(ctx, "scratch_scope"), (
        "TLContext must expose scratch_scope() per ADR-0063 D1"
    )
    with ctx.scratch_scope() as scope:
        assert scope is not None


# ── 2. Recycling within a loop bounds peak cursor ────────────────────


def test_scratch_scope_recycles_within_loop():
    """ADR-0063 D1 / Test Req 1: a loop inside scratch_scope keeps peak
    _scratch_cursor bounded by one iteration's footprint."""
    ctx = _ctx()
    # One iteration footprint: allocate three small handles via
    # _make_compute_out.
    one_iter_peaks = []
    for _ in range(10):
        with ctx.scratch_scope():
            ctx._make_compute_out(shape=(64,), dtype="f16")
            ctx._make_compute_out(shape=(64,), dtype="f16")
            ctx._make_compute_out(shape=(64,), dtype="f16")
            one_iter_peaks.append(ctx._scratch_cursor)
    # After the loop the cursor must be at its post-loop start (0 here),
    # not 10× one iteration.
    assert ctx._scratch_cursor == 0, (
        f"cursor must reset on each scope exit; final cursor = "
        f"{ctx._scratch_cursor}"
    )
    # Every iteration sees the same peak — proves recycling.
    assert len(set(one_iter_peaks)) == 1, (
        f"every iteration must reach the same peak (recycled); "
        f"got distinct peaks {one_iter_peaks}"
    )


# ── 3. Allocations outside the scope are not recycled ────────────────


def test_scratch_scope_preserves_outside_allocations():
    """ADR-0063 D3: handles allocated outside the scope (persistent
    arena) keep their addresses; only inside-scope allocations are
    rewound."""
    ctx = _ctx()
    # Persistent allocation BEFORE the scope.
    persistent = ctx._make_compute_out(shape=(64,), dtype="f16")
    cursor_after_persistent = ctx._scratch_cursor
    with ctx.scratch_scope():
        ctx._make_compute_out(shape=(64,), dtype="f16")
        ctx._make_compute_out(shape=(64,), dtype="f16")
    # On exit the cursor must rewind to the save-point (which equals
    # cursor_after_persistent), NOT to 0 — the persistent allocation is
    # preserved.
    assert ctx._scratch_cursor == cursor_after_persistent, (
        f"cursor must rewind to scope-enter point ({cursor_after_persistent}), "
        f"not 0; got {ctx._scratch_cursor}"
    )
    # A new allocation after the scope must not collide with the
    # persistent one.
    fresh = ctx._make_compute_out(shape=(64,), dtype="f16")
    assert fresh.addr != persistent.addr, (
        "post-scope allocation must not reuse the persistent address"
    )


# ── 4. Nested scopes rewind to their own save-points ─────────────────


def test_scratch_scope_nesting():
    """ADR-0063 D4: nested scopes rewind to the matching enter-point."""
    ctx = _ctx()
    with ctx.scratch_scope():
        ctx._make_compute_out(shape=(64,), dtype="f16")
        outer_save = ctx._scratch_cursor
        with ctx.scratch_scope():
            ctx._make_compute_out(shape=(64,), dtype="f16")
            ctx._make_compute_out(shape=(64,), dtype="f16")
            inner_peak = ctx._scratch_cursor
        # Inner exit: rewind to outer_save.
        assert ctx._scratch_cursor == outer_save, (
            f"inner exit must rewind to {outer_save}; "
            f"got {ctx._scratch_cursor}"
        )
        assert inner_peak > outer_save, (
            "sanity: inner scope must have allocated something"
        )
    # Outer exit: rewind all the way to 0.
    assert ctx._scratch_cursor == 0, (
        f"outer exit must rewind to 0; got {ctx._scratch_cursor}"
    )


# ── 5. Control: without scope, allocations pile up ───────────────────


def test_without_scope_grows_unbounded():
    """Sanity check that the bump allocator without scratch_scope grows
    linearly — the failure mode ADR-0063 §A.2 cites for the S=16 cap."""
    ctx = _ctx()
    for _ in range(10):
        ctx._make_compute_out(shape=(64,), dtype="f16")
        ctx._make_compute_out(shape=(64,), dtype="f16")
        ctx._make_compute_out(shape=(64,), dtype="f16")
    # 30 allocations × aligned(64 * 2 = 128 → 128) B = 3840 B
    assert ctx._scratch_cursor >= 30 * 128, (
        f"without scope, cursor must grow with allocation count; got "
        f"{ctx._scratch_cursor}"
    )


# ── 6. Overflow avoided by recycling ─────────────────────────────────


def test_overflow_avoided_with_scope():
    """ADR-0063 Test Req 3: a loop that would overflow 1 MiB without
    scope completes when wrapped in scratch_scope."""
    # 1 MiB pool / one alloc of 64 KiB = 16 max without scope.
    # 100 iterations of 64 KiB without recycling = 6.4 MiB → overflow.
    # With scratch_scope: peak ≈ 64 KiB per iter, well under 1 MiB.
    ctx = _ctx()
    big_shape = (32 * 1024,)  # 32 K elements × 2 B = 64 KiB per handle
    for _ in range(100):
        with ctx.scratch_scope():
            ctx._make_compute_out(shape=big_shape, dtype="f16")
    # Completes without raising — the recycle keeps us under the cap.
    assert ctx._scratch_cursor == 0


# ── 7. scratch_scope + copy_to integration (ADR-0063 §D3.1) ──────────


def test_scoped_loop_with_copy_to_keeps_cursor_bounded():
    """ADR-0063 §D3.1: the two-arena flash pattern.

    Persistent ``running`` allocated OUTSIDE the scope; per-iteration
    work allocates inside; the last act before scope exit is
    ``tl.copy_to(running, new_running)`` to persist state. After the
    loop the cursor must equal the post-persistent-allocation value
    (running survived) and never have grown beyond that + one
    iteration's footprint.
    """
    ctx = _ctx()
    # Persistent arena: one running-state handle outside any scope.
    running = ctx._make_compute_out(shape=(8, 64), dtype="f16")
    cursor_after_persist = ctx._scratch_cursor
    assert running.addr != 0

    n_iters = 50
    for _ in range(n_iters):
        with ctx.scratch_scope():
            # Simulate per-tile work: a few scoped allocations.
            ctx._make_compute_out(shape=(8, 64), dtype="f16")
            ctx._make_compute_out(shape=(8, 64), dtype="f16")
            new_running = ctx._make_compute_out(shape=(8, 64), dtype="f16")
            # Persist back to the outside-scope handle so its bytes
            # survive __exit__.
            ctx.copy_to(running, new_running)
        # After scope exit cursor returns to the persistent footprint —
        # scoped allocations are recycled, running survives.
        assert ctx._scratch_cursor == cursor_after_persist, (
            f"after scope exit cursor must rewind to "
            f"{cursor_after_persist}; got {ctx._scratch_cursor}"
        )

    assert ctx._scratch_cursor == cursor_after_persist


def test_persistent_handle_survives_after_copy_to_in_scope():
    """ADR-0063 §D2 / §D3: a handle allocated outside a scratch_scope
    must keep its address even after a copy_to(...) targeting it from
    inside the scope and after that scope exits."""
    ctx = _ctx()
    persistent = ctx._make_compute_out(shape=(64,), dtype="f16")
    persistent_addr_before = persistent.addr

    with ctx.scratch_scope():
        scoped = ctx._make_compute_out(shape=(64,), dtype="f16")
        ctx.copy_to(persistent, scoped)
    # Out of scope; persistent's address is unchanged.
    assert persistent.addr == persistent_addr_before
    # A fresh allocation after the scope must not collide with persistent.
    fresh = ctx._make_compute_out(shape=(64,), dtype="f16")
    assert fresh.addr != persistent.addr