kernbench2/scripts/paper/paper_plot_gqa.py

"""Report harness: render GQA figures from ``gqa_latency.json``.

Isolated under ``scripts/paper/`` for the 1H codesign report only.
Reads the JSON emitted by ``paper_gqa_latency.py`` and writes two PNGs into
``docs/report/1H-codesign-paper/figures/``:

  gqa_latency_by_panel.png      end-to-end latency per panel
  gqa_op_engine_breakdown.png   op-counts + engine occupancy per panel

Run (after paper_gqa_latency.py):
    python scripts/paper/paper_plot_gqa.py
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

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

_FIG_DIR = Path(__file__).resolve().parents[2] / "docs" / "report" / "1H-codesign-paper" / "figures"
_IN_JSON = _FIG_DIR / "gqa_latency.json"

_LABELS = {
    "single_kv_group_prefill_gqa_c8_p8":
        "prefill\nC=8, P=8\n(single KV group)",
}


def _load() -> list[dict]:
    return json.loads(_IN_JSON.read_text())["rows"]


def _plot_latency(rows: list[dict]) -> Path:
    labels = [_LABELS[r["panel"]] for r in rows]
    lat = [r["latency_ns"] for r in rows]
    fig, ax = plt.subplots(figsize=(7.0, 4.0))
    bars = ax.bar(labels, lat, color="#3b6ea5", width=0.6)
    ax.set_ylabel("end-to-end latency (ns)")
    ax.set_title("Fused GQA — end-to-end latency per panel")
    ax.bar_label(bars, fmt="%.0f", padding=3, fontsize=9)
    ax.grid(axis="y", ls=":", alpha=0.5)
    ax.set_ylim(0, max(lat) * 1.15)
    fig.tight_layout()
    out = _FIG_DIR / "gqa_latency_by_panel.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    return out


def _plot_breakdown(rows: list[dict]) -> Path:
    labels = [_LABELS[r["panel"]] for r in rows]
    x = range(len(rows))
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11.0, 4.2))

    # Left: op-count breakdown (grouped bars).
    keys = ["gemm_count", "ipcq_copy_count", "dma_read_count", "dma_write_count"]
    disp = ["GEMM", "IPCQ copy", "DMA read", "DMA write"]
    colors = ["#3b6ea5", "#c0504d", "#9bbb59", "#8064a2"]
    w = 0.2
    for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
        vals = [r["op_log_summary"][k] for r in rows]
        ax1.bar([xi + (i - 1.5) * w for xi in x], vals, width=w, label=d, color=c)
    ax1.set_xticks(list(x))
    ax1.set_xticklabels(labels, fontsize=8)
    ax1.set_ylabel("op count")
    ax1.set_title("Per-panel op-count breakdown")
    ax1.legend(fontsize=8)
    ax1.grid(axis="y", ls=":", alpha=0.5)

    # Right: engine occupancy (log scale — compute vs data movement).
    engs = ["pe_gemm", "pe_math", "pe_dma"]
    edisp = ["GEMM engine", "MATH engine", "DMA engine"]
    ecolors = ["#3b6ea5", "#e8a33d", "#c0504d"]
    w2 = 0.25
    for i, (e, d, c) in enumerate(zip(engs, edisp, ecolors)):
        vals = [max(r["engine_occupancy_ns"][e], 0.1) for r in rows]
        ax2.bar([xi + (i - 1) * w2 for xi in x], vals, width=w2, label=d, color=c)
    ax2.set_yscale("log")
    ax2.set_xticks(list(x))
    ax2.set_xticklabels(labels, fontsize=8)
    ax2.set_ylabel("summed engine occupancy (ns, log)")
    ax2.set_title("Compute vs data-movement occupancy")
    ax2.legend(fontsize=8)
    ax2.grid(axis="y", ls=":", alpha=0.5)

    fig.suptitle("Fused GQA — where the work goes", fontsize=12)
    fig.tight_layout()
    out = _FIG_DIR / "gqa_op_engine_breakdown.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    return out


def main() -> None:
    rows = _load()
    p1 = _plot_latency(rows)
    p2 = _plot_breakdown(rows)
    print(f"wrote {p1}")
    print(f"wrote {p2}")


if __name__ == "__main__":
    main()