diff --git a/scripts/paper/paper_plot_gemm_async_vs_composite.py b/scripts/paper/paper_plot_gemm_async_vs_composite.py new file mode 100644 index 0000000..94d3593 --- /dev/null +++ b/scripts/paper/paper_plot_gemm_async_vs_composite.py @@ -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 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()) diff --git a/src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async.json b/src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async.json new file mode 100644 index 0000000..e29adc4 --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async.json @@ -0,0 +1,338 @@ +[ + { + "M": 32, + "K": 32, + "N": 32, + "bench": "matmul-composite", + "variant": "load_ref", + "pe_window_ns": 171.394, + "engine_window_ns": 106.38400000000001, + "flops": 65536, + "tflops": 0.6160324860881335, + "n_records": 7 + }, + { + "M": 32, + "K": 32, + "N": 32, + "bench": "matmul-async", + "variant": null, + "pe_window_ns": 103.22199999999998, + "engine_window_ns": 71.71199999999999, + "flops": 65536, + "tflops": 0.9138777331548418, + "n_records": 4 + }, + { + "M": 32, + "K": 32, + "N": 32, + "bench": "matmul-async-chunked", + "variant": null, + "pe_window_ns": 103.22199999999998, + "engine_window_ns": 71.71199999999999, + "flops": 65536, + "tflops": 0.9138777331548418, + "n_records": 4 + }, + { + "M": 32, + "K": 32, + "N": 32, + "bench": "matmul-async-chunked-db", + "variant": null, + "pe_window_ns": 103.22199999999998, + "engine_window_ns": 71.71199999999999, + "flops": 65536, + "tflops": 0.9138777331548418, + "n_records": 4 + }, + { + "M": 32, + "K": 64, + "N": 32, + "bench": "matmul-composite", + "variant": "load_ref", + "pe_window_ns": 179.394, + "engine_window_ns": 106.38400000000001, + "flops": 131072, + "tflops": 1.232064972176267, + "n_records": 7 + }, + { + "M": 32, + "K": 64, + "N": 32, + "bench": "matmul-async", + "variant": null, + "pe_window_ns": 127.41399999999999, + "engine_window_ns": 87.904, + "flops": 131072, + "tflops": 1.4910811794685113, + "n_records": 4 + }, + { + "M": 32, + "K": 64, + "N": 32, + "bench": "matmul-async-chunked", + "variant": null, + "pe_window_ns": 127.41399999999999, + "engine_window_ns": 87.904, + "flops": 131072, + "tflops": 1.4910811794685113, + "n_records": 4 + }, + { + "M": 32, + "K": 64, + "N": 32, + "bench": "matmul-async-chunked-db", + "variant": null, + "pe_window_ns": 127.41399999999999, + "engine_window_ns": 87.904, + "flops": 131072, + "tflops": 1.4910811794685113, + "n_records": 4 + }, + { + "M": 32, + "K": 128, + "N": 32, + "bench": "matmul-composite", + "variant": "load_ref", + "pe_window_ns": 211.77800000000002, + "engine_window_ns": 122.76800000000003, + "flops": 262144, + "tflops": 2.135279551674703, + "n_records": 10 + }, + { + "M": 32, + "K": 128, + "N": 32, + "bench": "matmul-async", + "variant": null, + "pe_window_ns": 175.798, + "engine_window_ns": 120.28800000000001, + "flops": 262144, + "tflops": 2.1793030061186482, + "n_records": 4 + }, + { + "M": 32, + "K": 128, + "N": 32, + "bench": "matmul-async-chunked", + "variant": null, + "pe_window_ns": 185.81799999999998, + "engine_window_ns": 130.308, + "flops": 262144, + "tflops": 2.011726064401265, + "n_records": 8 + }, + { + "M": 32, + "K": 128, + "N": 32, + "bench": "matmul-async-chunked-db", + "variant": null, + "pe_window_ns": 185.81799999999998, + "engine_window_ns": 130.308, + "flops": 262144, + "tflops": 2.011726064401265, + "n_records": 8 + }, + { + "M": 32, + "K": 128, + "N": 128, + "bench": "matmul-composite", + "variant": "load_ref", + "pe_window_ns": 315.394, + "engine_window_ns": 226.38400000000001, + "flops": 1048576, + "tflops": 4.631846773623577, + "n_records": 34 + }, + { + "M": 32, + "K": 128, + "N": 128, + "bench": "matmul-async", + "variant": null, + "pe_window_ns": 394.102, + "engine_window_ns": 338.592, + "flops": 1048576, + "tflops": 3.096871751252245, + "n_records": 4 + }, + { + "M": 32, + "K": 128, + "N": 128, + "bench": "matmul-async-chunked", + "variant": null, + "pe_window_ns": 368.12200000000007, + "engine_window_ns": 312.6120000000001, + "flops": 1048576, + "tflops": 3.354241040011259, + "n_records": 8 + }, + { + "M": 32, + "K": 128, + "N": 128, + "bench": "matmul-async-chunked-db", + "variant": null, + "pe_window_ns": 368.12200000000007, + "engine_window_ns": 312.6120000000001, + "flops": 1048576, + "tflops": 3.354241040011259, + "n_records": 8 + }, + { + "M": 32, + "K": 3072, + "N": 32, + "bench": "matmul-composite", + "variant": "load_ref", + "pe_window_ns": 1701.4420000000007, + "engine_window_ns": 876.4320000000005, + "flops": 6291456, + "tflops": 7.178487321320988, + "n_records": 148 + }, + { + "M": 32, + "K": 3072, + "N": 32, + "bench": "matmul-async", + "variant": null, + "pe_window_ns": 2401.4620000000004, + "engine_window_ns": 1609.952, + "flops": 6291456, + "tflops": 3.9078531533859397, + "n_records": 4 + }, + { + "M": 32, + "K": 3072, + "N": 32, + "bench": "matmul-async-chunked", + "variant": null, + "pe_window_ns": 3264.8920000000217, + "engine_window_ns": 2473.3820000000214, + "flops": 6291456, + "tflops": 2.5436653133240017, + "n_records": 192 + }, + { + "M": 32, + "K": 3072, + "N": 32, + "bench": "matmul-async-chunked-db", + "variant": null, + "pe_window_ns": 3286.402000000022, + "engine_window_ns": 2494.8920000000217, + "flops": 6291456, + "tflops": 2.5217348085608298, + "n_records": 192 + }, + { + "M": 8, + "K": 128, + "N": 128, + "bench": "matmul-composite", + "variant": "load_ref", + "pe_window_ns": 291.394, + "engine_window_ns": 226.38400000000001, + "flops": 262144, + "tflops": 1.1579616934058943, + "n_records": 34 + }, + { + "M": 8, + "K": 128, + "N": 128, + "bench": "matmul-async", + "variant": null, + "pe_window_ns": 247.798, + "engine_window_ns": 216.288, + "flops": 262144, + "tflops": 1.212013611480988, + "n_records": 4 + }, + { + "M": 8, + "K": 128, + "N": 128, + "bench": "matmul-async-chunked", + "variant": null, + "pe_window_ns": 251.42399999999998, + "engine_window_ns": 214.914, + "flops": 262144, + "tflops": 1.2197623235340647, + "n_records": 8 + }, + { + "M": 8, + "K": 128, + "N": 128, + "bench": "matmul-async-chunked-db", + "variant": null, + "pe_window_ns": 251.42399999999998, + "engine_window_ns": 214.914, + "flops": 262144, + "tflops": 1.2197623235340647, + "n_records": 8 + }, + { + "M": 128, + "K": 8, + "N": 128, + "bench": "matmul-composite", + "variant": "load_ref", + "pe_window_ns": 462.01, + "engine_window_ns": 397.0, + "flops": 262144, + "tflops": 0.6603123425692695, + "n_records": 82 + }, + { + "M": 128, + "K": 8, + "N": 128, + "bench": "matmul-async", + "variant": null, + "pe_window_ns": 247.798, + "engine_window_ns": 216.288, + "flops": 262144, + "tflops": 1.212013611480988, + "n_records": 4 + }, + { + "M": 128, + "K": 8, + "N": 128, + "bench": "matmul-async-chunked", + "variant": null, + "pe_window_ns": 247.798, + "engine_window_ns": 216.288, + "flops": 262144, + "tflops": 1.212013611480988, + "n_records": 4 + }, + { + "M": 128, + "K": 8, + "N": 128, + "bench": "matmul-async-chunked-db", + "variant": null, + "pe_window_ns": 247.798, + "engine_window_ns": 216.288, + "flops": 262144, + "tflops": 1.212013611480988, + "n_records": 4 + } +] \ No newline at end of file diff --git a/src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async_tflops.png b/src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async_tflops.png new file mode 100644 index 0000000..e52c62c Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async_tflops.png differ diff --git a/src/kernbench/benches/matmul_async.py b/src/kernbench/benches/matmul_async.py new file mode 100644 index 0000000..9d9e1b1 --- /dev/null +++ b/src/kernbench/benches/matmul_async.py @@ -0,0 +1,57 @@ +"""Single-PE async-loading GEMM (tl.load + tl.dot) for composite-vs-async +dispatch comparison. + +Pairs with ``matmul_composite``: both run the same (M,K,N) sweep, but +this kernel uses atomic primitives: + + a = tl.load(a_ptr, ...) # async DMA into TCM (1 PE_CPU cmd) + b = tl.load(b_ptr, ...) # async DMA into TCM (1 PE_CPU cmd), overlaps with a + out = tl.dot(a, b) # 1 PE_CPU cmd; GEMM engine starts only after + # both loads complete (tl.dot _await_pending) + tl.store(out_ptr, out) # 1 PE_CPU cmd + +Compared to composite: + - dispatch count grows to 4 instead of ~2 + - GEMM cannot start until *all* of B is in TCM (no per-tile overlap of + streaming B-load with GEMM compute) + +Env vars: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE (mirror matmul_composite). +""" +import os + +from kernbench.benches.registry import bench +from kernbench.policy.placement.dp import DPPolicy + + +def _kernel_async(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"): + M, K, N = int(M), int(K), int(N) + # Pre-stage A: load A and block the user kernel until its DMA finishes + # before dispatching load(B). Without this barrier, load(A) and load(B) + # run concurrently on the DMA engine and share HBM bandwidth, so the + # "exclude load(A) from the engine window" methodology does not + # actually subtract A's time — A's traffic still shows up as halved-BW + # contention with B. Calling _await_pending makes the load(A) the + # explicit pre-staging step, matching the composite load_ref semantic + # (where the composite waits for A before its first tile fetch). + a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE) + tl._await_pending(a) + b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE) + out = tl.dot(a, b) + tl.store(int(out_ptr), out) + + +@bench( + name="matmul-async", + description="Single-PE async-loading GEMM (tl.load + tl.dot) for " + "composite-vs-async dispatch comparison.", +) +def run(torch): + M = int(os.environ.get("MATMUL_M", "256")) + K = int(os.environ.get("MATMUL_K", "256")) + N = int(os.environ.get("MATMUL_N", "256")) + DTYPE = os.environ.get("MATMUL_DTYPE", "f16") + dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) + a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a") + b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b") + out = torch.empty((M, N), dtype=DTYPE, dp=dp, name="out") + torch.launch("matmul_async", _kernel_async, a, b, out, M, K, N) diff --git a/src/kernbench/benches/matmul_async_chunked.py b/src/kernbench/benches/matmul_async_chunked.py new file mode 100644 index 0000000..8c9b786 --- /dev/null +++ b/src/kernbench/benches/matmul_async_chunked.py @@ -0,0 +1,90 @@ +"""Single-PE chunked-prefetching async-loading GEMM for composite-vs-async +dispatch comparison. + +Splits K into hardware-tile-sized chunks (CHUNK_K = TILE_K = 64) so that: + - A chunks are loaded first (A pre-staged), then awaited + - B chunks are all issued as async tl.load at once → DMA engine queues + them FIFO and starts streaming + - Per-chunk tl.dot is issued; each dot blocks on its specific b_chunk + only, so dot[i] can start as soon as b[i] lands in TCM, overlapping + GEMM[i] with DMA of b[i+1..] still in flight + +Per-chunk dot outputs are accumulated with tl.add via a running sum +(``out = out + tl.dot(...)``), which serializes on the MATH engine due +to accumulator data dependency. This is the closest user-level analog to +composite's per-tile pipeline that is achievable with the current tl +DSL (which lacks a GEMM accumulator parameter and slicing primitive). + +Env: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE (same as matmul_composite). +""" +import os + +from kernbench.benches.registry import bench +from kernbench.policy.placement.dp import DPPolicy + +# Mirror PeSchedulerComponent.TILE_K (also TILE_K in milestone_1h_gemm.py). +TILE_K = 64 + + +def _kernel_async_chunked(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"): + M, K, N = int(M), int(K), int(N) + bpe = 2 # f16 + + if K <= TILE_K: + # Under-tile in K: chunking would make each dot smaller, even + # worse off. Fall back to a single dot (naive async behaviour). + a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE) + tl._await_pending(a) + b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE) + out = tl.dot(a, b) + tl.store(int(out_ptr), out) + return + + n_chunks = K // TILE_K + a_chunk_bytes = M * TILE_K * bpe + b_chunk_bytes = TILE_K * N * bpe + + # ── Phase 1: pre-stage all A chunks (issue + wait) ── + a_chunks = [ + tl.load(int(a_ptr) + i * a_chunk_bytes, + shape=(M, TILE_K), dtype=DTYPE) + for i in range(n_chunks) + ] + tl._await_pending(*a_chunks) + + # ── Phase 2: queue all B-chunk async loads up front ── + # The DMA engine processes them FIFO; later chunks are in flight + # while the early dots run on GEMM engine. + b_chunks = [ + tl.load(int(b_ptr) + i * b_chunk_bytes, + shape=(TILE_K, N), dtype=DTYPE) + for i in range(n_chunks) + ] + + # ── Phase 3: per-chunk dot with running accumulator ── + # dot[i] only blocks on b_chunks[i] (tl.dot's _await_pending). While + # dot[i] runs on GEMM, b_chunks[i+1..] continue draining off DMA. + out = tl.dot(a_chunks[0], b_chunks[0]) + for i in range(1, n_chunks): + out = out + tl.dot(a_chunks[i], b_chunks[i]) + + tl.store(int(out_ptr), out) + + +@bench( + name="matmul-async-chunked", + description="Single-PE chunked-prefetching async-loading GEMM " + "(CHUNK_K = TILE_K = 64; A pre-staged, B chunks " + "prefetched, per-chunk tl.dot + tl.add accumulator).", +) +def run(torch): + M = int(os.environ.get("MATMUL_M", "256")) + K = int(os.environ.get("MATMUL_K", "256")) + N = int(os.environ.get("MATMUL_N", "256")) + DTYPE = os.environ.get("MATMUL_DTYPE", "f16") + dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) + a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a") + b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b") + out = torch.empty((M, N), dtype=DTYPE, dp=dp, name="out") + torch.launch("matmul_async_chunked", _kernel_async_chunked, + a, b, out, M, K, N) diff --git a/src/kernbench/benches/matmul_async_chunked_db.py b/src/kernbench/benches/matmul_async_chunked_db.py new file mode 100644 index 0000000..77f64f8 --- /dev/null +++ b/src/kernbench/benches/matmul_async_chunked_db.py @@ -0,0 +1,110 @@ +"""Single-PE chunked-prefetching async GEMM with depth=2 double-buffer. + +A TCM-bounded variant of ``matmul_async_chunked``: at most two B-chunk +loads are in flight at any time, so peak TCM occupancy for B is +``2 × TILE_K × N × bpe`` rather than the full ``K × N × bpe``. + +This is the architecturally required form for any real LLM workload +where ``K = context_length`` (4 K–32 K tokens): the +``matmul_async_chunked`` variant queues *all* B-chunk loads up front +and therefore needs the full B in TCM; that exceeds the 1 MiB kernel +scratch cap (topology.yaml ``pe_tcm.kernel_scratch_mb``) at +sub-attention scale. Double-buffer prefetch is the closest +user-orchestrated kernel pattern to composite's per-tile streaming. + +Structure (after pre-staging A): + - Issue B_chunk[0] + B_chunk[1] async (2 outstanding) + - For each chunk i: + - tl.dot(A_chunk[i], B_chunk[i]) -- blocks on b[i], then issues GEMM + - tl.add into running accumulator + - Issue B_chunk[i+2] async (if available) -- keeps in-flight count = 2 + +Compared to the full-prefetch variant, depth-2 changes *when* DMA +descriptors hit the queue, not their total number; the structural +dispatch cost amortization story therefore stays the same. The +purpose of the variant is to honour the TCM cap, so the comparison +to composite reflects a kernel that could actually run on real +hardware at LLM-scale K. + +Env: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE. +""" +import os + +from kernbench.benches.registry import bench +from kernbench.policy.placement.dp import DPPolicy + +TILE_K = 64 +PREFETCH_DEPTH = 2 + + +def _kernel_async_chunked_db(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"): + M, K, N = int(M), int(K), int(N) + bpe = 2 # f16 + + if K <= TILE_K: + # Single-chunk fallback (under-tile or one-tile case): degenerates + # to naive async, no prefetch opportunity. + a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE) + tl._await_pending(a) + b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE) + out = tl.dot(a, b) + tl.store(int(out_ptr), out) + return + + n_chunks = K // TILE_K + a_chunk_bytes = M * TILE_K * bpe + b_chunk_bytes = TILE_K * N * bpe + + # Pre-stage all A chunks; this is the "activation pre-staged" half of + # load_ref. A is loaded once at the start and reused across every + # output dot. + a_chunks = [ + tl.load(int(a_ptr) + i * a_chunk_bytes, + shape=(M, TILE_K), dtype=DTYPE) + for i in range(n_chunks) + ] + tl._await_pending(*a_chunks) + + # Prime the prefetch pipeline: issue the first PREFETCH_DEPTH B-chunk + # loads up front. Anything beyond depth-2 stays in HBM until we get + # to its dot. + b_chunks: list = [ + tl.load(int(b_ptr) + i * b_chunk_bytes, + shape=(TILE_K, N), dtype=DTYPE) + for i in range(min(PREFETCH_DEPTH, n_chunks)) + ] + + # First dot — blocks on b_chunks[0]; while it computes, b_chunks[1] + # is already in flight. + out = tl.dot(a_chunks[0], b_chunks[0]) + + for i in range(1, n_chunks): + # Issue the next prefetch BEFORE the dot that consumes b_chunks[i], + # so peak outstanding B loads stays at PREFETCH_DEPTH. + if i + PREFETCH_DEPTH - 1 < n_chunks: + b_chunks.append( + tl.load(int(b_ptr) + (i + PREFETCH_DEPTH - 1) * b_chunk_bytes, + shape=(TILE_K, N), dtype=DTYPE) + ) + out = out + tl.dot(a_chunks[i], b_chunks[i]) + + tl.store(int(out_ptr), out) + + +@bench( + name="matmul-async-chunked-db", + description="Single-PE chunked-prefetching async GEMM with depth=2 " + "double-buffer (TCM-bounded streaming analog of composite " + "load_ref).", +) +def run(torch): + M = int(os.environ.get("MATMUL_M", "256")) + K = int(os.environ.get("MATMUL_K", "256")) + N = int(os.environ.get("MATMUL_N", "256")) + DTYPE = os.environ.get("MATMUL_DTYPE", "f16") + dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) + a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a") + b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b") + out = torch.empty((M, N), dtype=DTYPE, dp=dp, name="out") + torch.launch("matmul_async_chunked_db", _kernel_async_chunked_db, + a, b, out, M, K, N) diff --git a/tests/test_bench_registry.py b/tests/test_bench_registry.py index c222f56..d11f1e8 100644 --- a/tests/test_bench_registry.py +++ b/tests/test_bench_registry.py @@ -11,10 +11,13 @@ EXPECTED_NAMES = [ "gemm-single-pe", "gpt3-qkv", "ipcq-allreduce", + "matmul-async", + "matmul-async-chunked", + "matmul-async-chunked-db", "matmul-composite", "milestone-1h-ccl", "milestone-1h-gemm", - "milestone-gqa-headline", + "milestone-1h-gqa", "qkv-gemm", "qkv-gemm-multi-pe", "va-offset-verify",