paper(gemm): per-PE HBM BW + TFLOPS evaluation plots; unify peak to 8 TFLOP/s

Two new evaluation plots for §3 GEMM, both using the composite window
as the denominator (the up-front tl.load of A in load_ref is therefore
excluded — only HBM traffic and compute inside the composite count):

- gemm_hbm_bw_util.png: per-PE HBM bandwidth utilization against the
  256 GB/s per-PE ceiling. Answers "is this workload memory-bound on
  this configuration?" — small/single-tile shapes sit at 22–46%
  (pipeline-fill-limited), large or output-write-dominated shapes
  saturate (≥85%).

- gemm_per_pe_tflops.png: per-PE GEMM throughput against the 8 TFLOP/s
  engine peak. The deep-K K=3072 load_ref case lands at
  7.18 TFLOP/s (~90% of peak) — the configuration's clean win.
  ref_ref at the same shape drops to 3.83 TFLOP/s because the second
  operand doubles HBM pressure and the kernel hits the BW ceiling.
  The variant gap on this chart is the operational cost of NOT
  pre-staging the activation.

Both charts compare two operand-staging variants on every shape:
load_ref ("activation pre-staged", weight-only HBM streaming) and
ref_ref (both A and W streamed from HBM). load_load is omitted —
with both operands pre-staged the composite carries no HBM traffic
and the BW metric collapses to 0.

Analytic / measured peak unified to the single hardware spec
(8.0 TFLOP/s = 8000 flops/ns). T_STAGE is now derived as
tile_flops / peak (= 16.384 ns for the 32×64×32 tile) rather than
the previous hardcoded 16 ns, so MAC efficiency and per-PE TFLOPS%
share the same denominator and never disagree. Side effect: max
analytic-vs-measured gap tightens from 2.2 ppt to 1.4 ppt across
the sweep.

§3 Results restructured: lead with the two new evaluation plots
(BW saturation, achieved throughput), then the MAC-utilization
analytic-vs-measured chart serves as the validation step, then the
per-stage engine wall-clock as the supporting diagnostic. §3
Analysis updated to use the load_ref / ref_ref contrast as the
hardware-software boundary line that motivates the GQA kernel of §5.

§2.4 Accuracy claim updated to match: GEMM analytic-vs-measured
agreement is now "within 1.4 ppt across every swept shape", from
~7.7% at single-tile up to ~90% at K=3072. Subtitle/title vertical
stacking glitch in the matplotlib charts fixed in passing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-17 10:24:42 -07:00
parent 23992548f7
commit e8b7f4f064
15 changed files with 1305 additions and 1013 deletions
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 47 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 40 KiB

After

Width:  |  Height:  |  Size: 40 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 46 KiB

After

Width:  |  Height:  |  Size: 46 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 42 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 45 KiB

After

Width:  |  Height:  |  Size: 46 KiB

