forked from ccf-ai-infra/TileOPs-Metax
218 lines
8.3 KiB
Python
218 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Plot a roofline chart for a TileOPs op from benchmark logs + a GPU profile.
|
|
|
|
Roofline model (docs/design/roofline.md §1.2):
|
|
|
|
memory_time = bytes_moved / hbm_bandwidth
|
|
compute_time = total_flops / peak_flops
|
|
sol_time = max(memory_time, compute_time)
|
|
efficiency = sol_time / actual_time
|
|
|
|
Data sources (all repo mechanisms, no ad-hoc numbers):
|
|
|
|
- ``hbm_bandwidth``, ``peak_flops``: ``tileops.perf.load_profile(<gpu>)`` —
|
|
``theoretical x calibration`` stored in ``tileops/perf/profiles/<gpu>.yaml``.
|
|
- ``total_flops``, ``bytes_moved``: the manifest roofline formula, evaluated
|
|
per workload via ``op.eval_roofline()``.
|
|
- ``actual_time``: latencies parsed from the markdown reports written by the
|
|
op's benchmark (``benchmarks/ops/*``).
|
|
|
|
For the C500 the profile's HBM ``effective`` is the STREAM-Triad-measured
|
|
value (1.487 TB/s), so it already reflects the sGPU slice. Pass
|
|
``--compute-frac`` to scale the compute peak by the slice's Compute
|
|
percentage (mx-smi Sliced GPU section), e.g. ``--compute-frac 0.5``.
|
|
|
|
Usage::
|
|
|
|
python scripts/plot_roofline.py \\
|
|
--op QuantSwiGLUFwdChannelCastTransposeOp \\
|
|
--logs logs/2026-08-03_baseline_profile_run.log \\
|
|
logs/2026-08-03_opt3_profile_run.log \\
|
|
--gpu c500 --compute-frac 0.5 --out logs/roofline.png
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
import matplotlib
|
|
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import torch
|
|
|
|
from tileops.manifest import load_manifest, load_workloads
|
|
from tileops.perf import load_profile
|
|
|
|
# Dataviz-validated categorical palette (light mode).
|
|
_BASELINE_COLOR = "#2a78d6" # series 1 blue
|
|
_OPT_COLOR = "#eb6834" # series 2 orange
|
|
_LINE_COLOR = "#52514e" # neutral ink for the roofline ceilings
|
|
_RIDGE_COLOR = "#008300"
|
|
|
|
|
|
def _parse_log(path: Path) -> dict[tuple, float]:
|
|
"""Parse the ``### tileops`` latency table of a benchmark report.
|
|
|
|
Returns {(num_tokens, two_hidden, num_per_tokens, without_transpose,
|
|
round_sf): latency_ms}.
|
|
"""
|
|
rows: dict[tuple, float] = {}
|
|
in_tileops = False
|
|
for line in path.read_text().splitlines():
|
|
if line.startswith("### "):
|
|
in_tileops = line.startswith("### tileops")
|
|
continue
|
|
if not in_tileops or not line.startswith("| "):
|
|
continue
|
|
parts = [p.strip() for p in line.split("|")[1:-1]]
|
|
if len(parts) == 8 and parts[0].isdigit():
|
|
key = (int(parts[0]), int(parts[1]), int(parts[2]),
|
|
parts[3] == "True", parts[4] == "True")
|
|
rows[key] = float(parts[5])
|
|
return rows
|
|
|
|
|
|
def _load_op_class(op_name: str):
|
|
entry = load_manifest()[op_name]
|
|
mod_path, _, cls_name = entry["source"]["op"].replace(".py", "").partition(".")
|
|
# source.op is a repo-relative path; import as a tileops module.
|
|
mod_path = mod_path.replace("/", ".")
|
|
module = importlib.import_module(mod_path)
|
|
return getattr(module, op_name), entry
|
|
|
|
|
|
def _workload_rows(op_name: str, entry: dict) -> list[tuple[dict, int, int]]:
|
|
"""Yield (workload, flops, bytes) for every manifest workload.
|
|
|
|
flops/bytes come from the manifest roofline via ``op.eval_roofline()``,
|
|
binding the op's shape/dtype attributes from the workload before the call
|
|
(mirrors what ``forward()`` does before codegen's evaluator runs).
|
|
"""
|
|
op_cls = _load_op_class(op_name)[0]
|
|
sig = entry.get("signature", {})
|
|
input_names = list((sig.get("inputs") or {}).keys())
|
|
param_specs = sig.get("params") or {}
|
|
|
|
out = []
|
|
for w in load_workloads(op_name):
|
|
kwargs = {}
|
|
for pname, pattr in param_specs.items():
|
|
if pname in w:
|
|
kwargs[pname] = w[pname]
|
|
elif isinstance(pattr, dict) and "default" in pattr:
|
|
kwargs[pname] = pattr["default"]
|
|
op = op_cls(**kwargs)
|
|
x_shape = tuple(w["x_shape"])
|
|
for inp in input_names:
|
|
setattr(op, f"{inp}_shape", x_shape)
|
|
op.dtype = getattr(torch, w["dtypes"][0])
|
|
flops, nbytes = op.eval_roofline()
|
|
out.append((w, flops, nbytes))
|
|
return out
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(description=__doc__)
|
|
ap.add_argument("--op", required=True, help="manifest op name")
|
|
ap.add_argument("--logs", nargs="+", required=True,
|
|
help="benchmark markdown report(s), oldest first")
|
|
ap.add_argument("--gpu", default="c500", help="profile name in tileops/perf/profiles/")
|
|
ap.add_argument("--compute-frac", type=float, default=1.0,
|
|
help="scale compute peak by the sGPU Compute fraction")
|
|
ap.add_argument("--out", default="roofline.png")
|
|
args = ap.parse_args()
|
|
|
|
profile = load_profile(args.gpu)
|
|
bw = profile["hbm"]["effective"] # bytes/s (measured)
|
|
p_peak = profile["tensor_core"]["fp32"]["effective"] # flop/s (spec) x slice
|
|
p_peak *= args.compute_frac
|
|
|
|
_, entry = _load_op_class(args.op)
|
|
rows = _workload_rows(args.op, entry)
|
|
|
|
logs = [Path(p) for p in args.logs]
|
|
parsed = [_parse_log(p) for p in logs]
|
|
labels = [p.stem for p in logs]
|
|
|
|
# Collect points per series.
|
|
series: list[tuple[str, str, list[tuple[float, float, float]]]] = []
|
|
for label, lat_map in zip(labels, parsed, strict=True):
|
|
pts = []
|
|
for (w, flops, nbytes) in rows:
|
|
key = (w["x_shape"][0], w["x_shape"][1], w["num_per_tokens"],
|
|
w.get("without_transpose", False), w.get("round_sf", False))
|
|
if key not in lat_map:
|
|
continue
|
|
ai = flops / nbytes # flop/byte
|
|
perf = flops / (lat_map[key] * 1e-3) # flop/s
|
|
sol = max(nbytes / bw, flops / p_peak) # s
|
|
pts.append((ai, perf / 1e12, sol / (lat_map[key] * 1e-3)))
|
|
series.append((label, _BASELINE_COLOR if len(series) == 0 else _OPT_COLOR, pts))
|
|
|
|
# ---- plot ----
|
|
ai_ridge = p_peak / bw
|
|
all_ai = [ai for _, _, pts in series for ai, _, _ in pts]
|
|
x_lo = min(all_ai) * 0.5
|
|
x_hi = max(all_ai) * 4.0
|
|
xs = [x_lo, x_hi]
|
|
|
|
fig, ax = plt.subplots(figsize=(9, 6))
|
|
ax.set_xscale("log")
|
|
ax.set_yscale("log")
|
|
|
|
# Memory ceiling: y = BW * AI (1 flop/byte at 1 B/s = 1 flop/s)
|
|
ax.plot(xs, [bw * x / 1e12 for x in xs], color=_LINE_COLOR, lw=1.5,
|
|
label=f"memory ceiling (BW {bw/1e12:.2f} TB/s)")
|
|
# Compute ceiling: y = P_peak
|
|
ax.axhline(p_peak / 1e12, color=_LINE_COLOR, lw=1.5, ls="--",
|
|
label=f"compute ceiling (P {p_peak/1e12:.1f} TFLOPS)")
|
|
# Ridge point
|
|
ax.plot([ai_ridge], [p_peak / 1e12], "o", color=_RIDGE_COLOR, ms=6,
|
|
label=f"ridge AI {ai_ridge:.2f} flop/B")
|
|
|
|
# Measured points
|
|
markers = ["s", "o"]
|
|
for (label, color, pts), marker in zip(series, markers, strict=True):
|
|
ax.scatter([p[0] for p in pts], [p[1] for p in pts], marker=marker,
|
|
s=60, color=color, label=label, zorder=5)
|
|
|
|
ax.set_xlabel("arithmetic intensity (flop/byte)")
|
|
ax.set_ylabel("performance (TFLOP/s)")
|
|
ax.set_title(f"Roofline — {args.op} ({args.gpu})")
|
|
ax.grid(True, which="both", ls=":", alpha=0.35)
|
|
ax.legend(loc="lower right")
|
|
|
|
fig.tight_layout()
|
|
fig.savefig(args.out, dpi=150)
|
|
print(f"plot saved to {args.out}")
|
|
|
|
# ---- table ----
|
|
print(f"\nGPU={args.gpu} BW_eff={bw/1e12:.3f} TB/s P_eff={p_peak/1e12:.1f} TFLOPS "
|
|
f"(compute-frac={args.compute_frac})")
|
|
print(f"{'workload':34s} {'AI':>6s} {'perf_TF':>8s} {'base_eff':>8s} {'opt_eff':>8s}")
|
|
for (w, flops, nbytes) in rows:
|
|
key = (w["x_shape"][0], w["x_shape"][1], w["num_per_tokens"],
|
|
w.get("without_transpose", False), w.get("round_sf", False))
|
|
effs = []
|
|
for lat_map in parsed:
|
|
if key in lat_map:
|
|
effs.append((max(nbytes / bw, flops / p_peak)
|
|
/ (lat_map[key] * 1e-3)))
|
|
if not effs:
|
|
continue
|
|
ai = flops / nbytes
|
|
perf = flops / (parsed[0][key] * 1e-3) / 1e12
|
|
label = f"{key[0]}x{key[1]} npt={key[2]} wt={key[3]} rs={key[4]}"
|
|
eff_str = " ".join(f"{e:7.3f}" for e in effs)
|
|
print(f"{label:34s} {ai:6.2f} {perf:8.3f} {eff_str:>8s}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|