plot_gpu_perf.py

import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib import font_manager
from matplotlib.patches import Patch


DEFAULT_RECORDS = Path("data/gpu_perf_records.jsonl")
DEFAULT_ARCH = Path("data/arch.json")
DEFAULT_OUTPUT_DIR = Path("figures")

ARCH_COLORS = {
    "Pascal": "#7f7f7f",
    "Volta": "#1f77b4",
    "Turing": "#17becf",
    "Ampere": "#2ca02c",
    "Ada Lovelace": "#ff7f0e",
    "Hopper": "#d62728",
    "Blackwell": "#9467bd",
    "Unknown": "#8c564b",
}

PLOT_ORDER = [
    ("ResNet50", "float16"),
    ("ResNet50", "float32"),
    ("ViT Transformer", "float16"),
    ("ViT Transformer", "float32"),
]

GENERATION_CLUSTER_ORDER = ["30xx", "40xx", "50xx"]
GENERATION_GPU_ORDER = {
    "3060": 0,
    "3080": 1,
    "3080_Ti": 2,
    "3090": 3,
    "3090_Ti": 4,
    "4080": 0,
    "vGPU_32GB": 1,
    "4090": 2,
    "vGPU_48GB": 3,
    "4090_D": 4,
    "5090": 0,
    "5090D": 1,
}

FONT_CANDIDATES = [
    "Arial Unicode MS",
    "PingFang SC",
    "Hiragino Sans GB",
    "Heiti SC",
    "STHeiti",
]


def configure_fonts() -> None:
    available: set[str] = set()
    for path in font_manager.findSystemFonts():
        try:
            available.add(font_manager.FontProperties(fname=path).get_name())
        except RuntimeError:
            continue
    for name in FONT_CANDIDATES:
        if name in available:
            plt.rcParams["font.family"] = name
            break
    plt.rcParams["axes.unicode_minus"] = False


def display_gpu_id(gpu_id: str, max_line_length: int | None = None) -> str:
    label = gpu_id.replace("_", " ")
    if max_line_length is None or len(label) <= max_line_length:
        return label

    lines: list[str] = []
    current = ""
    for word in label.split():
        candidate = word if not current else f"{current} {word}"
        if len(candidate) <= max_line_length:
            current = candidate
        else:
            if current:
                lines.append(current)
            current = word
    if current:
        lines.append(current)
    return "\n".join(lines)


def configure_gpu_tick_labels(ax: plt.Axes) -> None:
    ax.tick_params(axis="x", labelrotation=75, labelsize=8)
    for label in ax.get_xticklabels():
        label.set_horizontalalignment("right")
        label.set_multialignment("left")
        label.set_rotation_mode("anchor")


def load_records(path: Path) -> pd.DataFrame:
    return pd.read_json(path, lines=True)


def load_architecture_map(path: Path) -> dict[str, str]:
    data = json.loads(path.read_text(encoding="utf-8"))
    mapping: dict[str, str] = {}
    for architecture, items in data["architectures"].items():
        for item in items:
            mapping[item["gpu_id"]] = architecture
    return mapping


def add_architecture(df: pd.DataFrame, architecture_map: dict[str, str]) -> pd.DataFrame:
    result = df.copy()
    result["architecture"] = result["gpu_id"].map(architecture_map).fillna("Unknown")
    return result


def aggregate_perf(df: pd.DataFrame) -> pd.DataFrame:
    measured = df[df["iteration"] > 0]
    return (
        measured.groupby(["gpu_id", "architecture", "benchmark", "precision"], as_index=False)
        .agg(
            mean_images_per_second=("images_per_second", "mean"),
            std_images_per_second=("images_per_second", "std"),
        )
        .sort_values(["benchmark", "precision", "mean_images_per_second"], ascending=[True, True, False])
    )


def compute_speedup(agg: pd.DataFrame) -> pd.DataFrame:
    pivot = agg.pivot_table(
        index=["gpu_id", "architecture", "benchmark"],
        columns="precision",
        values="mean_images_per_second",
        aggfunc="first",
    ).reset_index()
    pivot["fp16_over_fp32"] = pivot["float16"] / pivot["float32"]
    return pivot.dropna(subset=["fp16_over_fp32"])


def assign_generation_cluster(gpu_id: str) -> str | None:
    compact = gpu_id.replace("_", "")
    if compact.startswith("30"):
        return "30xx"
    if compact.startswith("40") or compact.upper().startswith("VGPU"):
        return "40xx"
    if compact.startswith("50"):
        return "50xx"
    return None


