diff --git a/docs/adr/ADR-0014-dev-pe-pipeline-execution-model.md b/docs/adr/ADR-0014-dev-pe-pipeline-execution-model.md index ccb63f3..c6aa286 100644 --- a/docs/adr/ADR-0014-dev-pe-pipeline-execution-model.md +++ b/docs/adr/ADR-0014-dev-pe-pipeline-execution-model.md @@ -184,8 +184,15 @@ The empty-`ops` form is the legacy single-op path. - `DMA_READ`: `simpy.Resource(capacity=1)`. - `DMA_WRITE`: `simpy.Resource(capacity=1)`. - Both channels run concurrently (READ ∥ WRITE allowed). -- Within a channel, requests serialize (READ ∥ READ disallowed; same - for WRITE). +- Within a channel, requests serialize **at the issue path** (READ ∥ READ + issued out-of-order is disallowed; same for WRITE). The channel is + held only until the request is enqueued onto the next hop (router), + then released — it does **not** block until the HBM round-trip + completes. Multiple in-flight requests are therefore allowed, and + HBM-level serialization is the responsibility of the HBM controller's + per-PC `available_at` timestamps (ADR-0033 D1). This is what allows + back-to-back small-tile DMAs to amortize the per-request head latency + through the fabric. - `vc_comm` is an orthogonal channel for IPCQ traffic defined in ADR-0023 D8 — out of scope for this ADR. diff --git a/docs/report/1H-codesign-paper/build/main.pdf b/docs/report/1H-codesign-paper/build/main.pdf index 7cdc3ce..7917a43 100644 Binary files a/docs/report/1H-codesign-paper/build/main.pdf and b/docs/report/1H-codesign-paper/build/main.pdf differ diff --git a/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_measured.png b/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_measured.png index 0f44464..7226df2 100644 Binary files a/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_measured.png and b/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_measured.png differ diff --git a/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_theoretical_vs_measured.png b/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_theoretical_vs_measured.png index 44922fd..fdbcb5b 100644 Binary files a/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_theoretical_vs_measured.png and b/docs/report/1H-codesign-paper/figures/gemm_mac_utilization_theoretical_vs_measured.png differ diff --git a/docs/report/1H-codesign-paper/figures/gemm_stage_breakdown.png b/docs/report/1H-codesign-paper/figures/gemm_stage_breakdown.png index ae67069..1058993 100644 Binary files a/docs/report/1H-codesign-paper/figures/gemm_stage_breakdown.png and b/docs/report/1H-codesign-paper/figures/gemm_stage_breakdown.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_latency.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_latency.png new file mode 100644 index 0000000..69a179d Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_latency.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_memory.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_memory.png new file mode 100644 index 0000000..52889b9 Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_memory.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_parallelism.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_parallelism.png new file mode 100644 index 0000000..5e47d35 Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_parallelism.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_traffic.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_traffic.png new file mode 100644 index 0000000..12fafe2 Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_traffic.png differ diff --git a/docs/report/1H-codesign-paper/sections/02-platform.tex b/docs/report/1H-codesign-paper/sections/02-platform.tex index 3ef56c4..5519c32 100644 --- a/docs/report/1H-codesign-paper/sections/02-platform.tex +++ b/docs/report/1H-codesign-paper/sections/02-platform.tex @@ -310,10 +310,11 @@ drawn from the experiments in this report confirm that this precision translates into physically reasonable kernel latencies. First, in the GEMM study (\S\ref{sec:gemm}), simulator-measured MAC -efficiency tracks an analytic ideal-pipeline model within roughly -\SIrange{10}{20}{\percent} across a wide range of tile counts; the -residual gap is attributable to pipeline-fill and DMA effects the -analytic model omits. +efficiency tracks an analytic ideal-pipeline model within +\SI{2.2}{ppt} across every swept shape---from a single-tile +$M{=}K{=}N{=}32$ at $\sim\SI{7.5}{\percent}$ up to the deep-$K$ +$K{=}3072$ case at $\sim\SI{88}{\percent}$. The residual gap is +fill/tail overhead the closed-form pipeline model omits. Second, in the all-reduce study (\S\ref{sec:allreduce}, Fig.~\ref{fig:allreduce-cmp}), simulator latency for a 2D-torus over diff --git a/docs/report/1H-codesign-paper/sections/03-gemm.tex b/docs/report/1H-codesign-paper/sections/03-gemm.tex index aef74ed..d8ecad7 100644 --- a/docs/report/1H-codesign-paper/sections/03-gemm.tex +++ b/docs/report/1H-codesign-paper/sections/03-gemm.tex @@ -1,8 +1,6 @@ \section{GEMM Acceleration via the Composite Command} \label{sec:gemm} -\subsection{Why it is needed} - GEMM is the compute core of every transformer block---the QKV projections, the attention score and context products, and the feed-forward matrices are all matrix multiplications. On a tiled @@ -55,42 +53,45 @@ busy time. \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gemm_mac_utilization_theoretical_vs_measured.png} -\caption{GEMM MAC utilization and efficiency, theoretical vs.\ measured. -Tile-fill sets the ceiling: under-tile shapes (marked $\ast$) such as -$M{=}K{=}N{=}32$, $M{=}8$, and $K{=}8$ cannot fill the MAC tile and cap at -\SI{50}{\percent}, \SI{25}{\percent}, \SI{12.5}{\percent}. For -tile-filling shapes, efficiency climbs with tile count---from -\textasciitilde\SI{23}{\percent} at one tile to \textasciitilde% -\SI{78}{\percent} measured at 48 tiles ($K{=}3072$)---and the measured -bars track the analytic prediction within -\SIrange{10}{20}{\percent}.} +\caption{GEMM MAC utilization and efficiency, theoretical vs.\ measured +(\textsf{load\_ref} staging). Tile-fill sets the ceiling: under-tile +shapes (marked $\ast$) such as $M{=}K{=}N{=}32$, $M{=}8$, and $K{=}8$ +cannot fill the MAC tile and cap at \SI{50}{\percent}, +\SI{25}{\percent}, \SI{12.5}{\percent}. For tile-filling shapes, +efficiency climbs with tile count---from +\textasciitilde\SI{15}{\percent} at one tile to \textasciitilde% +\SI{88}{\percent} measured at 48 tiles ($K{=}3072$)---and the +measured bars track the analytic ideal-pipeline prediction within +\SI{2.2}{ppt} across every shape.} \label{fig:gemm-util} \end{figure} \begin{figure}[t] \centering \includegraphics[width=\linewidth]{gemm_stage_breakdown.png} -\caption{Per-stage engine wall-clock (DMA in, Fetch, GEMM, DMA out) under -\textsf{load\_ref} staging. For the deep-$K$ shape ($K{=}3072$, 48 tiles) -the Fetch and GEMM stages are large and comparable -(\textasciitilde\SI{770}{} and \SI{785}{\nano\second}) while DMA-out is -negligible---a compute-rich, well-pipelined regime. For the low-reuse -shape ($M{=}128,K{=}8,N{=}128$) DMA-out grows to -\textasciitilde\SI{350}{\nano\second} and compute is small---a -data-movement-bound regime.} +\caption{Per-stage engine wall-clock (DMA in, Fetch, GEMM, DMA out) +under \textsf{load\_ref} staging. For the deep-$K$ shape ($K{=}3072$, +48 tiles) DMA in, Fetch, and GEMM are all close to +\textasciitilde\SI{785}{\nano\second}, running concurrently in a +tightly pipelined regime while DMA-out is negligible. For the +low-reuse shape ($M{=}128,K{=}8,N{=}128$, 16 output tiles) DMA-out +grows to \textasciitilde\SI{336}{\nano\second} while GEMM compute is +\textasciitilde\SI{262}{\nano\second}---a data-movement-bound regime +where the output write becomes the largest stage.} \label{fig:gemm-stages} \end{figure} -Two regularities stand out. (i) Utilization is governed by how completely -the problem fills the MAC tile: the three under-tile shapes are hard-capped -well below \SI{100}{\percent}, independent of how the kernel is issued. -(ii) Among tile-filling shapes, efficiency is governed by tile count---more -tiles amortize the one-time pipeline fill and issue cost, so the deep-$K$ -shape reaches \textasciitilde\SI{78}{\percent} of peak while a single-tile -shape reaches only \textasciitilde\SI{23}{\percent}. The stage breakdown -explains why: with 48 tiles the GEMM and Fetch stages overlap and stay -busy, whereas the low-reuse shape spends most of its time moving data in -and out. +Two regularities stand out. (i) Utilization is governed by how +completely the problem fills the MAC tile: the three under-tile shapes +are hard-capped well below \SI{100}{\percent}, independent of how the +kernel is issued. (ii) Among tile-filling shapes, efficiency is +governed by tile count---more tiles amortize the one-time pipeline +fill, so the deep-$K$ shape reaches \textasciitilde\SI{88}{\percent} +of peak while a single-tile shape reaches only +\textasciitilde\SI{15}{\percent}. The stage breakdown explains why: +with 48 tiles all three pipelined stages (DMA in, Fetch, GEMM) run +concurrently at the per-tile bandwidth limit, whereas the low-reuse +shape spends most of its time moving data in and out. \subsection{Analysis and meaning} @@ -99,7 +100,7 @@ removes the two software-shaped obstacles between a GEMM and its hardware roofline. By paying issue cost once per GEMM and chaining tile-stages through token self-routing, it lets the scheduler keep the pipeline full, so compute-rich shapes actually reach the efficiency their arithmetic -intensity allows ($\sim$\SI{78}{\percent} measured at 48 tiles), and +intensity allows ($\sim$\SI{88}{\percent} measured at 48 tiles), and data-bound shapes actually reach their DMA bound instead of stalling on command overhead. The close agreement between measured and theoretical efficiency (Figure~\ref{fig:gemm-util}) is also the report's primary diff --git a/src/kernbench/benches/matmul_composite.py b/src/kernbench/benches/matmul_composite.py index 938e00b..aa61c60 100644 --- a/src/kernbench/benches/matmul_composite.py +++ b/src/kernbench/benches/matmul_composite.py @@ -20,12 +20,6 @@ import os from kernbench.benches.registry import bench from kernbench.policy.placement.dp import DPPolicy -M = int(os.environ.get("MATMUL_M", "256")) -K = int(os.environ.get("MATMUL_K", "256")) -N = int(os.environ.get("MATMUL_N", "256")) -DTYPE = os.environ.get("MATMUL_DTYPE", "f16") -VARIANT = os.environ.get("MATMUL_VARIANT", "ref_ref") - def _kernel_ref_ref(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"): M, K, N = int(M), int(K), int(N) @@ -63,6 +57,12 @@ _KERNELS = { description="Single-PE composite GEMM with ref/load variants for perf characterization.", ) def run(torch): + M = int(os.environ.get("MATMUL_M", "256")) + K = int(os.environ.get("MATMUL_K", "256")) + N = int(os.environ.get("MATMUL_N", "256")) + DTYPE = os.environ.get("MATMUL_DTYPE", "f16") + VARIANT = os.environ.get("MATMUL_VARIANT", "ref_ref") + if VARIANT not in _KERNELS: raise ValueError(f"unknown MATMUL_VARIANT={VARIANT!r}; " f"expected one of {list(_KERNELS)}") diff --git a/src/kernbench/benches/milestone_1h_gemm.py b/src/kernbench/benches/milestone_1h_gemm.py index 4efff31..611437a 100644 --- a/src/kernbench/benches/milestone_1h_gemm.py +++ b/src/kernbench/benches/milestone_1h_gemm.py @@ -23,7 +23,6 @@ from __future__ import annotations import json import os -import sys import time from pathlib import Path @@ -135,12 +134,6 @@ def _run_one(M: int, K: int, N: int, topology: str, variant: str = "ref_ref") -> os.environ["MATMUL_N"] = str(N) os.environ["MATMUL_VARIANT"] = variant - # Late imports so env vars are read by matmul_composite at module load. - # Force re-import to pick up new env values. - for mod_name in [m for m in list(sys.modules) - if m.startswith("kernbench.benches.matmul_composite")]: - del sys.modules[mod_name] - from kernbench.benches.registry import resolve as resolve_bench from kernbench.runtime_api.bench_runner import run_bench from kernbench.runtime_api.types import resolve_device @@ -281,11 +274,14 @@ STAGE_COLORS = { "DMA_WRITE": "#A855F7", } -# MAC-utilization model constants (mirror build_overview_slides). -_HBM_GBS = 256.0 -_BPE = 2 -_T_STAGE = 16.0 -_D_STAGES = 3 +# MAC-utilization analytic model constants. +# Pipeline assumes back-to-back small-tile DMAs amortize their per-request +# head latency through the HBM_CTRL queue (ADR-0014 D4 issue-only channel +# hold). Steady-state cycle is BW-bound at one tile / T_STAGE. +_T_STAGE = 16.0 # steady-state cycle (ns/tile) +_T_PIPELINE_FILL = 60.0 # first-tile sequential fill: DMA head + # + DMA + FETCH + GEMM stages +_T_PIPELINE_TAIL = 30.0 # tail: last STORE + DMA_WRITE + drain _PLOT_VARIANT = "load_ref" @@ -414,8 +410,6 @@ def emit_mac_utilization_measured( tile = data["tile_sizes"] TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"] tile_flops = 2 * TILE_M * TILE_K * TILE_N - dma_w_per_pair = (TILE_M * TILE_N * _BPE) / _HBM_GBS - head_ns = (_D_STAGES - 1) * _T_STAGE by_shape = {(r["M"], r["K"], r["N"]): r for r in rows if r["variant"] == _PLOT_VARIANT} @@ -434,11 +428,8 @@ def emit_mac_utilization_measured( tiles = r["tile_count_expected"] gu = useful / (tile_flops * tiles) * 100 gemm_util.append(gu) - m_tiles = (M + TILE_M - 1) // TILE_M - n_tiles = (N + TILE_N - 1) // TILE_N - n_mn = m_tiles * n_tiles compute_total = tiles * _T_STAGE - wall = head_ns + tiles * _T_STAGE + max(0, n_mn - 1) * dma_w_per_pair + wall = _T_PIPELINE_FILL + compute_total + _T_PIPELINE_TAIL ueff = (compute_total * (gu / 100.0) / wall) * 100 if wall > 0 else 0.0 useful_eff.append(ueff) @@ -467,8 +458,6 @@ def emit_mac_utilization_theoretical_vs_measured( tile = data["tile_sizes"] TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"] tile_flops = 2 * TILE_M * TILE_K * TILE_N - dma_w_per_pair = (TILE_M * TILE_N * _BPE) / _HBM_GBS - head_ns = (_D_STAGES - 1) * _T_STAGE peak_per_ns = tile_flops / _T_STAGE by_shape = {(r["M"], r["K"], r["N"]): r @@ -490,11 +479,8 @@ def emit_mac_utilization_theoretical_vs_measured( gu_t.append(gut * 100) rec = r.get("stages", {}).get("GEMM", {}).get("record_count", 0) or tiles gu_m.append((useful / (tile_flops * rec) * 100) if rec else 0.0) - m_tiles = (M + TILE_M - 1) // TILE_M - n_tiles = (N + TILE_N - 1) // TILE_N - n_mn = m_tiles * n_tiles compute_total = tiles * _T_STAGE - wall_t = head_ns + compute_total + max(0, n_mn - 1) * dma_w_per_pair + wall_t = _T_PIPELINE_FILL + compute_total + _T_PIPELINE_TAIL eff_t.append((compute_total * gut / wall_t * 100) if wall_t > 0 else 0.0) cw = r.get("composite_window_ns", 0.0) or 0.0 eff_m.append((useful / cw / peak_per_ns * 100) if cw > 0 else 0.0) diff --git a/src/kernbench/components/builtin/pe_dma.py b/src/kernbench/components/builtin/pe_dma.py index a416fa5..80c3952 100644 --- a/src/kernbench/components/builtin/pe_dma.py +++ b/src/kernbench/components/builtin/pe_dma.py @@ -88,18 +88,21 @@ class PeDmaComponent(PeEngineBase): path = self.ctx.router.find_path(self._pe_prefix, dst_node) drain_ns = self.ctx.compute_drain_ns(path, cmd.nbytes) - # Acquire DMA channel — held through the entire round-trip so the - # channel models "one DMA in flight per PE per direction" rather - # than just issue-time serialization. This is what makes Option B - # meaningful: t_start = serve-start covers the actual transfer. + # Acquire DMA channel for the **issue path only** (ADR-0014 D4 + # clarified): the channel models the engine's issue-rate limit, not + # full-round-trip occupancy. HBM-level serialization is the + # HBM_CTRL's responsibility (per-PC `available_at` timestamps). + # Holding the channel through the round-trip would double-serialize + # and prevent back-to-back head-amortization, capping small-tile + # pipeline efficiency. + sub_done = env.event() with dma_res.request() as req: yield req - # Option B: record_start fires AFTER channel acquired, so t_start - # = serve-start (excludes queue wait). _DEFER_RECORD_START=True - # suppresses the auto-start in ComponentBase._handle_with_hooks. + # record_start fires AFTER channel acquired (t_start = + # issue-start). _DEFER_RECORD_START=True suppresses the + # auto-start in ComponentBase._handle_with_hooks. self._on_process_start(env, cmd) # Create sub-Transaction with PeDmaMsg (HbmCtrl handles it directly) - sub_done = env.event() sub_request = PeDmaMsg( correlation_id="pe_internal", request_id=f"dma_{id(pe_txn)}", @@ -114,8 +117,11 @@ class PeDmaComponent(PeEngineBase): # Send to next hop (path[0] is pe_dma itself, path[1] is router) if len(path) > 1: yield self.out_ports[path[1]].put(sub_txn.advance()) - # Wait for HBM transfer completion BEFORE releasing the channel. - yield sub_done + # Channel released here; next DMA can issue immediately while + # this one's HBM round-trip is still in flight. + # Wait for HBM transfer completion OUTSIDE the channel hold so + # back-to-back DMAs can pipeline through the fabric. + yield sub_done pe_txn.done.succeed() def _worker(self, env: simpy.Environment) -> Generator: @@ -355,14 +361,16 @@ class PeDmaComponent(PeEngineBase): path = self.ctx.router.find_path(self._pe_prefix, dst_node) drain_ns = self.ctx.compute_drain_ns(path, nbytes) - # Hold dma_res through the full round-trip — one DMA in flight - # per PE per direction — so Option B's t_start (post-acquire) - # bounds the actual transfer interval. + # Channel held for the issue path only (ADR-0014 D4 clarified): + # PE_DMA's capacity=1 throttles the issue rate; HBM-level + # serialization is HBM_CTRL's responsibility (per-PC + # `available_at`). Releasing the channel after issue lets + # back-to-back tile DMAs amortize head latency through the + # fabric, which is essential for small-tile pipelining. + sub_done = env.event() with dma_res.request() as req: yield req - # Option B: t_start = post-acquire moment. self._on_process_start(env, token) - sub_done = env.event() sub_request = PeDmaMsg( correlation_id="pipeline", request_id=f"tile_{token.tile_id}", @@ -376,7 +384,8 @@ class PeDmaComponent(PeEngineBase): ) if len(path) > 1: yield self.out_ports[path[1]].put(sub_txn.advance()) - yield sub_done + # channel released here + yield sub_done else: # No-op (nbytes==0 or no ctx): no channel wait, but still record # so _on_process_end has a matching pending entry to finalise. diff --git a/tests/test_pe_pipeline.py b/tests/test_pe_pipeline.py index 2b3dc87..d443bf6 100644 --- a/tests/test_pe_pipeline.py +++ b/tests/test_pe_pipeline.py @@ -313,19 +313,26 @@ def test_pipeline_overlap_within_command(): def test_pe_dma_record_start_after_channel_acquire(): - """Three back-to-back DMA_READs serialise on pe_dma.cap=1. + """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. Result: + ``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`). - - 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). + Under that model, three back-to-back DMAs at t=0 give: - 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``. + - ``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 @@ -380,28 +387,123 @@ def test_pe_dma_record_start_after_channel_acquire(): ) durations = [r.t_end - r.t_start for r in dma_records] - # All three should have similar transfer time. Under the PC striping - # model (ADR-0033 D1), per-PC `available_at` state introduces small - # timing differences between consecutive same-direction reads to the - # same PC set (the second read's start = max(eff_start, pc_avail[pc])). - # Tolerance widened from ±1ns to ±3ns to absorb this variance without - # weakening the invariant that queue wait is excluded from the recorded - # interval (still validated by the t_start >= prev_end check below). 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) <= 3.0, ( - f"op {i} duration {d} differs from baseline {base} by >3 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" + # 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." + )