kernbench2/tests/test_pe_pipeline.py

"""Tests for ADR-0014 D6 PE pipeline: TileToken self-routing, pipeline overlap, e2e accuracy.

Test plan items:
  3. Phase 1 → Phase 2 end-to-end (op_log → DataExecutor → verify)
  4. TileToken self-routing (stage sequence, PipelineContext completion)
  5. Async pipeline overlap (intra-command tile overlap, FIFO ordering)
"""
from __future__ import annotations

import simpy

from kernbench.components.builtin.pe_types import (
    PipelineContext,
    PipelinePlan,
    Stage,
    StageType,
    TilePlan,
    TileToken,
)


# ── 4. TileToken self-routing ────────────────────────────────────────


def test_tile_token_advance():
    """TileToken.advance() increments stage_idx and returns next Stage."""
    stages = (
        Stage(StageType.DMA_READ, "pe_dma", {"src_addr": 0}),
        Stage(StageType.FETCH, "pe_fetch_store", {"direction": "read"}),
        Stage(StageType.GEMM, "pe_gemm", {"m": 32, "k": 64, "n": 32}),
    )
    plan = TilePlan(tile_id=0, stages=stages)
    ctx = PipelineContext(id="p1", total_tiles=1)
    token = TileToken(
        tile_id=0, pipeline_ctx=ctx, plan=plan,
        stage_idx=0, params=stages[0].params,
    )

    assert token.current_stage.stage_type == StageType.DMA_READ

    next_s = token.advance()
    assert next_s is not None
    assert next_s.stage_type == StageType.FETCH
    assert token.stage_idx == 1
    assert token.params == {"direction": "read"}

    next_s = token.advance()
    assert next_s is not None
    assert next_s.stage_type == StageType.GEMM
    assert token.stage_idx == 2

    # Last stage — advance returns None
    assert token.advance() is None
    assert token.stage_idx == 3


def test_pipeline_context_completion():
    """PipelineContext.complete_tile() fires done_event on last tile."""
    env = simpy.Environment()
    done = env.event()
    ctx = PipelineContext(id="p1", total_tiles=3, done_event=done)

    ctx.complete_tile()
    assert not done.triggered
    ctx.complete_tile()
    assert not done.triggered
    ctx.complete_tile()
    assert done.triggered


def test_pipeline_context_exactly_once():
    """PipelineContext tracks completed_tiles correctly."""
    ctx = PipelineContext(id="p1", total_tiles=2)
    assert ctx.completed_tiles == 0
    ctx.complete_tile()
    assert ctx.completed_tiles == 1
    ctx.complete_tile()
    assert ctx.completed_tiles == 2


def test_tile_token_self_routing_chain():
    """Simulated self-routing: component reads next stage from token."""
    stages = (
        Stage(StageType.DMA_READ, "dma", {}),
        Stage(StageType.FETCH, "fetch", {}),
        Stage(StageType.GEMM, "gemm", {}),
        Stage(StageType.STORE, "fetch", {}),
        Stage(StageType.DMA_WRITE, "dma", {}),
    )
    plan = TilePlan(tile_id=0, stages=stages)
    ctx = PipelineContext(id="p1", total_tiles=1)
    token = TileToken(
        tile_id=0, pipeline_ctx=ctx, plan=plan,
        stage_idx=0, params=stages[0].params,
    )

    visited = []
    while True:
        visited.append(token.current_stage.component)
        next_s = token.advance()
        if next_s is None:
            ctx.complete_tile()
            break

    assert visited == ["dma", "fetch", "gemm", "fetch", "dma"]
    assert ctx.completed_tiles == 1


# ── 5. Tiling plan generation ────────────────────────────────────────