def cluster_generation_perf(agg: pd.DataFrame) -> pd.DataFrame:
    result = agg.copy()
    result["generation_cluster"] = result["gpu_id"].map(assign_generation_cluster)
    result = result.dropna(subset=["generation_cluster"])
    result["generation_cluster"] = pd.Categorical(
        result["generation_cluster"],
        categories=GENERATION_CLUSTER_ORDER,
        ordered=True,
    )
    result["generation_order"] = result["gpu_id"].map(GENERATION_GPU_ORDER).fillna(999).astype(int)
    return result.sort_values(["generation_cluster", "generation_order", "gpu_id"]).drop(columns=["generation_order"])


def architecture_legend(df: pd.DataFrame) -> list[Patch]:
    architectures = [arch for arch in ARCH_COLORS if arch in set(df["architecture"])]
    return [Patch(facecolor=ARCH_COLORS[arch], label=arch) for arch in architectures]


def save_perf_rank(agg: pd.DataFrame, output_dir: Path) -> Path:
    fig, axes = plt.subplots(2, 2, figsize=(18, 11))

    for ax, (benchmark, precision) in zip(axes.flat, PLOT_ORDER):
        data = agg[(agg["benchmark"] == benchmark) & (agg["precision"] == precision)].copy()
        data = data.sort_values("mean_images_per_second", ascending=False)
        colors = [ARCH_COLORS.get(arch, ARCH_COLORS["Unknown"]) for arch in data["architecture"]]
        yerr = data["std_images_per_second"].fillna(0.0)
        labels = [display_gpu_id(gpu_id, max_line_length=14) for gpu_id in data["gpu_id"]]
        ax.bar(
            labels,
            data["mean_images_per_second"],
            yerr=yerr,
            color=colors,
            width=0.82,
            capsize=2.5,
            error_kw={"elinewidth": 0.9, "ecolor": "#333333", "alpha": 0.75},
        )
        ax.set_title(f"{benchmark} {precision.upper()} throughput")
        ax.set_ylabel("images/s, mean of iter 1-4")
        configure_gpu_tick_labels(ax)
        ax.grid(axis="y", alpha=0.25)

    fig.legend(
        handles=architecture_legend(agg),
        loc="lower center",
        ncol=4,
        frameon=False,
        bbox_to_anchor=(0.5, 0.01),
    )
    fig.tight_layout(rect=(0, 0.06, 1, 1))
    output = output_dir / "perf_rank_2x2.png"
    fig.savefig(output, dpi=180, bbox_inches="tight")
    plt.close(fig)
    return output


def save_speedup(speedup: pd.DataFrame, output_dir: Path) -> Path:
    benchmarks = list(speedup["benchmark"].drop_duplicates())
    fig, axes = plt.subplots(len(benchmarks), 1, figsize=(16, 9))
    if len(benchmarks) == 1:
        axes = [axes]

    for ax, benchmark in zip(axes, benchmarks):
        data = speedup[speedup["benchmark"] == benchmark].sort_values("fp16_over_fp32", ascending=False)
        colors = [ARCH_COLORS.get(arch, ARCH_COLORS["Unknown"]) for arch in data["architecture"]]
        labels = [display_gpu_id(gpu_id, max_line_length=14) for gpu_id in data["gpu_id"]]
        ax.bar(labels, data["fp16_over_fp32"], color=colors, width=0.82)
        ax.axhline(1.0, color="#444444", linewidth=1, linestyle="--")
        ax.set_title(f"{benchmark} FP16 / FP32 speedup")
        ax.set_ylabel("speedup ratio")
        configure_gpu_tick_labels(ax)
        ax.grid(axis="y", alpha=0.25)

    fig.legend(
        handles=architecture_legend(speedup),
        loc="lower center",
        ncol=4,
        frameon=False,
        bbox_to_anchor=(0.5, 0.01),
    )
    fig.tight_layout(rect=(0, 0.08, 1, 1))
    output = output_dir / "fp16_speedup.png"
    fig.savefig(output, dpi=180, bbox_inches="tight")
    plt.close(fig)
    return output


