ywkang
e8b7f4f064
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>
782 lines
30 KiB
Python
782 lines
30 KiB
Python
"""milestone-1h-gemm bench: GEMM evaluation harness (sweep + figures).
|
||
|
||
Self-contained milestone bench (ADR-0054). Holds the shape×variant sweep
|
||
and the figure renderers; the ``run(torch)`` entry at the bottom runs the
|
||
sweep (or reuses the committed JSON when ``MILESTONE_FAST=1``) and writes
|
||
every figure into ``benches/1H_milestone_output/gemm/``.
|
||
|
||
This is the single home for the GEMM eval logic: the figure tests import a
|
||
thin re-export shim (``tests/gemm/_gemm_plot_helpers.py``), as does the
|
||
``scripts/gemm_sweep.py`` wrapper.
|
||
|
||
The sweep drives ``matmul-composite`` across shapes×variants through the
|
||
same ``run_bench`` path the CLI uses, harvests ``result.engine.op_log``,
|
||
and writes the sweep JSON. The renderers read that JSON and emit matplotlib
|
||
PNGs. No simulation in the renderers — they are fast.
|
||
|
||
Chart set (mirrors the GEMM MAC slides in scripts/build_overview_slides.py):
|
||
- stage breakdown (load_ref operand staging)
|
||
- MAC utilization — measured (load_ref)
|
||
- MAC utilization — theoretical vs measured (load_ref)
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
from pathlib import Path
|
||
|
||
from kernbench.benches.registry import bench
|
||
from kernbench.policy.placement.dp import DPPolicy
|
||
|
||
ROOT = Path(__file__).resolve().parents[3]
|
||
DEFAULT_SWEEP_JSON = ROOT / "docs" / "diagrams" / "gemm_sweep.json"
|
||
DEFAULT_PLOTS_DIR = ROOT / "docs" / "diagrams" / "gemm_plots"
|
||
_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gemm"
|
||
|
||
# ── sweep configuration ────────────────────────────────────────────────
|
||
|
||
# Default sweep covering under-tile, single-tile, multi-tile, and asymmetric
|
||
# regimes. Each entry is "MxKxN" or a single int (square M=K=N).
|
||
# Override via env: SWEEP_SHAPES="16,32,16x2048x16,..."
|
||
DEFAULT_SHAPES = [
|
||
"32x32x32", # 1 tile, K=32 < TILE_K=64 → under-tile in K
|
||
"32x64x32", # 1 tile, exact single-tile fit
|
||
"32x128x32", # 2 tiles, aligned
|
||
"32x128x128", # 8 tiles, aligned
|
||
"32x3072x32", # 48 tiles, all K-axis (tall-skinny)
|
||
"8x128x128", # 8 tiles, but M=8 < TILE_M=32 → MAC array under-fed
|
||
"128x8x128", # 16 tiles, but K=8 < TILE_K=64 → MAC array under-fed
|
||
"512", # 2048 tiles, fully aligned — "well-pipelined" reference
|
||
]
|
||
|
||
# Operand-staging variants exercised per shape.
|
||
VARIANTS = ["ref_ref", "load_ref", "load_load"]
|
||
|
||
# Engines whose timings we collect (component_id suffix match).
|
||
ENGINES = ["pe_dma", "pe_fetch_store", "pe_gemm", "pe_math"]
|
||
|
||
# Per-stage breakdown labels (StageType enum names from pe_types.py).
|
||
STAGES = ["DMA_READ", "DMA_WRITE", "FETCH", "STORE", "GEMM", "MATH"]
|
||
|
||
# Scheduler tile sizes (mirror of PeSchedulerComponent.TILE_M/K/N).
|
||
TILE_M, TILE_K, TILE_N = 32, 64, 32
|
||
|
||
|
||
def _ceil(a: int, b: int) -> int:
|
||
return (a + b - 1) // b
|
||
|
||
|
||
def _engine_wall_ns(records, suffix: str) -> float:
|
||
"""Wall-clock interval the engine was active (union of overlapping ops)."""
|
||
intervals = [(r.t_start, r.t_end) for r in records
|
||
if r.component_id.endswith("." + suffix)]
|
||
if not intervals:
|
||
return 0.0
|
||
intervals.sort()
|
||
merged_end = intervals[0][1]
|
||
merged_start = intervals[0][0]
|
||
total = 0.0
|
||
for s, e in intervals[1:]:
|
||
if s <= merged_end:
|
||
merged_end = max(merged_end, e)
|
||
else:
|
||
total += merged_end - merged_start
|
||
merged_start, merged_end = s, e
|
||
total += merged_end - merged_start
|
||
return total
|
||
|
||
|
||
def _engine_occupancy_ns(records, suffix: str) -> float:
|
||
return sum(r.t_end - r.t_start for r in records
|
||
if r.component_id.endswith("." + suffix))
|
||
|
||
|
||
def _engine_count(records, suffix: str) -> int:
|
||
return sum(1 for r in records if r.component_id.endswith("." + suffix))
|
||
|
||
|
||
def _stage_occupancy_ns(records, stage_type: str) -> float:
|
||
return sum(
|
||
r.t_end - r.t_start
|
||
for r in records
|
||
if r.params.get("stage_type") == stage_type
|
||
)
|
||
|
||
|
||
def _stage_wall_ns(records, stage_type: str) -> float:
|
||
"""Interval-union wall-clock for records whose stage_type matches."""
|
||
intervals = sorted(
|
||
(r.t_start, r.t_end) for r in records
|
||
if r.params.get("stage_type") == stage_type
|
||
)
|
||
if not intervals:
|
||
return 0.0
|
||
total = 0.0
|
||
cs, ce = intervals[0]
|
||
for s, e in intervals[1:]:
|
||
if s <= ce:
|
||
ce = max(ce, e)
|
||
else:
|
||
total += ce - cs
|
||
cs, ce = s, e
|
||
total += ce - cs
|
||
return total
|
||
|
||
|
||
def _stage_count(records, stage_type: str) -> int:
|
||
return sum(1 for r in records if r.params.get("stage_type") == stage_type)
|
||
|
||
|
||
def _run_one(M: int, K: int, N: int, topology: str, variant: str = "ref_ref") -> dict:
|
||
os.environ["MATMUL_M"] = str(M)
|
||
os.environ["MATMUL_K"] = str(K)
|
||
os.environ["MATMUL_N"] = str(N)
|
||
os.environ["MATMUL_VARIANT"] = variant
|
||
|
||
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
|
||
from kernbench.sim_engine.engine import GraphEngine
|
||
from kernbench.topology.builder import resolve_topology
|
||
|
||
topo = resolve_topology(topology)
|
||
bench = resolve_bench("matmul-composite").run
|
||
device = resolve_device(None)
|
||
|
||
t0 = time.time()
|
||
result = run_bench(
|
||
topology=topo, bench_fn=bench, device=device,
|
||
engine_factory=lambda t, d: GraphEngine(
|
||
getattr(t, "topology_obj", t), enable_data=True,
|
||
),
|
||
)
|
||
wall = time.time() - t0
|
||
|
||
op_log = result.engine.op_log
|
||
if not result.completion.ok:
|
||
raise RuntimeError(f"bench failed at M={M},K={K},N={N}: {result.completion}")
|
||
|
||
# Bytes touched at f16 (2 B): full A + full B + full out (each operand
|
||
# streamed once through HBM by the composite plan).
|
||
bytes_total = (M * K + K * N + M * N) * 2
|
||
row = {
|
||
"M": M, "K": K, "N": N,
|
||
"variant": variant,
|
||
"flops": 2 * M * K * N,
|
||
"bytes_hbm": bytes_total,
|
||
"arith_intensity": (2 * M * K * N) / bytes_total, # flops/byte
|
||
"tile_count_expected": _ceil(M, TILE_M) * _ceil(N, TILE_N) * _ceil(K, TILE_K),
|
||
"sim_wall_clock_s": round(wall, 3),
|
||
"engines": {},
|
||
}
|
||
for eng in ENGINES:
|
||
row["engines"][eng] = {
|
||
"occupancy_ns": _engine_occupancy_ns(op_log, eng),
|
||
"wall_ns": _engine_wall_ns(op_log, eng),
|
||
"record_count": _engine_count(op_log, eng),
|
||
}
|
||
row["stages"] = {}
|
||
for stage in STAGES:
|
||
row["stages"][stage] = {
|
||
"occupancy_ns": _stage_occupancy_ns(op_log, stage),
|
||
"wall_ns": _stage_wall_ns(op_log, stage),
|
||
"record_count": _stage_count(op_log, stage),
|
||
}
|
||
# Kernel-window wall-clock = max t_end - min t_start over PE engine records.
|
||
pe_records = [r for r in op_log
|
||
if any(r.component_id.endswith("." + e) for e in ENGINES)]
|
||
if pe_records:
|
||
row["pe_window_ns"] = max(r.t_end for r in pe_records) \
|
||
- min(r.t_start for r in pe_records)
|
||
else:
|
||
row["pe_window_ns"] = 0.0
|
||
stage_records = [r for r in op_log
|
||
if r.params.get("stage_type") in STAGES]
|
||
if stage_records:
|
||
row["composite_window_ns"] = max(r.t_end for r in stage_records) \
|
||
- min(r.t_start for r in stage_records)
|
||
else:
|
||
row["composite_window_ns"] = 0.0
|
||
return row
|
||
|
||
|
||
def _parse_shapes(raw) -> list[tuple[int, int, int]]:
|
||
shapes: list[tuple[int, int, int]] = []
|
||
for s in raw:
|
||
s = s.strip()
|
||
if not s:
|
||
continue
|
||
if "x" in s.lower():
|
||
parts = s.lower().split("x")
|
||
shapes.append((int(parts[0]), int(parts[1]), int(parts[2])))
|
||
else:
|
||
v = int(s)
|
||
shapes.append((v, v, v))
|
||
return shapes
|
||
|
||
|
||
def run_sweep(out_json: Path | str = DEFAULT_SWEEP_JSON) -> Path:
|
||
"""Drive matmul-composite across shapes×variants; write the sweep JSON.
|
||
|
||
Honors ``SWEEP_SHAPES`` / ``SWEEP_TOPOLOGY`` env overrides (same as the
|
||
historical ``scripts/gemm_sweep.py``). Returns the JSON path written.
|
||
"""
|
||
shapes_env = os.environ.get("SWEEP_SHAPES")
|
||
raw = (shapes_env.split(",") if shapes_env else DEFAULT_SHAPES)
|
||
shapes = _parse_shapes(raw)
|
||
topology = os.environ.get("SWEEP_TOPOLOGY", "topology.yaml")
|
||
|
||
rows = []
|
||
for M, K, N in shapes:
|
||
for variant in VARIANTS:
|
||
print(f"[sweep] M={M} K={K} N={N} variant={variant} ...", flush=True)
|
||
row = _run_one(M, K, N, topology, variant=variant)
|
||
rows.append(row)
|
||
eng_dma = row["engines"]["pe_dma"]
|
||
eng_gem = row["engines"]["pe_gemm"]
|
||
print(f" tiles={row['tile_count_expected']:>6} "
|
||
f"pe_window={row['pe_window_ns']:8.1f}ns "
|
||
f"dma_occ={eng_dma['occupancy_ns']:9.1f} "
|
||
f"gemm_occ={eng_gem['occupancy_ns']:8.1f} "
|
||
f"(sim {row['sim_wall_clock_s']:.1f}s)")
|
||
|
||
out_json = Path(out_json)
|
||
out_json.parent.mkdir(parents=True, exist_ok=True)
|
||
out_json.write_text(json.dumps({
|
||
"tile_sizes": {"M": TILE_M, "K": TILE_K, "N": TILE_N},
|
||
"engines": ENGINES,
|
||
"stages": STAGES,
|
||
"variants": VARIANTS,
|
||
"rows": rows,
|
||
}, indent=2))
|
||
print(f"\n[sweep] wrote {out_json}")
|
||
return out_json
|
||
|
||
|
||
# ── figure rendering ───────────────────────────────────────────────────
|
||
|
||
# Shapes excluded from the figures (mirrors build_overview_slides).
|
||
EXCLUDED_SHAPES = {(512, 512, 512)}
|
||
|
||
# Stage bars shown (raw op_log stage_type keys) + display names + colors.
|
||
STAGE_KEYS = ["DMA_READ", "FETCH", "GEMM", "DMA_WRITE"]
|
||
STAGE_DISPLAY = {
|
||
"DMA_READ": "DMA in",
|
||
"FETCH": "Fetch",
|
||
"GEMM": "GEMM",
|
||
"DMA_WRITE": "DMA out",
|
||
}
|
||
STAGE_COLORS = {
|
||
"DMA_READ": "#3B82F6",
|
||
"FETCH": "#10B981",
|
||
"GEMM": "#F59E0B",
|
||
"DMA_WRITE": "#A855F7",
|
||
}
|
||
|
||
# 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, 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"
|
||
|
||
|
||
def _load_sweep_data(sweep_json: Path | str = DEFAULT_SWEEP_JSON) -> dict:
|
||
sweep_json = Path(sweep_json)
|
||
if not sweep_json.exists():
|
||
return {"rows": []}
|
||
data = json.loads(sweep_json.read_text())
|
||
data["rows"] = [
|
||
r for r in data.get("rows", [])
|
||
if (r["M"], r["K"], r["N"]) not in EXCLUDED_SHAPES
|
||
]
|
||
return data
|
||
|
||
|
||
def _shape_label(r: dict) -> str:
|
||
if r["M"] == r["K"] == r["N"]:
|
||
return f"M=K=N={r['M']}"
|
||
return f"M={r['M']} K={r['K']} N={r['N']}"
|
||
|
||
|
||
def _under_tile(M, K, N, tile_M, tile_K, tile_N) -> bool:
|
||
return M < tile_M or K < tile_K or N < tile_N
|
||
|
||
|
||
def _xtick_labels(shape_labels, tile_counts, flagged) -> list[str]:
|
||
out = []
|
||
for lbl, tc, fl in zip(shape_labels, tile_counts, flagged):
|
||
s = f"{lbl}\n({tc} tiles)"
|
||
if fl:
|
||
s += " *"
|
||
out.append(s)
|
||
return out
|
||
|
||
|
||
def _grouped_bar_png(
|
||
out_name: str, *, out_dir: Path, title: str, subtitle: str | None,
|
||
shape_labels, tile_counts, flagged, series: dict, colors: dict,
|
||
y_label: str, threshold: float | None = None, footnote: str | None = None,
|
||
) -> str:
|
||
"""Render one grouped-bar chart to out_dir/out_name; return the path."""
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
|
||
n_groups = len(shape_labels)
|
||
n_series = max(1, len(series))
|
||
x = np.arange(n_groups)
|
||
width = 0.8 / n_series
|
||
|
||
fig, ax = plt.subplots(figsize=(11, 6))
|
||
for i, (name, vals) in enumerate(series.items()):
|
||
offset = (i - (n_series - 1) / 2) * width
|
||
ax.bar(x + offset, vals, width, label=name, color=colors.get(name))
|
||
|
||
ax.set_xticks(x)
|
||
ax.set_xticklabels(
|
||
_xtick_labels(shape_labels, tile_counts, flagged), fontsize=8,
|
||
)
|
||
ax.set_ylabel(y_label)
|
||
# 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.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")
|
||
ax.grid(True, axis="y", alpha=0.3)
|
||
|
||
caption = "* = under-tile shape (M<TILE_M, K<TILE_K, or N<TILE_N)"
|
||
if footnote:
|
||
caption = footnote + "\n" + caption
|
||
fig.text(0.5, 0.01, caption, ha="center", fontsize=7, color="gray",
|
||
wrap=True)
|
||
|
||
fig.tight_layout(rect=(0, 0.05, 1, 1))
|
||
out_dir = Path(out_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
out = out_dir / out_name
|
||
fig.savefig(out, dpi=120)
|
||
plt.close(fig)
|
||
return str(out)
|
||
|
||
|
||
def emit_stage_breakdown(
|
||
sweep_json: Path | str = DEFAULT_SWEEP_JSON,
|
||
out_dir: Path | str = DEFAULT_PLOTS_DIR,
|
||
) -> str | None:
|
||
"""Per-stage engine wall-clock per shape (load_ref operand staging)."""
|
||
data = _load_sweep_data(sweep_json)
|
||
rows = [r for r in data["rows"] if r.get("variant") == _PLOT_VARIANT]
|
||
if not rows:
|
||
return None
|
||
tile = data["tile_sizes"]
|
||
shape_labels = [_shape_label(r) for r in rows]
|
||
flagged = [_under_tile(r["M"], r["K"], r["N"], tile["M"], tile["K"], tile["N"])
|
||
for r in rows]
|
||
tile_counts = [r["tile_count_expected"] for r in rows]
|
||
series = {
|
||
STAGE_DISPLAY[s]: [r.get("stages", {}).get(s, {}).get("wall_ns", 0.0)
|
||
for r in rows]
|
||
for s in STAGE_KEYS
|
||
}
|
||
colors = {STAGE_DISPLAY[s]: STAGE_COLORS[s] for s in STAGE_KEYS}
|
||
return _grouped_bar_png(
|
||
"gemm_stage_breakdown.png", out_dir=Path(out_dir),
|
||
title="GEMM stage breakdown",
|
||
subtitle=(f"Per-stage engine wall-clock (DMA in / Fetch / GEMM / "
|
||
f"DMA out), {_PLOT_VARIANT} staging. "
|
||
f"Tile {tile['M']}x{tile['K']}x{tile['N']}."),
|
||
shape_labels=shape_labels, tile_counts=tile_counts, flagged=flagged,
|
||
series=series, colors=colors, y_label="ns",
|
||
footnote="Bars = engine wall-clock interval (merged overlaps).",
|
||
)
|
||
|
||
|
||
def emit_mac_utilization_measured(
|
||
sweep_json: Path | str = DEFAULT_SWEEP_JSON,
|
||
out_dir: Path | str = DEFAULT_PLOTS_DIR,
|
||
) -> str | None:
|
||
"""GEMM util % and useful pipeline-eff % (analytical model, load_ref)."""
|
||
data = _load_sweep_data(sweep_json)
|
||
rows = data["rows"]
|
||
if not rows:
|
||
return None
|
||
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
|
||
|
||
by_shape = {(r["M"], r["K"], r["N"]): r
|
||
for r in rows if r["variant"] == _PLOT_VARIANT}
|
||
shapes = list(by_shape)
|
||
if not shapes:
|
||
return None
|
||
shape_labels = [_shape_label(by_shape[k]) for k in shapes]
|
||
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]
|
||
M, K, N = r["M"], r["K"], r["N"]
|
||
useful = 2 * M * K * N
|
||
tiles = r["tile_count_expected"]
|
||
gu = useful / (tile_flops * tiles) * 100
|
||
gemm_util.append(gu)
|
||
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)
|
||
|
||
series = {"GEMM util %": gemm_util, "Useful eff %": useful_eff}
|
||
colors = {"GEMM util %": "#10B981", "Useful eff %": "#F59E0B"}
|
||
return _grouped_bar_png(
|
||
"gemm_mac_utilization_measured.png", out_dir=Path(out_dir),
|
||
title="GEMM MAC utilization — load_ref",
|
||
subtitle=("GEMM util = useful FLOPs / (tile FLOPs x tiles); "
|
||
"Useful eff = GEMM util x ideal pipeline efficiency."),
|
||
shape_labels=shape_labels, tile_counts=tile_counts, flagged=flagged,
|
||
series=series, colors=colors, y_label="%", threshold=100.0,
|
||
footnote="Theoretical ideal-pipeline model (not simulator data).",
|
||
)
|
||
|
||
|
||
def emit_mac_utilization_theoretical_vs_measured(
|
||
sweep_json: Path | str = DEFAULT_SWEEP_JSON,
|
||
out_dir: Path | str = DEFAULT_PLOTS_DIR,
|
||
) -> str | None:
|
||
"""Theoretical vs simulator-measured GEMM util / useful eff (load_ref)."""
|
||
data = _load_sweep_data(sweep_json)
|
||
rows = data["rows"]
|
||
if not rows:
|
||
return None
|
||
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
|
||
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}
|
||
shapes = list(by_shape)
|
||
if not shapes:
|
||
return None
|
||
shape_labels = [_shape_label(by_shape[k]) for k in shapes]
|
||
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]
|
||
|
||
gu_t, gu_m, eff_t, eff_m = [], [], [], []
|
||
for k in shapes:
|
||
r = by_shape[k]
|
||
M, K, N = r["M"], r["K"], r["N"]
|
||
useful = 2 * M * K * N
|
||
tiles = r["tile_count_expected"]
|
||
gut = useful / (tile_flops * tiles)
|
||
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_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 / _GEMM_PEAK_FLOPS_PER_NS * 100) if cw > 0 else 0.0)
|
||
|
||
series = {
|
||
"GEMM util % (theoretical)": gu_t,
|
||
"GEMM util % (measured)": gu_m,
|
||
"Theoretical eff %": eff_t,
|
||
"Measured eff %": eff_m,
|
||
}
|
||
colors = {
|
||
"GEMM util % (theoretical)": "#10B981",
|
||
"GEMM util % (measured)": "#6EE7B7",
|
||
"Theoretical eff %": "#F59E0B",
|
||
"Measured eff %": "#3B82F6",
|
||
}
|
||
return _grouped_bar_png(
|
||
"gemm_mac_utilization_theoretical_vs_measured.png", out_dir=Path(out_dir),
|
||
title="GEMM MAC utilization — theoretical vs measured (load_ref)",
|
||
subtitle=("theoretical model vs simulator op_log; agreement "
|
||
"validates the analytical pipeline model."),
|
||
shape_labels=shape_labels, tile_counts=tile_counts, flagged=flagged,
|
||
series=series, colors=colors, y_label="%", threshold=100.0,
|
||
)
|
||
|
||
|
||
# ── 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,
|
||
) -> list[str]:
|
||
"""Render every GEMM figure that has data; return the paths written."""
|
||
paths = []
|
||
for fn in (emit_stage_breakdown,
|
||
emit_mac_utilization_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)
|
||
return paths
|
||
|
||
|
||
# ── bench entry ────────────────────────────────────────────────────────
|
||
|
||
|
||
@bench(
|
||
name="milestone-1h-gemm",
|
||
description="1H milestone: regenerate all GEMM results + figures.",
|
||
)
|
||
def run(torch) -> None:
|
||
"""Run the GEMM sweep (or reuse committed JSON) and render every figure.
|
||
|
||
``MILESTONE_FAST=1`` reuses the committed ``DEFAULT_SWEEP_JSON`` (seconds);
|
||
otherwise the full sweep runs into ``out_dir/gemm_sweep.json`` (minutes).
|
||
The sweep drives its own engines, so a sentinel tensor is submitted at the
|
||
end to satisfy the run_bench contract (ADR-0045 D4).
|
||
"""
|
||
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||
fast = bool(os.environ.get("MILESTONE_FAST"))
|
||
if fast:
|
||
sweep_json = DEFAULT_SWEEP_JSON
|
||
else:
|
||
sweep_json = run_sweep(out_json=_OUTPUT_DIR / "gemm_sweep.json")
|
||
paths = emit_all_gemm_plots(sweep_json=sweep_json, out_dir=_OUTPUT_DIR)
|
||
print(f" milestone-1h-gemm: {len(paths)} figures -> {_OUTPUT_DIR} "
|
||
f"(fast={fast})")
|
||
|
||
torch.zeros(
|
||
(1, 1), dtype="f16",
|
||
dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
|
||
name="milestone_gemm_sentinel",
|
||
)
|