def test_gemm_plan_tile_count():
    """generate_gemm_plan produces correct number of tiles."""
    from kernbench.components.builtin.tiling import generate_gemm_plan

    plan = generate_gemm_plan(
        M=64, K=128, N=64,
        tile_m=32, tile_k=64, tile_n=32,
        bytes_per_element=2,
        A_addr=0, B_addr=0x1000, C_addr=0x2000,
        pe_prefix="sip0.cube0.pe0",
    )
    # M_tiles=2, K_tiles=2, N_tiles=2 → 2*2*2 = 8 tiles
    assert len(plan.tiles) == 8
    assert plan.m_tiles == 2
    assert plan.k_tiles == 2
    assert plan.n_tiles == 2


def test_gemm_plan_stage_sequence():
    """Each GEMM tile has correct stage sequence."""
    from kernbench.components.builtin.tiling import generate_gemm_plan

    plan = generate_gemm_plan(
        M=32, K=64, N=32,
        tile_m=32, tile_k=64, tile_n=32,
        bytes_per_element=2,
        A_addr=0, B_addr=0x1000, C_addr=0x2000,
        pe_prefix="sip0.cube0.pe0",
    )
    # Single tile (1x1x1), last_k=True → includes DMA_WRITE
    assert len(plan.tiles) == 1
    tile = plan.tiles[0]
    stage_types = [s.stage_type for s in tile.stages]
    assert stage_types == [
        StageType.DMA_READ, StageType.DMA_READ,  # A and B
        StageType.FETCH, StageType.GEMM, StageType.STORE,
        StageType.DMA_WRITE,
    ]


def test_gemm_plan_intermediate_k_no_dma_write():
    """Intermediate K-tiles don't have DMA_WRITE or STORE stage.

    The C accumulator stays in RegFile across the K loop; STORE +
    DMA_WRITE only fire on the last K-tile per (m,n).
    """
    from kernbench.components.builtin.tiling import generate_gemm_plan

    plan = generate_gemm_plan(
        M=32, K=128, N=32,  # K_tiles=2
        tile_m=32, tile_k=64, tile_n=32,
        bytes_per_element=2,
        A_addr=0, B_addr=0x1000, C_addr=0x2000,
        pe_prefix="sip0.cube0.pe0",
    )
    assert len(plan.tiles) == 2

    # First tile (k=0): no STORE, no DMA_WRITE — accumulator stays in RegFile
    t0_types = [s.stage_type for s in plan.tiles[0].stages]
    assert StageType.STORE not in t0_types
    assert StageType.DMA_WRITE not in t0_types

    # Last tile (k=1, last_k=True): has both STORE and DMA_WRITE
    t1_types = [s.stage_type for s in plan.tiles[1].stages]
    assert StageType.STORE in t1_types
    assert StageType.DMA_WRITE in t1_types


def test_gemm_plan_pinned_operand_skips_dma_read():
    """When a_pinned=True, A's per-tile DMA_READ is omitted.

    Same for b_pinned. FETCH is unaffected — it still stages from TCM
    into RegFile.
    """
    from kernbench.components.builtin.tiling import generate_gemm_plan

    # Baseline: neither pinned — both A and B get DMA_READ per tile.
    base = generate_gemm_plan(
        M=32, K=128, N=32,  # K_tiles=2
        tile_m=32, tile_k=64, tile_n=32,
        bytes_per_element=2,
        A_addr=0, B_addr=0x1000, C_addr=0x2000,
        pe_prefix="sip0.cube0.pe0",
    )
    for tile in base.tiles:
        operands = [s.params.get("operand") for s in tile.stages
                    if s.stage_type == StageType.DMA_READ]
        assert operands == ["A", "B"], \
            f"baseline tile should DMA_READ A and B, got {operands}"

    # a_pinned: no A DMA_READ.
    plan_a = generate_gemm_plan(
        M=32, K=128, N=32,
        tile_m=32, tile_k=64, tile_n=32,
        bytes_per_element=2,
        A_addr=0, B_addr=0x1000, C_addr=0x2000,
        pe_prefix="sip0.cube0.pe0",
        a_pinned=True,
    )
    for tile in plan_a.tiles:
        operands = [s.params.get("operand") for s in tile.stages
                    if s.stage_type == StageType.DMA_READ]
        assert operands == ["B"], \
            f"a_pinned should leave only B DMA_READ, got {operands}"
        # FETCH must still exist
        assert any(s.stage_type == StageType.FETCH for s in tile.stages)

    # Both pinned: no DMA_READ at all.
    plan_both = generate_gemm_plan(
        M=32, K=128, N=32,
        tile_m=32, tile_k=64, tile_n=32,
        bytes_per_element=2,
        A_addr=0, B_addr=0x1000, C_addr=0x2000,
        pe_prefix="sip0.cube0.pe0",
        a_pinned=True, b_pinned=True,
    )
    for tile in plan_both.tiles:
        dma_reads = [s for s in tile.stages
                     if s.stage_type == StageType.DMA_READ]
        assert dma_reads == [], \
            f"both pinned should skip all DMA_READ, got {dma_reads}"


