"""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",
    )