kernbench2/tests/test_pe_scheduler_strict_fifo.py

"""Phase 1 spec tests for ADR-0065 P4 — strict-FIFO RW hazard tracker.

PE_SCHEDULER serializes composites that share read/write handles
(ADR-0065 D6.3 / DDD §8): a new composite whose read or write handles
intersect an in-flight composite's ``rw_handles`` waits until all earlier
overlapping composites complete. Non-overlapping composites proceed.

These unit-test the ``_RwHazardTracker`` directly with a SimPy env. Legacy
composites (``rw_handles == ()``) never register and never block, so the
existing benches' Stage sequences are untouched.

Phase 1 (this commit): tests only. FAIL until P4:
  - ``_RwHazardTracker`` does not exist in ``pe_scheduler``.
"""
from __future__ import annotations

import simpy

from kernbench.common.pe_commands import (
    CompletionHandle,
    CompositeCmd,
    OpSpec,
    Scope,
    TensorHandle,
)


def _h(addr: int) -> TensorHandle:
    return TensorHandle(id=f"h{addr}", addr=addr, shape=(8,), dtype="f16",
                        nbytes=16, space="tcm")


def _cmd(cid: str, rw=(), reads=()) -> CompositeCmd:
    """A composite that writes ``rw`` (rw_handles) and reads ``reads``
    (as MATH-op operands)."""
    if reads:
        ops = tuple(
            OpSpec(kind="math", scope=Scope.KERNEL, operands={"src": r}, out=r)
            for r in reads
        )
    else:
        ops = (OpSpec(kind="math", scope=Scope.KERNEL, operands={}, out=None),)
    return CompositeCmd(completion=CompletionHandle(id=cid), ops=ops,
                        rw_handles=tuple(rw))


def _tracker():
    from kernbench.components.builtin.pe_scheduler import _RwHazardTracker

    return _RwHazardTracker()


def test_admit_immediate_when_no_inflight():
    env = simpy.Environment()
    tr = _tracker()
    state: dict = {}

    def driver():
        yield from tr.admit(env, _cmd("a", rw=(_h(1),)), env.event())
        state["t"] = env.now

    env.process(driver())
    env.run()
    assert state["t"] == 0


def test_admit_blocks_until_overlapping_inflight_retires():
    env = simpy.Environment()
    tr = _tracker()
    O = _h(1)
    eA = env.event()
    state: dict = {}

    def b_proc():
        # B reads O, which in-flight A writes (RAW) → must block.
        yield from tr.admit(env, _cmd("b", reads=(O,)), env.event())
        state["B"] = env.now

    def driver():
        yield from tr.admit(env, _cmd("a", rw=(O,)), eA)
        state["A"] = env.now
        env.process(b_proc())
        yield env.timeout(10)
        tr.retire("a")
        eA.succeed()                      # A completes at t=10
        yield env.timeout(1)

    env.process(driver())
    env.run()
    assert state["A"] == 0
    assert state["B"] >= 10, f"B admitted at {state.get('B')}, must wait for A"


def test_admit_no_block_when_disjoint():
    env = simpy.Environment()
    tr = _tracker()
    O, X = _h(1), _h(2)
    state: dict = {}

    def driver():
        yield from tr.admit(env, _cmd("a", rw=(O,)), env.event())
        yield from tr.admit(env, _cmd("b", reads=(X,)), env.event())
        state["B"] = env.now

    env.process(driver())
    env.run()
    assert state["B"] == 0, "disjoint composite must not block"


def test_write_after_write_overlap_blocks():
    """Two composites both writing the same handle (WAW) serialize."""
    env = simpy.Environment()
    tr = _tracker()
    O = _h(1)
    eA = env.event()
    state: dict = {}

    def b_proc():
        yield from tr.admit(env, _cmd("b", rw=(O,)), env.event())
        state["B"] = env.now

    def driver():
        yield from tr.admit(env, _cmd("a", rw=(O,)), eA)
        env.process(b_proc())
        yield env.timeout(5)
        tr.retire("a")
        eA.succeed()
        yield env.timeout(1)

    env.process(driver())
    env.run()
    assert state["B"] >= 5


def test_legacy_composite_never_registers():
    """A composite with empty rw_handles must not block a later composite
    that touches the same input handles (legacy composites carry no RW
    metadata → no serialization, byte-equal behavior)."""
    env = simpy.Environment()
    tr = _tracker()
    O = _h(1)
    state: dict = {}

    def driver():
        # A is "legacy": reads O but declares no rw_handles → not registered.
        yield from tr.admit(env, _cmd("a", reads=(O,)), env.event())
        # B reads O too; since A registered nothing, B must not block.
        yield from tr.admit(env, _cmd("b", reads=(O,)), env.event())
        state["B"] = env.now

    env.process(driver())
    env.run()
    assert state["B"] == 0