def test_math_plan_stage_sequence():
    """Math plan has READ→FETCH→MATH→STORE→WRITE sequence."""
    from kernbench.components.builtin.tiling import generate_math_plan

    plan = generate_math_plan(
        M=32, N=32,
        tile_m=32, tile_n=32,
        bytes_per_element=2,
        math_op="exp",
        src_addr=0, dst_addr=0x1000,
        pe_prefix="sip0.cube0.pe0",
    )
    assert len(plan.tiles) == 1
    stage_types = [s.stage_type for s in plan.tiles[0].stages]
    assert stage_types == [
        StageType.DMA_READ, StageType.FETCH, StageType.MATH,
        StageType.STORE, StageType.DMA_WRITE,
    ]


# ── 5. Async pipeline (SimPy simulation) ─────────────────────────────


def test_pipeline_overlap_within_command():
    """Tiles within same command overlap: tile1 DMA while tile0 in GEMM."""
    env = simpy.Environment()
    done_event = env.event()
    ctx = PipelineContext(id="p1", total_tiles=2, done_event=done_event)

    # Track when each tile enters each stage
    stage_times: dict[tuple[int, str], float] = {}

    def mock_component(env, inbox, stage_name, latency_ns, out_ports):
        while True:
            token = yield inbox.get()
            stage_times[(token.tile_id, stage_name)] = env.now
            yield env.timeout(latency_ns)
            next_s = token.advance()
            if next_s is not None:
                yield out_ports[next_s.component].put(token)
            else:
                token.pipeline_ctx.complete_tile()

    dma_q = simpy.Store(env)
    gemm_q = simpy.Store(env)

    out_ports = {"dma": dma_q, "gemm": gemm_q}
    env.process(mock_component(env, dma_q, "dma", 10.0, out_ports))
    env.process(mock_component(env, gemm_q, "gemm", 20.0, out_ports))

    # Create 2 tiles: DMA → GEMM
    for i in range(2):
        stages = (
            Stage(StageType.DMA_READ, "dma", {}),
            Stage(StageType.GEMM, "gemm", {}),
        )
        plan = TilePlan(tile_id=i, stages=stages)
        token = TileToken(
            tile_id=i, pipeline_ctx=ctx, plan=plan,
            stage_idx=0, params={},
        )
        dma_q.put(token)

    env.run()
    assert done_event.triggered

    # tile0 DMA starts at 0, finishes at 10
    # tile1 DMA starts at 10, finishes at 20
    # tile0 GEMM starts at 10, finishes at 30
    # tile1 GEMM starts at 20 (wait for DMA) but actually at 30 (gemm queue)
    assert stage_times[(0, "dma")] == 0.0
    assert stage_times[(0, "gemm")] == 10.0
    assert stage_times[(1, "dma")] == 10.0
    # tile1 gemm starts when tile0 gemm finishes (serialized at gemm queue)
    assert stage_times[(1, "gemm")] == 30.0


# ── 6. Option B: pe_dma record_start fires post channel-acquire ────────


