Honest measured pipeline efficiency: two timing fixes

Two related issues caused measured pipeline efficiency to look
worse than the simulator's actual behavior:

1. DMA timing recorded too early. The op-log start timestamp
   for a DMA op fired when the request entered the queue, and
   the DMA channel was released as soon as the request was
   issued. Back-to-back DMAs therefore appeared to grab the
   channel simultaneously, with per-op duration drifting
   upward as queue depth grew - an artifact, not real cost.

   Fix: defer the start timestamp until after the channel is
   acquired, and hold the channel through the full HBM
   round-trip until the response returns. Per-op duration is
   now constant and equal to the actual transfer interval;
   serialization is visible as queue wait, not as inflated
   service time.

2. Sweep timing window folded in pre-composite work. The PE
   timing window spanned every PE engine record, which
   included the upfront pinned-operand DMA issued before the
   composite GEMM begins. For large-K shapes that one-shot
   load can be nearly half of the window, conflating
   operand-staging cost with composite-pipeline behavior.

   Fix: add a second window scoped to the composite pipeline
   by filtering op_log records to those tagged with a
   tile-pipeline stage; the legacy operand-load path is
   untagged and naturally excluded. For 32x3072x32 load_ref
   the window drops from 1765ns to 992ns and measured eff
   lines up with the steady-state DMA-bound stage limit
   instead of being penalized for the one-time load.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-14 14:19:17 -07:00
parent 83ea97b05f
commit f6d262e359
7 changed files with 543 additions and 263 deletions
+92
View File
@@ -307,3 +307,95 @@ def test_pipeline_overlap_within_command():
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 serialise on pe_dma.cap=1.
With ``_DEFER_RECORD_START = True`` on PeDmaComponent, each op's
``t_start`` is captured right after ``yield req`` succeeds. Result:
- op N's ``(t_end - t_start)`` is the *actual transfer time* — same
across all three ops (no queueing inflation).
- op N+1's ``t_start`` >= op N's ``t_end - epsilon`` (waited for the
previous holder to release the channel before being recorded).
Counter-example (the bug this fix addresses): if ``record_start`` fired
on command entry, all three ops would share ``t_start == 0`` and the
second/third would show inflated ``t_end - t_start``.
"""
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]
# All three should have the same actual transfer time within ±1 ns.
base = durations[0]
assert base > 0, f"first dma duration must be positive, got {base}"
for i, d in enumerate(durations):
assert abs(d - base) <= 1.0, (
f"op {i} duration {d} differs from baseline {base} by >1 ns "
f"— record_start may still be including queue wait"
)
# Each subsequent op's t_start must be at or after the previous op's
# t_end (modulo a few ns of scheduler overhead) — i.e. the wait is
# *excluded* from the recorded interval, not folded into it.
for i in range(1, len(dma_records)):
prev_end = dma_records[i - 1].t_end
cur_start = dma_records[i].t_start
assert cur_start >= prev_end - 1.0, (
f"op {i} t_start={cur_start} began before op {i-1} t_end={prev_end} "
f"— channel was not actually held, fix is incorrect"
)