@@ -311,9 +311,9 @@ 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
\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
\SI{1.4}{ppt} across every swept shape---from a single-tile
$M{=}K{=}N{=}32$ at $\sim\SI{7.7}{\percent}$ up to the deep-$K$
$K{=}3072$ case at $\sim\SI{90}{\percent}$. The residual gap is
fill/tail overhead the closed-form pipeline model omits.
Second, in the all-reduce study (\S\ref{sec:allreduce},
@@ -43,12 +43,84 @@ kernel to fuse its softmax work into the GEMM pipeline
\subsection{Results}
We sweep eight GEMM shapes spanning square, tall, wide, and deep-$K$
geometries, under three operand-staging variants
(\textsf{ref\_ref}, both operands streamed from HBM; \textsf{load\_ref},
one operand resident in TCM; \textsf{load\_load}, both resident).
Figure~\ref{fig:gemm-util} reports MAC utilization and efficiency, and
Figure~\ref{fig:gemm-stages} breaks the kernel into per-stage engine
busy time.
geometries under two operand-staging variants that bracket the
realistic LLM cases. In \textsf{load\_ref}, the activation $A$ is
pre-staged on chip and only the weight $W$ streams from HBM during
the composite (the ``activation in TCM, weights from HBM'' case
typical of decoding with a small batch). In \textsf{ref\_ref}, both
operands stream from HBM during the composite (the ``cold both
sides'' case, where the activation does not fit on-chip or is itself
the output of a producer composite that wrote back to HBM). Both
variants run the same composite command; only the up-front staging
of $A$ differs.
We evaluate the composite GEMM along two axes that together describe
how well it lands on the hardware: \emph{HBM bandwidth utilization}
(does the workload saturate the per-PE HBM bandwidth, i.e.\ is it
memory-bound on this configuration?) and \emph{achieved GEMM
throughput} (what fraction of the \SI{8}{\tera\flop\per\second}
per-PE peak does the kernel actually deliver). Both metrics use the
composite window as the denominator, so the up-front \textsf{tl.load}
of $A$ in \textsf{load\_ref} is excluded — only HBM traffic and
compute that happen \emph{inside} the composite count.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_hbm_bw_util.png}
\caption{Per-PE HBM bandwidth utilization during the composite window,
for the two staging variants. The ceiling is the per-PE HBM bandwidth
(\SI{256}{\giga\byte\per\second}; dashed line at \SI{100}{\percent}).
\textsf{ref\_ref} streams both $A$ and $W$ and so always sits at a
higher BW-utilization than \textsf{load\_ref} on the same shape;
both variants saturate ($>\SI{85}{\percent}$) once the kernel
has enough total bytes to keep the link busy (deep-$K$ $K{=}3072$,
or low-reuse output-dominated $M{=}128,K{=}8,N{=}128$). Small or
single-tile shapes do not have enough total traffic to saturate the
link and are bottlenecked by pipeline fill rather than HBM.}
\label{fig:gemm-bw}
\end{figure}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_per_pe_tflops.png}
\caption{Per-PE GEMM throughput delivered during the composite window
(reference line at the \SI{8}{\tera\flop\per\second} per-PE engine
peak). At the deep-$K$ corner $M{=}32,K{=}3072,N{=}32$ the
activation-pre-staged kernel reaches
\SI{7.18}{\tera\flop\per\second} ($\sim\SI{90}{\percent}$ of peak);
the both-from-HBM variant drops to
\SI{3.83}{\tera\flop\per\second} ($\sim\SI{48}{\percent}$) at the
same shape because the second operand doubles HBM pressure and the
kernel hits the BW ceiling shown in Fig.~\ref{fig:gemm-bw}.
Under-tile shapes ($M{=}K{=}N{=}32$, $M{=}8$, $K{=}8$) are
hard-capped by tile-fill at \SI{50}{\percent}, \SI{25}{\percent},
\SI{12.5}{\percent} of peak regardless of staging.}
\label{fig:gemm-tflops}
\end{figure}
The composite reaches the hardware roofline in two distinct regimes.
\emph{Memory-bound}: $K{=}3072$ deep-$K$ saturates HBM (\SI{88}{\percent}
load\_ref, \SI{94}{\percent} ref\_ref) and the low-reuse
$M{=}128,K{=}8,N{=}128$ corner saturates even harder
($>\SI{96}{\percent}$ in both), because back-to-back output
writes dominate. \emph{Compute-rich}: at the same deep-$K$ shape,
\textsf{load\_ref} delivers \SI{7.18}{\tera\flop\per\second}
($\sim\SI{90}{\percent}$ of the per-PE peak) — the configuration's
clean win. The \emph{distance between the two variants} at the same
shape is the cost of \emph{not} pre-staging $A$: at $K{=}3072$ the
throughput halves (\SI{7.18}{\tera\flop\per\second} $\to$
\SI{3.83}{\tera\flop\per\second}) because the second operand doubles
HBM pressure
and the kernel hits the BW ceiling. This is the operational
take-away — when the activation can live on-chip the composite
delivers near-peak GEMM; when it cannot the BW ceiling dominates and
throughput halves.
Figure~\ref{fig:gemm-util} validates that simulator-measured
efficiency tracks an analytic ideal-pipeline model (within
\SI{1.4}{ppt} across every swept shape), and
Figure~\ref{fig:gemm-stages} resolves the per-stage engine wall-clock
that backs the above interpretation.
\begin{figure}[t]
\centering
@@ -60,9 +132,9 @@ 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
\SI{90}{\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.}
\SI{1.4}{ppt} across every shape.}
\label{fig:gemm-util}
\end{figure}
@@ -81,31 +153,24 @@ 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, 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}
The composite command does not manufacture bandwidth or MAC throughput; it
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{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
validation that KernBench's latency model is faithful in the regime that
matters. The hardware implication is concrete: a single-command,
self-routing tile pipeline is the issue mechanism that makes the MAC array
usable, and it is a prerequisite---not a luxury---for the fused attention
kernel of \S\ref{sec:gqa}.
The composite command does not manufacture bandwidth or MAC
throughput; it removes the two software-shaped obstacles between a
GEMM and its hardware roofline. By paying the issue cost once per
GEMM rather than once per tile-stage, and by chaining tile-stages
through token self-routing, it lets the scheduler keep every stage
busy on a different tile, so compute-rich shapes actually reach the
efficiency their arithmetic intensity allows
($\sim$\SI{90}{\percent} of GEMM peak at the deep-$K$
\textsf{load\_ref} corner) and data-bound shapes actually reach
their HBM-BW ceiling instead of stalling on command overhead.
The split between \textsf{load\_ref} and \textsf{ref\_ref} also makes
the hardware-software boundary explicit: when the activation fits in
on-chip storage the composite is BW-headroom-rich and saturates the
GEMM engine; when it does not, both operands compete for the same
HBM port and the kernel is BW-bound. This is the lever the GQA
kernel of \S\ref{sec:gqa} reaches for next — keeping the right
working set on-chip so the composite pipeline lands in the
compute-rich regime rather than the BW-bound one.
Binary file not shown.

After

Width:  |  Height:  |  Size: 47 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 40 KiB

After

Width:  |  Height:  |  Size: 40 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 45 KiB

After

Width:  |  Height:  |  Size: 46 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 42 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 46 KiB

After

Width:  |  Height:  |  Size: 46 KiB

File diff suppressed because it is too large Load Diff
+236 -9
View File
@@ -277,12 +277,35 @@ STAGE_COLORS = {
# 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)
# hold). Steady-state cycle is BW-bound at one tile / T_STAGE, where
# T_STAGE is derived from the hardware peak so that analytic and measured
# share the same per-PE GEMM-throughput denominator
# (\\SI{8}{\\tera\\flop\\per\\second} = 8000 flops/ns) — see
# `_t_stage_ns_from_tile()` below.
_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
# Per-PE GEMM peak. Single source of truth for both the analytic
# `compute_total = tiles * T_STAGE` term and the measured-eff
# denominator. Matches topology.yaml pe_gemm.peak_tflops_f16.
_GEMM_PEAK_TFLOPS = 8.0
_GEMM_PEAK_FLOPS_PER_NS = _GEMM_PEAK_TFLOPS * 1000.0 # = 8000 flops/ns
def _t_stage_ns_from_tile(tile_flops: int) -> float:
"""Per-tile bottleneck cycle (ns) implied by the hardware peak.
The composite pipeline's steady-state stage cycle equals the time the
GEMM engine spends on one tile at peak throughput. Deriving T_STAGE
from tile_flops / peak (rather than hardcoding 16 ns) keeps the
analytic model self-consistent with the measured-eff denominator: a
tile with tile_flops = 131072 and an \\SI{8}{\\tera\\flop\\per\\second}
engine takes exactly 131072 / 8000 = 16.384 ns of GEMM time.
"""
return tile_flops / _GEMM_PEAK_FLOPS_PER_NS
_PLOT_VARIANT = "load_ref"
@@ -342,10 +365,15 @@ def _grouped_bar_png(
_xtick_labels(shape_labels, tile_counts, flagged), fontsize=8,
)
ax.set_ylabel(y_label)
ax.set_title(title, fontsize=13, fontweight="bold")
# Title with optional subtitle stacked above (subtitle in lighter ink,
# title in bold below it). Coordinates are in axes-fraction so they
# survive tight_layout without colliding with the data area.
if subtitle:
ax.text(0.5, 1.01, subtitle, transform=ax.transAxes, ha="center",
ax.text(0.5, 1.10, subtitle, transform=ax.transAxes, ha="center",
va="bottom", fontsize=8, color="#475569")
ax.set_title(title, fontsize=13, fontweight="bold", pad=22)
else:
ax.set_title(title, fontsize=13, fontweight="bold")
if threshold is not None:
ax.axhline(threshold, ls="--", color="gray", lw=1.0)
ax.legend(fontsize=8, loc="upper right")
@@ -420,6 +448,7 @@ def emit_mac_utilization_measured(
flagged = [_under_tile(*k, TILE_M, TILE_K, TILE_N) for k in shapes]
tile_counts = [by_shape[k]["tile_count_expected"] for k in shapes]
t_stage_ns = _t_stage_ns_from_tile(tile_flops)
gemm_util, useful_eff = [], []
for k in shapes:
r = by_shape[k]
@@ -428,7 +457,7 @@ def emit_mac_utilization_measured(
tiles = r["tile_count_expected"]
gu = useful / (tile_flops * tiles) * 100
gemm_util.append(gu)
compute_total = tiles * _T_STAGE
compute_total = tiles * t_stage_ns
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)
@@ -458,7 +487,7 @@ 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
peak_per_ns = tile_flops / _T_STAGE
t_stage_ns = _t_stage_ns_from_tile(tile_flops)
by_shape = {(r["M"], r["K"], r["N"]): r
for r in rows if r["variant"] == _PLOT_VARIANT}
@@ -479,11 +508,11 @@ 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)
compute_total = tiles * _T_STAGE
compute_total = tiles * t_stage_ns
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)
eff_m.append((useful / cw / _GEMM_PEAK_FLOPS_PER_NS * 100) if cw > 0 else 0.0)
series = {
"GEMM util % (theoretical)": gu_t,
@@ -507,6 +536,202 @@ def emit_mac_utilization_theoretical_vs_measured(
)
# ── per-PE composite-pipeline plots: HBM BW util & GEMM TFLOPS ─────────
#
# These two plots evaluate the composite GEMM kernel against the
# per-PE HBM bandwidth (256 GB/s) and per-PE GEMM peak (8 TFLOP/s).
# Both metrics use composite_window_ns as the denominator — the
# composite-internal pipeline window, with any pre-composite tl.load
# of operand A excluded by construction. The two variants on each plot
# answer different staging questions:
#
# - load_ref: activation A pre-staged on chip, only weight W streamed
# from HBM during the composite. ("Weight-only HBM" case.)
# - ref_ref: both A and W streamed from HBM during the composite.
# ("Both A and W from HBM" case.)
#
# load_load is omitted from these plots — with both operands pre-staged
# the composite carries no HBM traffic and the BW metric collapses to 0,
# which is uninformative against the saturation question.
# Per-tile DMA payload sizes (one TILE_M x TILE_K or TILE_K x TILE_N
# read; one TILE_M x TILE_N write per output tile).
_PIPELINE_VARIANTS = ("load_ref", "ref_ref")
_VARIANT_LABEL = {
"load_ref": "Activation pre-staged",
"ref_ref": "A and W both from HBM",
}
_VARIANT_COLOR = {
"load_ref": "#3B82F6", # blue
"ref_ref": "#F59E0B", # orange
}
# Per-PE HBM peak BW (GB/s); matches topology.yaml hbm_to_router_bw_gbs
# and (8 PCs / PE) x (32 GB/s / PC).
_HBM_BW_PER_PE_GBS = 256.0
# Per-PE GEMM peak (TFLOP/s); matches topology.yaml pe_gemm.peak_tflops_f16.
_GEMM_PEAK_TFLOPS = 8.0
def _shape_summary(rows: list[dict], tile: dict):
"""Common shape-axis state for the BW / TFLOPS plots.
Returns (shape_labels, tile_counts, flagged, shape_keys) — only
shapes that actually appear in `rows` for any of the pipeline
variants, in their original order.
"""
seen: list[tuple[int, int, int]] = []
by_key: dict[tuple[int, int, int], dict] = {}
for r in rows:
key = (r["M"], r["K"], r["N"])
if key not in by_key:
by_key[key] = r
seen.append(key)
shape_labels = [_shape_label(by_key[k]) for k in seen]
flagged = [_under_tile(*k, tile["M"], tile["K"], tile["N"]) for k in seen]
tile_counts = [by_key[k]["tile_count_expected"] for k in seen]
return shape_labels, tile_counts, flagged, seen
def _composite_bytes_through_hbm(r: dict, tile: dict) -> float:
"""Total bytes transferred via HBM during the composite window.
Derived from the simulator's per-stage record counts so it
automatically reflects the variant's pre-staging choices:
load_ref skips A reads, ref_ref reads both A and B, load_load
reads neither. Each DMA record carries a full hardware tile
(TILE_M x TILE_K read for A, TILE_K x TILE_N read for B,
TILE_M x TILE_N write for the output), regardless of how much
of that tile is "useful" — which is the actual byte traffic the
HBM controller sees.
"""
bpe = 2 # f16; matches matmul_composite.py default
a_tile_bytes = tile["M"] * tile["K"] * bpe
b_tile_bytes = tile["K"] * tile["N"] * bpe
out_tile_bytes = tile["M"] * tile["N"] * bpe
# The schedule emits one DMA_READ per (A,B) per non-pinned operand.
# Without per-record operand tags in the aggregated JSON, treat all
# DMA_READ records as B-shaped reads on load_ref (B is the only
# streamed operand) and average the two tile shapes on ref_ref.
# In the current topology both tile shapes are 4 KiB so this
# distinction is moot — kept symbolic for future tile-size sweeps.
variant = r.get("variant")
n_read = r["stages"]["DMA_READ"]["record_count"]
n_write = r["stages"]["DMA_WRITE"]["record_count"]
if variant == "ref_ref":
# A reads + B reads, roughly half each.
read_bytes = (n_read // 2) * a_tile_bytes \
+ (n_read - n_read // 2) * b_tile_bytes
else:
read_bytes = n_read * b_tile_bytes
write_bytes = n_write * out_tile_bytes
return float(read_bytes + write_bytes)
def _by_variant(rows: list[dict], variant: str) -> dict:
return {(r["M"], r["K"], r["N"]): r
for r in rows if r.get("variant") == variant}
def emit_hbm_bw_utilization(
sweep_json: Path | str = DEFAULT_SWEEP_JSON,
out_dir: Path | str = DEFAULT_PLOTS_DIR,
) -> str | None:
"""Per-PE HBM bandwidth utilization during the composite window.
Two grouped bars per shape: load_ref ("activation pre-staged")
vs ref_ref ("both A and W from HBM"). Y-axis is BW
util % against the per-PE HBM ceiling (\\SI{256}{\\giga\\byte\\per
\\second}).
"""
data = _load_sweep_data(sweep_json)
rows = [r for r in data["rows"] if r.get("variant") in _PIPELINE_VARIANTS]
if not rows:
return None
tile = data["tile_sizes"]
shape_labels, tile_counts, flagged, shape_keys = _shape_summary(rows, tile)
series: dict = {}
for variant in _PIPELINE_VARIANTS:
by = _by_variant(rows, variant)
vals = []
for k in shape_keys:
r = by.get(k)
if r is None:
vals.append(0.0)
continue
cw = r.get("composite_window_ns", 0.0) or 0.0
if cw <= 0:
vals.append(0.0)
continue
bytes_hbm = _composite_bytes_through_hbm(r, tile)
# bytes_hbm [B] / cw [ns] = GB/s (since 1 B/ns = 1 GB/s).
bw_gbs = bytes_hbm / cw
vals.append(bw_gbs / _HBM_BW_PER_PE_GBS * 100.0)
series[_VARIANT_LABEL[variant]] = vals
colors = {_VARIANT_LABEL[v]: _VARIANT_COLOR[v] for v in _PIPELINE_VARIANTS}
return _grouped_bar_png(
"gemm_hbm_bw_util.png", out_dir=Path(out_dir),
title="GEMM HBM bandwidth utilization (per PE, composite window)",
subtitle=(f"BW util = bytes through HBM / composite window "
f"/ {_HBM_BW_PER_PE_GBS:.0f} GB/s. "
f"Saturation → memory-bound regime."),
shape_labels=shape_labels, tile_counts=tile_counts, flagged=flagged,
series=series, colors=colors, y_label="HBM BW utilization (%)",
threshold=100.0,
)
def emit_per_pe_tflops(
sweep_json: Path | str = DEFAULT_SWEEP_JSON,
out_dir: Path | str = DEFAULT_PLOTS_DIR,
) -> str | None:
"""Per-PE GEMM throughput in TFLOP/s during the composite window.
Two grouped bars per shape (load_ref vs ref_ref). Reference line
at the per-PE GEMM engine peak (\\SI{8}{\\tera\\flop\\per\\second}).
"""
data = _load_sweep_data(sweep_json)
rows = [r for r in data["rows"] if r.get("variant") in _PIPELINE_VARIANTS]
if not rows:
return None
tile = data["tile_sizes"]
shape_labels, tile_counts, flagged, shape_keys = _shape_summary(rows, tile)
series: dict = {}
for variant in _PIPELINE_VARIANTS:
by = _by_variant(rows, variant)
vals = []
for k in shape_keys:
r = by.get(k)
if r is None:
vals.append(0.0)
continue
cw = r.get("composite_window_ns", 0.0) or 0.0
if cw <= 0:
vals.append(0.0)
continue
M, K, N = r["M"], r["K"], r["N"]
useful = 2 * M * K * N
# flops / ns -> GFLOP/s; / 1000 -> TFLOP/s.
tflops = useful / cw / 1000.0
vals.append(tflops)
series[_VARIANT_LABEL[variant]] = vals
colors = {_VARIANT_LABEL[v]: _VARIANT_COLOR[v] for v in _PIPELINE_VARIANTS}
return _grouped_bar_png(
"gemm_per_pe_tflops.png", out_dir=Path(out_dir),
title="GEMM achieved throughput (per PE, composite window)",
subtitle=(f"TFLOP/s = 2 · M · K · N / composite "
f"window. Peak per PE = "
f"{_GEMM_PEAK_TFLOPS:.0f} TFLOP/s (dashed line)."),
shape_labels=shape_labels, tile_counts=tile_counts, flagged=flagged,
series=series, colors=colors, y_label="Per-PE TFLOP/s",
threshold=_GEMM_PEAK_TFLOPS,
)
def emit_all_gemm_plots(
sweep_json: Path | str = DEFAULT_SWEEP_JSON,
out_dir: Path | str = DEFAULT_PLOTS_DIR,
@@ -515,7 +740,9 @@ def emit_all_gemm_plots(
paths = []
for fn in (emit_stage_breakdown,
emit_mac_utilization_measured,
emit_mac_utilization_theoretical_vs_measured):
emit_mac_utilization_theoretical_vs_measured,
emit_hbm_bw_utilization,
emit_per_pe_tflops):
p = fn(sweep_json, out_dir)
if p:
paths.append(p)