def test_pe_dma_record_start_after_channel_acquire():
    """Three back-to-back DMA_READs at the same PE_DMA.

    With ``_DEFER_RECORD_START = True`` on PeDmaComponent, each op's
    ``t_start`` is captured right after ``yield req`` succeeds (Option B:
    record_start fires post-channel-acquire, not at command-entry).
    The DMA channel is held only for the **issue path** (ADR-0014 D4),
    not for the full HBM round-trip — HBM-level serialization is the
    HBM_CTRL's responsibility (per-PC `available_at`).

    Under that model, three back-to-back DMAs at t=0 give:

      - ``t_start`` values monotonic and tightly clustered (channel
        release is fast → each subsequent op's ``t_start`` is at or
        immediately after the previous one's, NOT after the previous
        one's HBM round-trip completes).
      - ``t_end`` values reflect HBM_CTRL serialization (later ops can
        have longer duration as they queue at the per-PC).

    Counter-example (the bug this rules out): if ``record_start`` fired
    on command entry, all three ops would share ``t_start == 0``.
    """
    from pathlib import Path

    from kernbench.common.pe_commands import DmaReadCmd, PeInternalTxn, TensorHandle
    from kernbench.policy.address.phyaddr import PhysAddr
    from kernbench.sim_engine.engine import GraphEngine
    from kernbench.topology.builder import load_topology

    TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"

    def _hbm_pa() -> int:
        slice_bytes = 48 * (1 << 30) // 8
        pa = PhysAddr.pe_hbm_addr(
            sip_id=0, die_id=0, pe_id=0,
            pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes,
        )
        return pa.encode()

    # enable_data=True wires the OpLogger into every component.
    engine = GraphEngine(load_topology(TOPOLOGY_PATH), enable_data=True)
    pe_dma_id = "sip0.cube0.pe0.pe_dma"
    pe_dma = engine._components[pe_dma_id]
    env = engine._env

    # Three back-to-back DMA_READ commands fed straight into pe_dma's inbox
    # at t=0 so they all race for the cap=1 channel.
    handles = [
        TensorHandle(id=f"r{i}", addr=0x1000 + i * 0x1000,
                     shape=(64, 32), dtype="f16", nbytes=4096)
        for i in range(3)
    ]
    cmds = [
        DmaReadCmd(handle=h, src_addr=_hbm_pa(), nbytes=4096)
        for h in handles
    ]
    txns = [PeInternalTxn(command=c, done=env.event()) for c in cmds]

    def submit_all():
        for txn in txns:
            yield pe_dma._inbox.put(txn)

    env.process(submit_all())
    env.run()

    # Pull the three dma_read records out of the op log in order
    dma_records = [
        r for r in engine.op_log
        if r.op_name == "dma_read" and r.component_id == pe_dma_id
    ]
    assert len(dma_records) == 3, (
        f"expected 3 dma_read records, got {len(dma_records)}: {dma_records}"
    )

    durations = [r.t_end - r.t_start for r in dma_records]
    base = durations[0]
    assert base > 0, f"first dma duration must be positive, got {base}"

    # t_start values must be monotonic and tightly clustered. The window
    # over all three t_starts is bounded by the per-op channel-issue cost,
    # NOT by the HBM round-trip — i.e. record_start fires post channel
    # acquire AND channel release is fast.
    t_starts = [r.t_start for r in dma_records]
    for i in range(1, len(t_starts)):
        assert t_starts[i] >= t_starts[i - 1], (
            f"op {i} t_start={t_starts[i]} regressed before "
            f"op {i-1} t_start={t_starts[i-1]} — record order broken"
        )
    t_start_window = t_starts[-1] - t_starts[0]
    assert t_start_window < base, (
        f"t_start window {t_start_window:.1f}ns across 3 back-to-back "
        f"ops is >= a single op's duration {base:.1f}ns — PE_DMA "
        f"channel is being held beyond the issue path (ADR-0014 D4)."
    )


# ── 7. Back-to-back DMA pipelining (head amortization) ───────────────