def save_arch_distribution(agg: pd.DataFrame, output_dir: Path) -> Path:
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    arch_order = [arch for arch in ARCH_COLORS if arch in set(agg["architecture"])]
    arch_index = {arch: index for index, arch in enumerate(arch_order)}

    for ax, (benchmark, precision) in zip(axes.flat, PLOT_ORDER):
        data = agg[(agg["benchmark"] == benchmark) & (agg["precision"] == precision)].copy()
        data["x"] = data["architecture"].map(arch_index)
        for arch, group in data.groupby("architecture"):
            x = group["x"]
            jitter = pd.Series(range(len(group)), index=group.index).map(lambda i: ((i % 5) - 2) * 0.045)
            ax.scatter(
                x + jitter,
                group["mean_images_per_second"],
                color=ARCH_COLORS.get(arch, ARCH_COLORS["Unknown"]),
                s=44,
                alpha=0.82,
                edgecolor="white",
                linewidth=0.45,
            )
            median = group["mean_images_per_second"].median()
            ax.hlines(median, arch_index[arch] - 0.23, arch_index[arch] + 0.23, color="#222222", linewidth=2)

        ax.set_title(f"{benchmark} {precision.upper()} by architecture")
        ax.set_ylabel("images/s, mean of iter 1-4")
        ax.set_xticks(range(len(arch_order)))
        ax.set_xticklabels(arch_order, rotation=35, ha="right")
        ax.grid(axis="y", alpha=0.25)

    fig.legend(
        handles=architecture_legend(agg),
        loc="lower center",
        ncol=4,
        frameon=False,
        bbox_to_anchor=(0.5, 0.01),
    )
    fig.tight_layout(rect=(0, 0.07, 1, 1))
    output = output_dir / "arch_distribution.png"
    fig.savefig(output, dpi=180, bbox_inches="tight")
    plt.close(fig)
    return output


def save_generation_cluster(clustered: pd.DataFrame, output_dir: Path) -> Path:
    fig, axes = plt.subplots(2, 2, figsize=(16, 10))

    for ax, (benchmark, precision) in zip(axes.flat, PLOT_ORDER):
        data = clustered[(clustered["benchmark"] == benchmark) & (clustered["precision"] == precision)].copy()
        labels = [display_gpu_id(gpu_id, max_line_length=14) for gpu_id in data["gpu_id"]]
        colors = [ARCH_COLORS.get(arch, ARCH_COLORS["Unknown"]) for arch in data["architecture"]]
        yerr = data["std_images_per_second"].fillna(0.0)

        positions: list[float] = []
        tick_labels: list[str] = []
        group_centers: list[float] = []
        group_labels: list[str] = []
        cursor = 0.0
        for cluster in GENERATION_CLUSTER_ORDER:
            group = data[data["generation_cluster"] == cluster]
            if group.empty:
                continue
            group_positions = [cursor + i for i in range(len(group))]
            positions.extend(group_positions)
            tick_labels.extend(labels[data.index.get_loc(idx)] for idx in group.index)
            group_centers.append((group_positions[0] + group_positions[-1]) / 2)
            group_labels.append(cluster)
            cursor = group_positions[-1] + 1.8

        ax.bar(
            positions,
            data["mean_images_per_second"],
            yerr=yerr,
            color=colors,
            width=0.72,
            capsize=2.5,
            error_kw={"elinewidth": 0.9, "ecolor": "#333333", "alpha": 0.75},
        )
        ax.set_title(f"{benchmark} {precision.upper()} cluster comparison")
        ax.set_ylabel("images/s, mean of iter 1-4")
        ax.set_xticks(positions)
        ax.set_xticklabels(tick_labels)
        configure_gpu_tick_labels(ax)
        ax.grid(axis="y", alpha=0.25)

        _, top = ax.get_ylim()
        ax.set_ylim(top=top * 1.08)
        top = ax.get_ylim()[1]
        for center, label in zip(group_centers, group_labels):
            ax.text(center, top * 0.96, label, ha="center", va="top", fontsize=10, fontweight="bold")

    fig.legend(
        handles=architecture_legend(clustered),
        loc="lower center",
        ncol=4,
        frameon=False,
        bbox_to_anchor=(0.5, 0.01),
    )
    fig.tight_layout(rect=(0, 0.07, 1, 1))
    output = output_dir / "cluster_30_40_50.png"
    fig.savefig(output, dpi=180, bbox_inches="tight")
    plt.close(fig)
    return output


def generate_figures(records_path: Path, arch_path: Path, output_dir: Path) -> list[Path]:
    configure_fonts()
    output_dir.mkdir(parents=True, exist_ok=True)
    df = add_architecture(load_records(records_path), load_architecture_map(arch_path))
    agg = aggregate_perf(df)
    speedup = compute_speedup(agg)
    clustered = cluster_generation_perf(agg)
    return [
        save_perf_rank(agg, output_dir),
        save_speedup(speedup, output_dir),
        save_arch_distribution(agg, output_dir),
        save_generation_cluster(clustered, output_dir),
    ]


def main() -> None:
    parser = argparse.ArgumentParser(description="Plot GPU performance analysis figures.")
    parser.add_argument("--records", default=str(DEFAULT_RECORDS), help="Parsed GPU performance JSONL.")
    parser.add_argument("--arch", default=str(DEFAULT_ARCH), help="Architecture classification JSON.")
    parser.add_argument("--output-dir", default=str(DEFAULT_OUTPUT_DIR), help="Figure output directory.")
    args = parser.parse_args()

    for path in generate_figures(Path(args.records), Path(args.arch), Path(args.output_dir)):
        print(path)


if __name__ == "__main__":
    main()