kernbench2/scripts/paper/paper_plot_gqa_decode_composite_breakdown.py

"""Latency breakdown bar chart for the Case-6 composite-command decode study.

Reads sweep_decode_composite.json and writes a single-figure stacked bar
chart comparing three variants — primitive hand-tiled (16×16×16),
composite GEMM, composite + softmax_merge — at the S_kv = 131 072 point:

  bottom stack: PE_CPU dispatch time (from pe_cpu_dispatch_cycles, ns at
                1 GHz — ADR-0064 Rev2 D3)
  top stack:    engine time (latency_ns − dispatch cycles) — DMA / GEMM /
                MATH / IPCQ work the engine flushes on the critical path

The dispatch/engine split is a first-order breakdown; in reality the
two paths overlap partially. The note on the plot calls this out.

Run (after the composite bench sweep):
    python scripts/paper/paper_plot_gqa_decode_composite_breakdown.py
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

_REPO_ROOT = Path(__file__).resolve().parents[2]
_FIG_DIR = (
    _REPO_ROOT / "src" / "kernbench" / "benches"
    / "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json"
_PAPER_FIG_DIR = (
    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
)

# The three variants to compare (coarse `primitive` intentionally dropped).
_ORDER = ("primitive_tiled", "composite", "composite_extended")
_LABELS = {
    "primitive_tiled":    "primitive hand-tiled\n(16×16×16)",
    "composite":          "composite\nGEMM",
    "composite_extended": "composite +\nsoftmax_merge",
}
_S_KV_TARGET = 131_072


def main() -> None:
    sweep = json.loads(_SWEEP_JSON.read_text())
    rows = {(r["variant"], r["S_kv"]): r for r in sweep["rows"]}

    dispatch_us = []
    engine_us = []
    totals_us = []
    for v in _ORDER:
        r = rows[(v, _S_KV_TARGET)]
        disp_ns = r["pe_cpu_dispatch_cycles"]
        total_ns = r["latency_ns"]
        eng_ns = max(0.0, total_ns - disp_ns)
        dispatch_us.append(disp_ns / 1e3)
        engine_us.append(eng_ns / 1e3)
        totals_us.append(total_ns / 1e3)

    xs = list(range(len(_ORDER)))
    labels = [_LABELS[v] for v in _ORDER]

    fig, ax = plt.subplots(figsize=(7.5, 5.0))

    bars_eng = ax.bar(
        xs, engine_us,
        color="#4f8a4f", edgecolor="#2a4a2a", label="engine (DMA + GEMM + MATH + IPCQ)",
    )
    bars_disp = ax.bar(
        xs, dispatch_us, bottom=engine_us,
        color="#c0504d", edgecolor="#5a2624", label="PE_CPU dispatch",
    )

    # Segment value labels — engine at mid, dispatch at mid of its stack.
    for i, (eng, disp, tot) in enumerate(zip(engine_us, dispatch_us, totals_us)):
        ax.text(i, eng / 2, f"{eng:.1f} µs",
                ha="center", va="center", fontsize=9, color="white")
        if disp / max(totals_us) > 0.04:  # only label if visible
            ax.text(i, eng + disp / 2, f"{disp:.1f} µs",
                    ha="center", va="center", fontsize=9, color="white")
        else:
            ax.annotate(f"{disp:.2f} µs",
                        xy=(i, tot), xytext=(0, 6),
                        textcoords="offset points",
                        ha="center", va="bottom", fontsize=8,
                        color="#5a2624")
        offset_pts = 14 if disp / max(totals_us) > 0.04 else 20
        ax.annotate(f"total {tot:.1f}",
                    xy=(i, tot), xytext=(0, offset_pts),
                    textcoords="offset points",
                    ha="center", va="bottom", fontsize=9,
                    color="#333", fontweight="bold")

    ax.set_xticks(xs)
    ax.set_xticklabels(labels, fontsize=10)
    ax.set_ylabel("time (µs)")
    ax.set_title(
        f"Case-6 decode latency breakdown at $S_{{kv}}=${_S_KV_TARGET // 1024}K\n"
        "(engine dominates all three — primitive hand-tiled adds "
        f"{dispatch_us[0]:.0f} µs PE_CPU dispatch overhead)",
        fontsize=11,
    )
    ax.legend(loc="upper right", fontsize=9)
    ax.grid(True, axis="y", ls=":", alpha=0.5)
    ax.set_ylim(0, max(totals_us) * 1.15)

    fig.text(
        0.5, 0.02,
        "First-order breakdown: dispatch and engine paths overlap partially on the real "
        "critical path;\ntreat the split as an upper bound on dispatch's contribution.",
        ha="center", fontsize=8, color="#666",
    )
    fig.tight_layout(rect=(0, 0.06, 1, 1))

    out = _FIG_DIR / "gqa_decode_long_ctx_composite_breakdown.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    print(f"wrote {out}")

    if _PAPER_FIG_DIR.is_dir():
        dst = _PAPER_FIG_DIR / out.name
        dst.write_bytes(out.read_bytes())
        print(f"copied {dst}")


if __name__ == "__main__":
    main()