def test_pe_dma_back_to_back_pipelining():
    """Back-to-back DMA_READs should pipeline so that per-tile head latency
    is paid once, not N times.

    Setup: 8 back-to-back 4 KiB DMA_READ commands to the same PE_DMA's
    inbox at t=0. HBM_CTRL serializes via its per-PC availability
    timestamps; PE_DMA should release its channel after issuing the
    request, NOT hold it through the full HBM round-trip.

    Expected (post-fix model):
      - First op pays full round-trip: head + transmission.
      - Subsequent ops pay only transmission time (head amortized): each
        op's t_end advances by ~the per-request HBM service time, NOT by
        the full round-trip.
      - Total wall (last t_end - first t_start) is close to the
        transmission-bound limit: ``head_once + N * service``, not
        ``N * (head + service)``.

    Anti-counter-example: under the current PE_DMA round-trip-hold model
    each op N+1's t_start equals op N's t_end, so the wall is exactly
    ``N * (head + service)``. That is what this test rules out.

    Concrete numbers (32 KiB probe = 141 ns; 4 KiB single-request inferred
    at ~29 ns): the current model gives wall ~232 ns for 8 ops; the
    fixed model should give wall ≲ 160 ns (head ~13 + 8 * ~16 ns).
    """
    from pathlib import Path

    from kernbench.common.pe_commands import DmaReadCmd, PeInternalTxn, TensorHandle
    from kernbench.policy.address.phyaddr import PhysAddr
    from kernbench.sim_engine.engine import GraphEngine
    from kernbench.topology.builder import load_topology

    TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"

    def _hbm_pa(offset: int) -> int:
        slice_bytes = 48 * (1 << 30) // 8
        pa = PhysAddr.pe_hbm_addr(
            sip_id=0, die_id=0, pe_id=0,
            pe_local_hbm_offset=offset, slice_size_bytes=slice_bytes,
        )
        return pa.encode()

    engine = GraphEngine(load_topology(TOPOLOGY_PATH), enable_data=True)
    pe_dma_id = "sip0.cube0.pe0.pe_dma"
    pe_dma = engine._components[pe_dma_id]
    env = engine._env

    N_OPS = 8
    BYTES_PER_OP = 4096

    handles = [
        TensorHandle(id=f"r{i}", addr=0x1000 + i * 0x1000,
                     shape=(32, 64), dtype="f16", nbytes=BYTES_PER_OP)
        for i in range(N_OPS)
    ]
    # Spread across distinct cache lines so HBM_CTRL hits different PCs
    # and we measure the BW-only floor, not per-PC bank conflict.
    cmds = [
        DmaReadCmd(handle=h, src_addr=_hbm_pa(0x1000 * (1 + i)),
                   nbytes=BYTES_PER_OP)
        for i, h in enumerate(handles)
    ]
    txns = [PeInternalTxn(command=c, done=env.event()) for c in cmds]

    def submit_all():
        for txn in txns:
            yield pe_dma._inbox.put(txn)

    env.process(submit_all())
    env.run()

    dma_records = [
        r for r in engine.op_log
        if r.op_name == "dma_read" and r.component_id == pe_dma_id
    ]
    assert len(dma_records) == N_OPS, (
        f"expected {N_OPS} dma_read records, got {len(dma_records)}"
    )
    dma_records.sort(key=lambda r: r.t_start)

    single_op_ns = dma_records[0].t_end - dma_records[0].t_start
    wall = dma_records[-1].t_end - dma_records[0].t_start

    # Under the round-trip-hold model: wall == N_OPS * single_op_ns.
    # Under the issue-only-hold model: wall ≈ single_op_ns + (N-1)*service
    # where service ≈ BYTES_PER_OP / hbm_bw_per_pe ≈ 16 ns.
    # We assert wall < 0.75 * (N_OPS * single_op_ns) — i.e. amortization
    # of head must save at least 25% vs strict serialization.
    strict_serialized = N_OPS * single_op_ns
    assert wall < 0.75 * strict_serialized, (
        f"wall={wall:.1f}ns vs strict-serialized={strict_serialized:.1f}ns "
        f"(N_OPS={N_OPS}, single_op={single_op_ns:.1f}ns). "
        f"PE_DMA is still holding channel through full round-trip — "
        f"head latency is not amortized across back-to-back DMAs."
    )