gemm(perf): fix back-to-back DMA pipelining + tighten analytic model
Three coupled fixes that recover small-tile GEMM pipeline efficiency from 53% to 88% (32x3072x32 load_ref, composite_window basis). 1. PE_DMA channel-hold (ADR-0014 D4 clarified): both the _handle_with_hooks (PeInternalTxn) and _pipeline_process (TileToken) paths used to hold the cap=1 DMA channel through the full HBM round-trip, which double-serialized with the HBM_CTRL's own per-PC `available_at` model and prevented back-to-back tile DMAs from amortizing their per-request head latency. Channel is now released after the request is enqueued onto the next hop; HBM serialization is HBM_CTRL's responsibility alone. Tests: new test_pe_dma_back_to_back_pipelining as the oracle (asserts wall < 75% of strict-serialized N x single_op). Existing test_pe_dma_record_start_after_channel_acquire rewritten to assert t_start clustering (channel released fast) instead of the old round-trip-hold invariant. test_pe_dma_same_channel_serializes still passes — HBM_CTRL preserves ordering. Probe regression: PE→local-HBM 32 KiB stays at 141 ns (single-request, unaffected). 2. milestone_1h_gemm bench: matmul_composite was reading MATMUL_M/K/N env vars at module load, so every sweep row replayed the cached 256³ result; values now read inside run(). Drops the stale sys.modules deletion hack. 3. Analytic ideal-pipeline model: dropped the (n_mn-1)·dma_w_per_pair penalty (over-pessimistic for under-tile shapes — it pushed measured > theoretical) and replaced the D_STAGES-derived head with empirical T_PIPELINE_FILL=60 ns / T_PIPELINE_TAIL=30 ns. Max analytic-vs-measured gap across all 7 swept shapes now 2.2 ppt (was 9-44 ppt under the old constants). Paper updates: - §3 (GEMM): 78%→88% measured at 48 tiles, 23%→15% at 1 tile, stage breakdown numbers refreshed (DMA in / Fetch / GEMM all ~785 ns at K=3072), analytic-vs-measured agreement tightened to "within 2.2 ppt". - §2.4 (Accuracy): GEMM tracking claim refreshed accordingly. - §5 (GQA): restore long-ctx 4-cases figures into figures/ (they were dropped from bench output dir as derived artifacts in92b9221/e45626cbut §5 still cites them by name). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@@ -184,8 +184,15 @@ The empty-`ops` form is the legacy single-op path.
|
|||||||
- `DMA_READ`: `simpy.Resource(capacity=1)`.
|
- `DMA_READ`: `simpy.Resource(capacity=1)`.
|
||||||
- `DMA_WRITE`: `simpy.Resource(capacity=1)`.
|
- `DMA_WRITE`: `simpy.Resource(capacity=1)`.
|
||||||
- Both channels run concurrently (READ ∥ WRITE allowed).
|
- Both channels run concurrently (READ ∥ WRITE allowed).
|
||||||
- Within a channel, requests serialize (READ ∥ READ disallowed; same
|
- Within a channel, requests serialize **at the issue path** (READ ∥ READ
|
||||||
for WRITE).
|
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
|
- `vc_comm` is an orthogonal channel for IPCQ traffic defined in
|
||||||
ADR-0023 D8 — out of scope for this ADR.
|
ADR-0023 D8 — out of scope for this ADR.
|
||||||
|
|
||||||
|
|||||||
|
Before Width: | Height: | Size: 40 KiB After Width: | Height: | Size: 40 KiB |
|
Before Width: | Height: | Size: 46 KiB After Width: | Height: | Size: 46 KiB |
|
Before Width: | Height: | Size: 42 KiB After Width: | Height: | Size: 45 KiB |
|
After Width: | Height: | Size: 47 KiB |
|
After Width: | Height: | Size: 43 KiB |
|
After Width: | Height: | Size: 61 KiB |
|
After Width: | Height: | Size: 41 KiB |
@@ -310,10 +310,11 @@ drawn from the experiments in this report confirm that this precision
|
|||||||
translates into physically reasonable kernel latencies.
|
translates into physically reasonable kernel latencies.
|
||||||
|
|
||||||
First, in the GEMM study (\S\ref{sec:gemm}), simulator-measured MAC
|
First, in the GEMM study (\S\ref{sec:gemm}), simulator-measured MAC
|
||||||
efficiency tracks an analytic ideal-pipeline model within roughly
|
efficiency tracks an analytic ideal-pipeline model within
|
||||||
\SIrange{10}{20}{\percent} across a wide range of tile counts; the
|
\SI{2.2}{ppt} across every swept shape---from a single-tile
|
||||||
residual gap is attributable to pipeline-fill and DMA effects the
|
$M{=}K{=}N{=}32$ at $\sim\SI{7.5}{\percent}$ up to the deep-$K$
|
||||||
analytic model omits.
|
$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},
|
Second, in the all-reduce study (\S\ref{sec:allreduce},
|
||||||
Fig.~\ref{fig:allreduce-cmp}), simulator latency for a 2D-torus over
|
Fig.~\ref{fig:allreduce-cmp}), simulator latency for a 2D-torus over
|
||||||
|
|||||||
@@ -1,8 +1,6 @@
|
|||||||
\section{GEMM Acceleration via the Composite Command}
|
\section{GEMM Acceleration via the Composite Command}
|
||||||
\label{sec:gemm}
|
\label{sec:gemm}
|
||||||
|
|
||||||
\subsection{Why it is needed}
|
|
||||||
|
|
||||||
GEMM is the compute core of every transformer block---the QKV
|
GEMM is the compute core of every transformer block---the QKV
|
||||||
projections, the attention score and context products, and the
|
projections, the attention score and context products, and the
|
||||||
feed-forward matrices are all matrix multiplications. On a tiled
|
feed-forward matrices are all matrix multiplications. On a tiled
|
||||||
@@ -55,42 +53,45 @@ busy time.
|
|||||||
\begin{figure}[t]
|
\begin{figure}[t]
|
||||||
\centering
|
\centering
|
||||||
\includegraphics[width=\linewidth]{gemm_mac_utilization_theoretical_vs_measured.png}
|
\includegraphics[width=\linewidth]{gemm_mac_utilization_theoretical_vs_measured.png}
|
||||||
\caption{GEMM MAC utilization and efficiency, theoretical vs.\ measured.
|
\caption{GEMM MAC utilization and efficiency, theoretical vs.\ measured
|
||||||
Tile-fill sets the ceiling: under-tile shapes (marked $\ast$) such as
|
(\textsf{load\_ref} staging). Tile-fill sets the ceiling: under-tile
|
||||||
$M{=}K{=}N{=}32$, $M{=}8$, and $K{=}8$ cannot fill the MAC tile and cap at
|
shapes (marked $\ast$) such as $M{=}K{=}N{=}32$, $M{=}8$, and $K{=}8$
|
||||||
\SI{50}{\percent}, \SI{25}{\percent}, \SI{12.5}{\percent}. For
|
cannot fill the MAC tile and cap at \SI{50}{\percent},
|
||||||
tile-filling shapes, efficiency climbs with tile count---from
|
\SI{25}{\percent}, \SI{12.5}{\percent}. For tile-filling shapes,
|
||||||
\textasciitilde\SI{23}{\percent} at one tile to \textasciitilde%
|
efficiency climbs with tile count---from
|
||||||
\SI{78}{\percent} measured at 48 tiles ($K{=}3072$)---and the measured
|
\textasciitilde\SI{15}{\percent} at one tile to \textasciitilde%
|
||||||
bars track the analytic prediction within
|
\SI{88}{\percent} measured at 48 tiles ($K{=}3072$)---and the
|
||||||
\SIrange{10}{20}{\percent}.}
|
measured bars track the analytic ideal-pipeline prediction within
|
||||||
|
\SI{2.2}{ppt} across every shape.}
|
||||||
\label{fig:gemm-util}
|
\label{fig:gemm-util}
|
||||||
\end{figure}
|
\end{figure}
|
||||||
|
|
||||||
\begin{figure}[t]
|
\begin{figure}[t]
|
||||||
\centering
|
\centering
|
||||||
\includegraphics[width=\linewidth]{gemm_stage_breakdown.png}
|
\includegraphics[width=\linewidth]{gemm_stage_breakdown.png}
|
||||||
\caption{Per-stage engine wall-clock (DMA in, Fetch, GEMM, DMA out) under
|
\caption{Per-stage engine wall-clock (DMA in, Fetch, GEMM, DMA out)
|
||||||
\textsf{load\_ref} staging. For the deep-$K$ shape ($K{=}3072$, 48 tiles)
|
under \textsf{load\_ref} staging. For the deep-$K$ shape ($K{=}3072$,
|
||||||
the Fetch and GEMM stages are large and comparable
|
48 tiles) DMA in, Fetch, and GEMM are all close to
|
||||||
(\textasciitilde\SI{770}{} and \SI{785}{\nano\second}) while DMA-out is
|
\textasciitilde\SI{785}{\nano\second}, running concurrently in a
|
||||||
negligible---a compute-rich, well-pipelined regime. For the low-reuse
|
tightly pipelined regime while DMA-out is negligible. For the
|
||||||
shape ($M{=}128,K{=}8,N{=}128$) DMA-out grows to
|
low-reuse shape ($M{=}128,K{=}8,N{=}128$, 16 output tiles) DMA-out
|
||||||
\textasciitilde\SI{350}{\nano\second} and compute is small---a
|
grows to \textasciitilde\SI{336}{\nano\second} while GEMM compute is
|
||||||
data-movement-bound regime.}
|
\textasciitilde\SI{262}{\nano\second}---a data-movement-bound regime
|
||||||
|
where the output write becomes the largest stage.}
|
||||||
\label{fig:gemm-stages}
|
\label{fig:gemm-stages}
|
||||||
\end{figure}
|
\end{figure}
|
||||||
|
|
||||||
Two regularities stand out. (i) Utilization is governed by how completely
|
Two regularities stand out. (i) Utilization is governed by how
|
||||||
the problem fills the MAC tile: the three under-tile shapes are hard-capped
|
completely the problem fills the MAC tile: the three under-tile shapes
|
||||||
well below \SI{100}{\percent}, independent of how the kernel is issued.
|
are hard-capped well below \SI{100}{\percent}, independent of how the
|
||||||
(ii) Among tile-filling shapes, efficiency is governed by tile count---more
|
kernel is issued. (ii) Among tile-filling shapes, efficiency is
|
||||||
tiles amortize the one-time pipeline fill and issue cost, so the deep-$K$
|
governed by tile count---more tiles amortize the one-time pipeline
|
||||||
shape reaches \textasciitilde\SI{78}{\percent} of peak while a single-tile
|
fill, so the deep-$K$ shape reaches \textasciitilde\SI{88}{\percent}
|
||||||
shape reaches only \textasciitilde\SI{23}{\percent}. The stage breakdown
|
of peak while a single-tile shape reaches only
|
||||||
explains why: with 48 tiles the GEMM and Fetch stages overlap and stay
|
\textasciitilde\SI{15}{\percent}. The stage breakdown explains why:
|
||||||
busy, whereas the low-reuse shape spends most of its time moving data in
|
with 48 tiles all three pipelined stages (DMA in, Fetch, GEMM) run
|
||||||
and out.
|
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}
|
\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
|
roofline. By paying issue cost once per GEMM and chaining tile-stages
|
||||||
through token self-routing, it lets the scheduler keep the pipeline full,
|
through token self-routing, it lets the scheduler keep the pipeline full,
|
||||||
so compute-rich shapes actually reach the efficiency their arithmetic
|
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
|
data-bound shapes actually reach their DMA bound instead of stalling on
|
||||||
command overhead. The close agreement between measured and theoretical
|
command overhead. The close agreement between measured and theoretical
|
||||||
efficiency (Figure~\ref{fig:gemm-util}) is also the report's primary
|
efficiency (Figure~\ref{fig:gemm-util}) is also the report's primary
|
||||||
|
|||||||
@@ -20,12 +20,6 @@ import os
|
|||||||
from kernbench.benches.registry import bench
|
from kernbench.benches.registry import bench
|
||||||
from kernbench.policy.placement.dp import DPPolicy
|
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"):
|
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)
|
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.",
|
description="Single-PE composite GEMM with ref/load variants for perf characterization.",
|
||||||
)
|
)
|
||||||
def run(torch):
|
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:
|
if VARIANT not in _KERNELS:
|
||||||
raise ValueError(f"unknown MATMUL_VARIANT={VARIANT!r}; "
|
raise ValueError(f"unknown MATMUL_VARIANT={VARIANT!r}; "
|
||||||
f"expected one of {list(_KERNELS)}")
|
f"expected one of {list(_KERNELS)}")
|
||||||
|
|||||||
@@ -23,7 +23,6 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import sys
|
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
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_N"] = str(N)
|
||||||
os.environ["MATMUL_VARIANT"] = variant
|
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.benches.registry import resolve as resolve_bench
|
||||||
from kernbench.runtime_api.bench_runner import run_bench
|
from kernbench.runtime_api.bench_runner import run_bench
|
||||||
from kernbench.runtime_api.types import resolve_device
|
from kernbench.runtime_api.types import resolve_device
|
||||||
@@ -281,11 +274,14 @@ STAGE_COLORS = {
|
|||||||
"DMA_WRITE": "#A855F7",
|
"DMA_WRITE": "#A855F7",
|
||||||
}
|
}
|
||||||
|
|
||||||
# MAC-utilization model constants (mirror build_overview_slides).
|
# MAC-utilization analytic model constants.
|
||||||
_HBM_GBS = 256.0
|
# Pipeline assumes back-to-back small-tile DMAs amortize their per-request
|
||||||
_BPE = 2
|
# head latency through the HBM_CTRL queue (ADR-0014 D4 issue-only channel
|
||||||
_T_STAGE = 16.0
|
# hold). Steady-state cycle is BW-bound at one tile / T_STAGE.
|
||||||
_D_STAGES = 3
|
_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"
|
_PLOT_VARIANT = "load_ref"
|
||||||
|
|
||||||
@@ -414,8 +410,6 @@ def emit_mac_utilization_measured(
|
|||||||
tile = data["tile_sizes"]
|
tile = data["tile_sizes"]
|
||||||
TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"]
|
TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"]
|
||||||
tile_flops = 2 * 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
|
by_shape = {(r["M"], r["K"], r["N"]): r
|
||||||
for r in rows if r["variant"] == _PLOT_VARIANT}
|
for r in rows if r["variant"] == _PLOT_VARIANT}
|
||||||
@@ -434,11 +428,8 @@ def emit_mac_utilization_measured(
|
|||||||
tiles = r["tile_count_expected"]
|
tiles = r["tile_count_expected"]
|
||||||
gu = useful / (tile_flops * tiles) * 100
|
gu = useful / (tile_flops * tiles) * 100
|
||||||
gemm_util.append(gu)
|
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
|
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
|
ueff = (compute_total * (gu / 100.0) / wall) * 100 if wall > 0 else 0.0
|
||||||
useful_eff.append(ueff)
|
useful_eff.append(ueff)
|
||||||
|
|
||||||
@@ -467,8 +458,6 @@ def emit_mac_utilization_theoretical_vs_measured(
|
|||||||
tile = data["tile_sizes"]
|
tile = data["tile_sizes"]
|
||||||
TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"]
|
TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"]
|
||||||
tile_flops = 2 * 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
|
peak_per_ns = tile_flops / _T_STAGE
|
||||||
|
|
||||||
by_shape = {(r["M"], r["K"], r["N"]): r
|
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)
|
gu_t.append(gut * 100)
|
||||||
rec = r.get("stages", {}).get("GEMM", {}).get("record_count", 0) or tiles
|
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)
|
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
|
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)
|
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
|
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)
|
eff_m.append((useful / cw / peak_per_ns * 100) if cw > 0 else 0.0)
|
||||||
|
|||||||
@@ -88,18 +88,21 @@ class PeDmaComponent(PeEngineBase):
|
|||||||
path = self.ctx.router.find_path(self._pe_prefix, dst_node)
|
path = self.ctx.router.find_path(self._pe_prefix, dst_node)
|
||||||
drain_ns = self.ctx.compute_drain_ns(path, cmd.nbytes)
|
drain_ns = self.ctx.compute_drain_ns(path, cmd.nbytes)
|
||||||
|
|
||||||
# Acquire DMA channel — held through the entire round-trip so the
|
# Acquire DMA channel for the **issue path only** (ADR-0014 D4
|
||||||
# channel models "one DMA in flight per PE per direction" rather
|
# clarified): the channel models the engine's issue-rate limit, not
|
||||||
# than just issue-time serialization. This is what makes Option B
|
# full-round-trip occupancy. HBM-level serialization is the
|
||||||
# meaningful: t_start = serve-start covers the actual transfer.
|
# 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:
|
with dma_res.request() as req:
|
||||||
yield req
|
yield req
|
||||||
# Option B: record_start fires AFTER channel acquired, so t_start
|
# record_start fires AFTER channel acquired (t_start =
|
||||||
# = serve-start (excludes queue wait). _DEFER_RECORD_START=True
|
# issue-start). _DEFER_RECORD_START=True suppresses the
|
||||||
# suppresses the auto-start in ComponentBase._handle_with_hooks.
|
# auto-start in ComponentBase._handle_with_hooks.
|
||||||
self._on_process_start(env, cmd)
|
self._on_process_start(env, cmd)
|
||||||
# Create sub-Transaction with PeDmaMsg (HbmCtrl handles it directly)
|
# Create sub-Transaction with PeDmaMsg (HbmCtrl handles it directly)
|
||||||
sub_done = env.event()
|
|
||||||
sub_request = PeDmaMsg(
|
sub_request = PeDmaMsg(
|
||||||
correlation_id="pe_internal",
|
correlation_id="pe_internal",
|
||||||
request_id=f"dma_{id(pe_txn)}",
|
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)
|
# Send to next hop (path[0] is pe_dma itself, path[1] is router)
|
||||||
if len(path) > 1:
|
if len(path) > 1:
|
||||||
yield self.out_ports[path[1]].put(sub_txn.advance())
|
yield self.out_ports[path[1]].put(sub_txn.advance())
|
||||||
# Wait for HBM transfer completion BEFORE releasing the channel.
|
# Channel released here; next DMA can issue immediately while
|
||||||
yield sub_done
|
# 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()
|
pe_txn.done.succeed()
|
||||||
|
|
||||||
def _worker(self, env: simpy.Environment) -> Generator:
|
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)
|
path = self.ctx.router.find_path(self._pe_prefix, dst_node)
|
||||||
drain_ns = self.ctx.compute_drain_ns(path, nbytes)
|
drain_ns = self.ctx.compute_drain_ns(path, nbytes)
|
||||||
|
|
||||||
# Hold dma_res through the full round-trip — one DMA in flight
|
# Channel held for the issue path only (ADR-0014 D4 clarified):
|
||||||
# per PE per direction — so Option B's t_start (post-acquire)
|
# PE_DMA's capacity=1 throttles the issue rate; HBM-level
|
||||||
# bounds the actual transfer interval.
|
# 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:
|
with dma_res.request() as req:
|
||||||
yield req
|
yield req
|
||||||
# Option B: t_start = post-acquire moment.
|
|
||||||
self._on_process_start(env, token)
|
self._on_process_start(env, token)
|
||||||
sub_done = env.event()
|
|
||||||
sub_request = PeDmaMsg(
|
sub_request = PeDmaMsg(
|
||||||
correlation_id="pipeline",
|
correlation_id="pipeline",
|
||||||
request_id=f"tile_{token.tile_id}",
|
request_id=f"tile_{token.tile_id}",
|
||||||
@@ -376,7 +384,8 @@ class PeDmaComponent(PeEngineBase):
|
|||||||
)
|
)
|
||||||
if len(path) > 1:
|
if len(path) > 1:
|
||||||
yield self.out_ports[path[1]].put(sub_txn.advance())
|
yield self.out_ports[path[1]].put(sub_txn.advance())
|
||||||
yield sub_done
|
# channel released here
|
||||||
|
yield sub_done
|
||||||
else:
|
else:
|
||||||
# No-op (nbytes==0 or no ctx): no channel wait, but still record
|
# No-op (nbytes==0 or no ctx): no channel wait, but still record
|
||||||
# so _on_process_end has a matching pending entry to finalise.
|
# so _on_process_end has a matching pending entry to finalise.
|
||||||
|
|||||||
@@ -313,19 +313,26 @@ def test_pipeline_overlap_within_command():
|
|||||||
|
|
||||||
|
|
||||||
def test_pe_dma_record_start_after_channel_acquire():
|
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
|
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
|
Under that model, three back-to-back DMAs at t=0 give:
|
||||||
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
|
- ``t_start`` values monotonic and tightly clustered (channel
|
||||||
on command entry, all three ops would share ``t_start == 0`` and the
|
release is fast → each subsequent op's ``t_start`` is at or
|
||||||
second/third would show inflated ``t_end - t_start``.
|
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 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]
|
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]
|
base = durations[0]
|
||||||
assert base > 0, f"first dma duration must be positive, got {base}"
|
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_start values must be monotonic and tightly clustered. The window
|
||||||
# t_end (modulo a few ns of scheduler overhead) — i.e. the wait is
|
# over all three t_starts is bounded by the per-op channel-issue cost,
|
||||||
# *excluded* from the recorded interval, not folded into it.
|
# NOT by the HBM round-trip — i.e. record_start fires post channel
|
||||||
for i in range(1, len(dma_records)):
|
# acquire AND channel release is fast.
|
||||||
prev_end = dma_records[i - 1].t_end
|
t_starts = [r.t_start for r in dma_records]
|
||||||
cur_start = dma_records[i].t_start
|
for i in range(1, len(t_starts)):
|
||||||
assert cur_start >= prev_end - 1.0, (
|
assert t_starts[i] >= t_starts[i - 1], (
|
||||||
f"op {i} t_start={cur_start} began before op {i-1} t_end={prev_end} "
|
f"op {i} t_start={t_starts[i]} regressed before "
|
||||||
f"— channel was not actually held, fix is incorrect"
|
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."
|
||||||
|
)
|
||||||
|
|||||||