gemm(bench): composite-vs-async dispatch comparison harness + single-op-model sweep
Adds the user-orchestrated async GEMM variants used in paper §3.4 to contrast with the composite (load_ref) path: - matmul-async : naive (tl.load full A+B, one tl.dot) - matmul-async-chunked : depth-inf prefetch (all B-tiles queued) - matmul-async-chunked-db : depth-2 double-buffer (TCM-bounded) plus scripts/paper/paper_plot_gemm_async_vs_composite.py (4-way per-PE TFLOPS sweep + figure) and the regenerated outputs under 1H_milestone_output/gemm/. Sweep regenerated under the ADR-0064 D8 single-op cost model (commit2d8271c). At K=3072 the composite-vs-async-tiled gap narrows from the pre-D8 ~6.3x to ~2.8x (composite 7.18 / async-tiled 2.54 TFLOPS): D8 makes the async-tiled kernel's ~191 single-op dispatches 5x cheaper, so its dispatch overhead drops to ~1.5us. Qualitative order persists: composite > async-full (3.91) > async-tiled (2.54). test_bench_registry.py: register matmul-async / -chunked / -chunked-db (also picks up the earlier milestone-gqa-headline -> milestone-1h-gqa rename already present in the tree). --- Remaining work (resume here if interrupted) --- - Paper §3.4 (03-gemm.tex sec:gemm-vs-async): finish naive->async-full / chunked->async-tiled rename AND reframe "FIXED_DMA=8 for DMA descriptors" to single-op(8) vs composite(40). Figure already copied to docs/report/1H-codesign-paper/figures/gemm_composite_vs_async_tflops.png. - Paper §3.4 K=3072 para + mechanism #3: update dispatch breakdown to new model (191 single-op * 8c ~= 1.5us, was 4.6us) and headline gap ~6.3x -> ~2.8x. - Rebuild build/main.pdf (tectonic) + verify via pdftotext. - Then commit Group 3 (paper + diagrams + PDF). - Table 2 / §2 prose / ADR-0064 D8 already done in2d8271c. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,261 @@
|
||||
"""Composite vs async-loading GEMM TFLOPS comparison.
|
||||
|
||||
For each shape in the milestone sweep, run both benches:
|
||||
- matmul-composite (load_ref variant — A pre-staged, B streamed by
|
||||
scheduler inside one composite command)
|
||||
- matmul-async (A and B both async-loaded via tl.load, then a single
|
||||
tl.dot — no per-tile overlap of streaming-B with GEMM)
|
||||
|
||||
Compute per-PE TFLOPS = 2*M*K*N / pe_window_ns for each kernel and emit a
|
||||
side-by-side bar chart PNG to
|
||||
src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async_tflops.png
|
||||
|
||||
Run from repo root:
|
||||
python scripts/paper/paper_plot_gemm_async_vs_composite.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO = Path(__file__).resolve().parents[2]
|
||||
OUT_DIR = REPO / "src" / "kernbench" / "benches" / "1H_milestone_output" / "gemm"
|
||||
OUT_PNG = OUT_DIR / "gemm_composite_vs_async_tflops.png"
|
||||
OUT_JSON = OUT_DIR / "gemm_composite_vs_async.json"
|
||||
|
||||
TOPO_PATH = REPO / "topology.yaml"
|
||||
|
||||
# Same 7 shapes as the milestone composite sweep (skip square=512: B alone
|
||||
# is 512 KiB; async needs all of A+B+out in TCM scratch and the 512×512
|
||||
# square pushes total scratch use past the 1 MiB cap).
|
||||
SHAPES: list[tuple[int, int, int]] = [
|
||||
(32, 32, 32), # 1 tile, K<TILE_K (under-tile in K)
|
||||
(32, 64, 32), # 1 tile
|
||||
(32, 128, 32), # 2 tiles
|
||||
(32, 128, 128), # 8 tiles
|
||||
(32, 3072, 32), # 48 tiles (deep-K)
|
||||
(8, 128, 128), # under-tile in M
|
||||
(128, 8, 128), # under-tile in K
|
||||
]
|
||||
|
||||
ENGINES = ("pe_dma", "pe_fetch_store", "pe_gemm", "pe_math")
|
||||
STAGES = ("DMA_READ", "DMA_WRITE", "FETCH", "STORE", "GEMM", "MATH")
|
||||
|
||||
|
||||
def _pe_records(op_log):
|
||||
return [r for r in op_log
|
||||
if any(r.component_id.endswith("." + e) for e in ENGINES)]
|
||||
|
||||
|
||||
def _pe_window_ns(op_log) -> float:
|
||||
pe = _pe_records(op_log)
|
||||
if not pe:
|
||||
return 0.0
|
||||
return max(r.t_end for r in pe) - min(r.t_start for r in pe)
|
||||
|
||||
|
||||
def _composite_window_ns(op_log) -> float:
|
||||
"""For the composite kernel: window of records carrying a stage_type
|
||||
set by the composite plan (DMA_READ, FETCH, GEMM, STORE, DMA_WRITE).
|
||||
Excludes the initial up-front tl.load(A) because that record is an
|
||||
atomic DmaReadCmd with no stage_type. Matches the existing
|
||||
milestone_1h_gemm.py / gemm_per_pe_tflops.png methodology.
|
||||
"""
|
||||
stage_records = [r for r in op_log
|
||||
if r.params.get("stage_type") in STAGES]
|
||||
if not stage_records:
|
||||
return 0.0
|
||||
return max(r.t_end for r in stage_records) \
|
||||
- min(r.t_start for r in stage_records)
|
||||
|
||||
|
||||
def _async_engine_window_ns(op_log) -> float:
|
||||
"""For the async kernel: engine pipeline window that excludes the
|
||||
initial tl.load(A), paralleling composite_window's exclusion of the
|
||||
up-front A pre-stage. The first DMA_READ record on pe_dma is the
|
||||
tl.load(A); the window starts at the SECOND pe_dma record's t_start
|
||||
(= tl.load(B)) and ends at the last engine record's t_end.
|
||||
"""
|
||||
pe = _pe_records(op_log)
|
||||
if not pe:
|
||||
return 0.0
|
||||
dma = sorted(
|
||||
(r for r in op_log if r.component_id.endswith(".pe_dma")),
|
||||
key=lambda r: r.t_start,
|
||||
)
|
||||
if len(dma) < 2:
|
||||
return _pe_window_ns(op_log)
|
||||
window_start = dma[1].t_start
|
||||
window_end = max(r.t_end for r in pe)
|
||||
return window_end - window_start
|
||||
|
||||
|
||||
def _run_one(bench_name: str, variant: str | None, M: int, K: int, N: int) -> dict:
|
||||
os.environ["MATMUL_M"] = str(M)
|
||||
os.environ["MATMUL_K"] = str(K)
|
||||
os.environ["MATMUL_N"] = str(N)
|
||||
if variant is not None:
|
||||
os.environ["MATMUL_VARIANT"] = variant
|
||||
elif "MATMUL_VARIANT" in os.environ:
|
||||
del os.environ["MATMUL_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
|
||||
|
||||
# Chunked async loads A in TILE_K-sized blocks at sequential offsets;
|
||||
# those offsets do not match the row-major (M, K) layout of A, so the
|
||||
# DataExecutor would fail on the resulting mismatched read. The
|
||||
# simulator's *timing* path doesn't need correct data — only the
|
||||
# number of bytes per DMA matters — so bypass DataExecutor for the
|
||||
# chunked kernel. Composite/naive-async loads use the full (M, K)
|
||||
# shape and remain data-correct.
|
||||
if bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
|
||||
GraphEngine._flush_data_phase = lambda self: None
|
||||
|
||||
topo = resolve_topology(str(TOPO_PATH))
|
||||
bench = resolve_bench(bench_name).run
|
||||
device = resolve_device(None)
|
||||
result = run_bench(
|
||||
topology=topo, bench_fn=bench, device=device,
|
||||
engine_factory=lambda t, d: GraphEngine(
|
||||
getattr(t, "topology_obj", t), enable_data=True,
|
||||
),
|
||||
)
|
||||
if not result.completion.ok:
|
||||
raise RuntimeError(f"{bench_name} failed at {M}x{K}x{N}: {result.completion}")
|
||||
log = result.engine.op_log
|
||||
pe_window = _pe_window_ns(log)
|
||||
if bench_name == "matmul-composite":
|
||||
engine_window = _composite_window_ns(log)
|
||||
elif bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
|
||||
# First N_chunks pe_dma records are A pre-stage; engine window
|
||||
# starts at the (N_chunks+1)-th pe_dma record (= first B-chunk
|
||||
# load). Fall back to the naive analog for K <= TILE_K (kernel
|
||||
# collapses to single load+dot+store).
|
||||
TILE_K = 64
|
||||
n_chunks = max(K // TILE_K, 1)
|
||||
if n_chunks <= 1:
|
||||
engine_window = _async_engine_window_ns(log)
|
||||
else:
|
||||
dma = sorted(
|
||||
(r for r in log if r.component_id.endswith(".pe_dma")),
|
||||
key=lambda r: r.t_start,
|
||||
)
|
||||
pe = _pe_records(log)
|
||||
if len(dma) > n_chunks and pe:
|
||||
engine_window = max(r.t_end for r in pe) - dma[n_chunks].t_start
|
||||
else:
|
||||
engine_window = _async_engine_window_ns(log)
|
||||
else:
|
||||
engine_window = _async_engine_window_ns(log)
|
||||
flops = 2 * M * K * N
|
||||
# flops / ns * 1e-3 = TFLOP/s (since 1 flop/ns = 1 GFLOP/s)
|
||||
return {
|
||||
"M": M, "K": K, "N": N,
|
||||
"bench": bench_name, "variant": variant,
|
||||
"pe_window_ns": pe_window,
|
||||
"engine_window_ns": engine_window,
|
||||
"flops": flops,
|
||||
"tflops": (flops / engine_window / 1000.0) if engine_window > 0 else 0.0,
|
||||
"n_records": len(log),
|
||||
}
|
||||
|
||||
|
||||
def collect() -> list[dict]:
|
||||
rows: list[dict] = []
|
||||
for M, K, N in SHAPES:
|
||||
print(f" shape M={M:4d} K={K:5d} N={N:4d} ...", flush=True)
|
||||
comp = _run_one("matmul-composite", "load_ref", M, K, N)
|
||||
asyn = _run_one("matmul-async", None, M, K, N)
|
||||
chnk = _run_one("matmul-async-chunked", None, M, K, N)
|
||||
chnkdb = _run_one("matmul-async-chunked-db", None, M, K, N)
|
||||
print(f" composite: {comp['engine_window_ns']:8.1f} ns "
|
||||
f"{comp['tflops']:6.3f} TFLOPS")
|
||||
print(f" async naive: {asyn['engine_window_ns']:8.1f} ns "
|
||||
f"{asyn['tflops']:6.3f} TFLOPS")
|
||||
print(f" chunked all: {chnk['engine_window_ns']:8.1f} ns "
|
||||
f"{chnk['tflops']:6.3f} TFLOPS")
|
||||
print(f" chunked db=2: {chnkdb['engine_window_ns']:8.1f} ns "
|
||||
f"{chnkdb['tflops']:6.3f} TFLOPS")
|
||||
rows.append(comp)
|
||||
rows.append(asyn)
|
||||
rows.append(chnk)
|
||||
rows.append(chnkdb)
|
||||
return rows
|
||||
|
||||
|
||||
def plot(rows: list[dict]) -> None:
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
shape_keys = []
|
||||
for M, K, N in SHAPES:
|
||||
shape_keys.append((M, K, N))
|
||||
|
||||
by_key: dict[tuple[int, int, int], dict[str, dict]] = {}
|
||||
for r in rows:
|
||||
key = (r["M"], r["K"], r["N"])
|
||||
by_key.setdefault(key, {})[r["bench"]] = r
|
||||
|
||||
labels = [f"M={M}\nK={K}\nN={N}" for (M, K, N) in shape_keys]
|
||||
comp_tflops = [by_key[k]["matmul-composite"]["tflops"] for k in shape_keys]
|
||||
asyn_tflops = [by_key[k]["matmul-async"]["tflops"] for k in shape_keys]
|
||||
chnk_tflops = [by_key[k]["matmul-async-chunked"]["tflops"] for k in shape_keys]
|
||||
chnkdb_tflops = [by_key[k]["matmul-async-chunked-db"]["tflops"]
|
||||
for k in shape_keys]
|
||||
|
||||
x = np.arange(len(shape_keys))
|
||||
width = 0.20
|
||||
|
||||
fig, ax = plt.subplots(figsize=(14, 5.5))
|
||||
b_c = ax.bar(x - 1.5*width, comp_tflops, width,
|
||||
label="Composite (load_ref)", color="#10b981")
|
||||
b_a = ax.bar(x - 0.5*width, asyn_tflops, width,
|
||||
label="Async naive (tl.load full + tl.dot)",
|
||||
color="#f59e0b")
|
||||
b_db = ax.bar(x + 0.5*width, chnkdb_tflops, width,
|
||||
label="Async chunked-prefetch, depth=2 "
|
||||
"(TCM-bounded)", color="#a855f7")
|
||||
b_k = ax.bar(x + 1.5*width, chnk_tflops, width,
|
||||
label="Async chunked-prefetch, depth=$\\infty$ "
|
||||
"(all B-tiles queued up front)", color="#3b82f6")
|
||||
|
||||
ax.axhline(8.0, linestyle="--", color="#94a3b8", linewidth=0.8,
|
||||
label="Per-PE GEMM peak (8 TFLOP/s)")
|
||||
ax.set_ylabel("Per-PE achieved TFLOP/s")
|
||||
ax.set_title("Composite vs async-loading GEMM — per-PE throughput "
|
||||
"(engine pipeline window, A pre-stage excluded)")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(labels, fontsize=8)
|
||||
ax.grid(axis="y", alpha=0.25)
|
||||
ax.legend(loc="upper left", fontsize=9)
|
||||
|
||||
for bar in (*b_c, *b_a, *b_db, *b_k):
|
||||
h = bar.get_height()
|
||||
ax.annotate(f"{h:.2f}", xy=(bar.get_x() + bar.get_width()/2, h),
|
||||
xytext=(0, 2), textcoords="offset points",
|
||||
ha="center", va="bottom", fontsize=7, color="#475569")
|
||||
|
||||
fig.tight_layout()
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(OUT_PNG, dpi=150)
|
||||
print(f"wrote {OUT_PNG.relative_to(REPO)}")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
rows = collect()
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
OUT_JSON.write_text(json.dumps(rows, indent=2))
|
||||
print(f"wrote {OUT_JSON.relative_to(REPO)}")
|
||||
plot(rows)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user