diff --git a/dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp b/dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp new file mode 100644 index 0000000000..11b7fc44f0 --- /dev/null +++ b/dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp @@ -0,0 +1,290 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * Stream-K GEMM Dispatcher ctypes Library + * + * Provides C API for Python ctypes integration for the STREAM-K GEMM variant. + * Kernel header included via -include at compile time. + * + * Stream-K is a single GEMM (one A/B/C, one M/N/K) like regular GEMM, so this + * lib keeps the exact same C ABI as gemm_ctypes_lib.cpp -- ``dispatcher_run_gemm`` + * takes host A/B/C and M/N/K. The difference is internal: the generated launch + * has a Stream-K-specific signature + * + * static float launch(const ck_tile::StreamKHostArgs& args, const stream_config& stream); + * + * which allocates the reduction workspace internally (DeviceMem) and uses the + * Atomic reduction strategy. The single-problem registry path + * (g_dispatcher->run / GemmHostArgs) and the generated_tile_backend wrapper both + * hard-code the plain GemmHostArgs launch, so this lib bypasses the registry and + * calls SelectedKernel::launch(args, stream) directly, reporting the kernel name + * from the compile-time KERNEL_NAME macro. + * + * Because the C ABI matches the regular lib, the Python side reuses + * GemmDispatcherLib / GpuGemmRunner unchanged -- only the .so internals differ. + * + * Usage from Python: + * lib = ctypes.CDLL("libdispatcher_streamk_gemm.so") + * lib.dispatcher_init() + * lib.dispatcher_run_gemm(...) + */ + +#include +#include +#include +#include +#include +#include +#include + +// Kernel header included via -include compiler flag (with CK_TILE_SINGLE_KERNEL_INCLUDE). +// Defines: ADataType, BDataType, CDataType, AccDataType, SelectedKernel, KERNEL_NAME +// and transitively brings in ck_tile::StreamKHostArgs and ck_tile::stream_config. + +// GPU architecture - can be overridden via -DGFX_ARCH="gfx90a" at compile time +#ifndef GFX_ARCH +#define GFX_ARCH "gfx942" +#endif + +static bool g_initialized = false; + +// Read an integer benchmark knob from the environment, falling back to +// `fallback` when unset or unparseable. +static int env_int(const char* name, int fallback) +{ + const char* v = std::getenv(name); + if(v == nullptr || *v == '\0') + return fallback; + char* end = nullptr; + const long out = std::strtol(v, &end, 10); + if(end == v) + return fallback; + return static_cast(out); +} + +// Read a boolean benchmark knob ("0"/"false"/"off", any case => false, else true). +static bool env_bool(const char* name, bool fallback) +{ + const char* v = std::getenv(name); + if(v == nullptr || *v == '\0') + return fallback; + std::string s(v); + for(char& c : s) + if(c >= 'A' && c <= 'Z') + c = static_cast(c - 'A' + 'a'); + return !(s == "0" || s == "false" || s == "off"); +} + +extern "C" { + +/** + * Initialize the stream-k GEMM library. + * + * The stream-k path does not use the dispatcher/registry (it launches the + * force-included kernel directly), so this is a lightweight no-op kept for ABI + * parity with the regular GEMM lib. Returns 0 on success. + */ +int dispatcher_initialize() +{ + g_initialized = true; + return 0; +} + +/** + * Initialize dispatcher (alias) + */ +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Run a Stream-K GEMM on GPU by launching the force-included kernel directly. + * + * hipMalloc A/B/C, copy A and B host->device, memset C (the Atomic reduction + * strategy accumulates into C, so it must start zeroed), build a + * ck_tile::StreamKHostArgs whose strides are derived from the kernel's actual + * ALayout/BLayout/CLayout (no layout hardcoding) and launch. The launch + * allocates the reduction workspace internally and resets C between timed + * iterations. C is then copied back. + * + * The host buffers must be laid out to match each operand's layout (the Python + * runner arranges A/B/C as RowMajor=C-contiguous, ColumnMajor=F-contiguous). + * + * Returns: 0 on success, -1 on HIP error / generic throw, -2 if the kernel + * reports the arguments are unsupported. + */ +int dispatcher_run_gemm( + const void* A, const void* B, void* C, int64_t M, int64_t N, int64_t K, float* time_ms) +{ + if(!g_initialized || !A || !B || !C || M <= 0 || N <= 0 || K <= 0) + { + return -1; + } + + const ADataType* A_host = static_cast(A); + const BDataType* B_host = static_cast(B); + CDataType* C_host = static_cast(C); + + ADataType* A_dev = nullptr; + BDataType* B_dev = nullptr; + CDataType* C_dev = nullptr; + + auto cleanup_gpu_mem = [&]() { + if(A_dev) + (void)hipFree(A_dev); + if(B_dev) + (void)hipFree(B_dev); + if(C_dev) + (void)hipFree(C_dev); + }; + + if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + // Strides are DERIVED from the kernel's actual layouts (ALayout/BLayout/CLayout + // come from the force-included generated header) -- nothing layout-specific is + // hardcoded, so every layout (rcr/rrr/ccr/crr/...) works. A RowMajor R x C + // matrix has leading dim C; a ColumnMajor one has leading dim R. + // A is M x K, B is K x N, C is M x N. + using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; + const ck_tile::index_t lda = + static_cast(std::is_same_v ? K : M); + const ck_tile::index_t ldb = + static_cast(std::is_same_v ? N : K); + const ck_tile::index_t ldc = + static_cast(std::is_same_v ? N : M); + // k_batch is fixed to 1 inside StreamKHostArgs. + ck_tile::StreamKHostArgs args(static_cast(A_dev), + static_cast(B_dev), + static_cast(C_dev), + static_cast(M), + static_cast(N), + static_cast(K), + /*stride_A=*/lda, + /*stride_B=*/ldb, + /*stride_C=*/ldc); + + // Benchmark parameters. warmup/repeat default to old Tile Engine's values + // (warmup=50, repeat=100); a generous warmup keeps the GPU clock ramped, and + // 100 timed iterations give a stable median. These were the knobs behind the + // regular bridge's spurious "perf gap" (#8123): the old default of warmup=3/ + // repeat=10 measured a cold, un-ramped clock. Each knob is env-overridable so + // a caller can match another harness without recompiling. + // + // Divergence from the regular path (generated_tile_backend.hpp): flush_cache_ + // and rotating_count_ default OFF here. The Stream-K Atomic reduction + // accumulates into C, and the generated launch's launch_kernel_time_mask + // preprocess re-zeros only the original args.e_ptr -- rotating C across + // multiple buffers would leave the rotated copies un-zeroed and corrupt the + // accumulation. Leave rotating_count_=1 unless a caller knows the kernel + // re-zeros every rotated buffer. + ck_tile::stream_config stream_cfg; + stream_cfg.stream_id_ = nullptr; + stream_cfg.time_kernel_ = true; + stream_cfg.log_level_ = 0; + stream_cfg.cold_niters_ = env_int("CK_TILE_BENCH_WARMUP", 50); + stream_cfg.nrepeat_ = env_int("CK_TILE_BENCH_REPEAT", 100); + stream_cfg.is_gpu_timer_ = true; + stream_cfg.flush_cache_ = env_bool("CK_TILE_BENCH_FLUSH", false); + stream_cfg.rotating_count_ = env_int("CK_TILE_BENCH_ROTATING", 1); + + float exec_time = 0.0f; + try + { + exec_time = SelectedKernel::launch(args, stream_cfg); + } + catch(const std::exception& e) + { + cleanup_gpu_mem(); + if(std::string(e.what()).find("not supported") != std::string::npos) + { + if(time_ms) + { + *time_ms = -1.0f; + } + return -2; // Arguments not supported by this kernel + } + return -1; + } + + if(hipMemcpy(C_host, C_dev, M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + if(time_ms) + { + *time_ms = exec_time; + } + + cleanup_gpu_mem(); + return 0; +} + +/** + * Get kernel information (legacy single-kernel ABI). + * + * Returns the compile-time KERNEL_NAME of the force-included kernel header. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Get the name of the kernel at a given registry index (multi-kernel ABI). + * + * Each stream-k .so force-includes exactly one kernel header, so index 0 reports + * KERNEL_NAME and any other index is out of range. Mirrors the regular GEMM lib's + * name ABI so the Python bridge can use the same name-lookup path. + * Returns 0 on success, -1 on bad args or out-of-range index. + */ +int dispatcher_get_kernel_name_at(int index, char* buffer, int buffer_size) +{ + if(!buffer || buffer_size <= 0 || index != 0) + { + return -1; + } + + std::strncpy(buffer, KERNEL_NAME, static_cast(buffer_size) - 1); + buffer[buffer_size - 1] = '\0'; + return 0; +} + +/** + * Get the number of kernels in this .so (always 1 for the stream-k single-include lib). + */ +int dispatcher_get_kernel_count() { return 1; } + +/** + * Cleanup library resources (no-op; kept for ABI parity). + */ +void dispatcher_cleanup() { g_initialized = false; } + +} // extern "C" diff --git a/dispatcher/parity_diag/regression/ab_efficient_sweep.py b/dispatcher/parity_diag/regression/ab_efficient_sweep.py new file mode 100644 index 0000000000..c7a332e9b9 --- /dev/null +++ b/dispatcher/parity_diag/regression/ab_efficient_sweep.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +"""Efficient A/B sweep: bridge .so vs Old-TE binary, all layouts + fp16/bf16. + +Faster successor to run_alllayout_sweep.py: the bridge side batches all shapes +for a stem into ONE run_one_gemm_kernel.py worker call (one Python+numpy+CDLL +startup per stem instead of one per measurement). Old-TE binaries are run once +per shape; their internal warmup=50/repeat=100 already yields a stable median, +matching the prior methodology. + +- Bridge .so : main worktree dispatcher/build/examples (built from the FIXED source). +- Old-TE bin : develop-parity worktree build/bin (develop branch), per user instruction. + +Writes allresult_fp16_bf16.csv with resume support (keyed on stem,shape). + +CSV fields: stem,pipeline,dtype,layout,shape,bridge_tflops,old_tflops,gap_pct, + bridge_verified,oldte_built +""" +import csv, json, os, re, subprocess, sys, time +from pathlib import Path + +ROOT = Path("/home/AMD/muozturk/New_project/rocm-libraries/projects/composablekernel") +DISP = ROOT / "dispatcher" +WORKER = ROOT / "tile_engine/ops/gemm/run_one_gemm_kernel.py" +SO_DIR = DISP / "build" / "examples" +GEN_DIR = DISP / "build" / "generated_kernels" +OLD_BIN_DIR = Path( + "/home/AMD/muozturk/New_project/rocm-libraries/.claude/worktrees" + "/develop-parity/projects/composablekernel/build/bin" +) +REG = DISP / "parity_diag" / "regression" +STEMS_FILE = REG / "stems_selected.txt" +CSV_OUT = REG / "allresult_fp16_bf16.csv" + +PYPATH = os.pathsep.join([str(DISP / "python"), str(ROOT / "tile_engine/ops/gemm")]) +DEVICE = os.environ.get("PARITY_DEVICE", "0") + +SHAPES = [(512, 512, 512), (1024, 1024, 1024), (2048, 2048, 2048), + (1024, 512, 256), (4096, 4096, 4096)] + +FIELDS = ["stem", "pipeline", "dtype", "layout", "shape", + "bridge_tflops", "old_tflops", "gap_pct", + "bridge_verified", "oldte_built"] + +_TFLOPS_RE = re.compile(r'"tflops\(TFlops\)":\s*([0-9.]+)') + + +def pipeline_of(stem): + for p in ("compv3", "compv4", "mem"): + if f"_{p}_" in stem: + return p + return "other" + + +def base_env(): + env = os.environ.copy() + env["HIP_VISIBLE_DEVICES"] = DEVICE + env["GEMM_PYPATH"] = PYPATH + env["LD_LIBRARY_PATH"] = "/opt/rocm/lib:" + env.get("LD_LIBRARY_PATH", "") + return env + + +def run_bridge_all(stem): + """One batched worker call over all SHAPES. Returns {shape_str: tflops|None}.""" + so = SO_DIR / f"libgemm_{stem}.so" + out = {f"{M}x{N}x{K}": None for (M, N, K) in SHAPES} + if not so.exists(): + return out + # Staleness guard: a .so older than its generated header was built from an + # obsolete codegen and must NOT be measured -- doing so reports phantom + # regressions (the big 256-tile gaps in allresult_fp16_bf16_2.csv were all + # stale binaries that recovered to parity on rebuild). Treat stale as missing. + hdr = GEN_DIR / f"gemm_{stem}.hpp" + if hdr.exists() and so.stat().st_mtime < hdr.stat().st_mtime: + print(f" STALE .so (older than header), skipping: {stem}", file=sys.stderr, flush=True) + return out + items = [{"so_path": str(so), "problem": {"M": M, "N": N, "K": K}, + "kernel_name": f"gemm_{stem}"} for (M, N, K) in SHAPES] + payload = json.dumps({"items": items, "verify": False}) + try: + p = subprocess.run([sys.executable, str(WORKER)], input=payload.encode(), + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + env=base_env(), timeout=900) + except subprocess.TimeoutExpired: + return out + for line in p.stdout.decode().strip().splitlines(): + try: + d = json.loads(line) + except json.JSONDecodeError: + continue + idx = d.get("idx") + if isinstance(idx, int) and 0 <= idx < len(SHAPES) and d.get("ok"): + M, N, K = SHAPES[idx] + out[f"{M}x{N}x{K}"] = d.get("tflops") + return out + + +def run_oldte(stem, M, N, K): + binp = OLD_BIN_DIR / f"benchmark_gemm_universal_{stem}" + if not binp.exists(): + return None + try: + p = subprocess.run([str(binp), f"-m={M}", f"-n={N}", f"-k={K}", + "-warmup=50", "-repeat=100"], + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + env=base_env(), timeout=300) + except subprocess.TimeoutExpired: + return None + m = _TFLOPS_RE.search(p.stdout.decode()) + return float(m.group(1)) if m else None + + +def main(): + stems = [l.strip() for l in STEMS_FILE.read_text().splitlines() if l.strip()] + total = len(stems) * len(SHAPES) + done = set() + if CSV_OUT.exists(): + with open(CSV_OUT) as f: + for row in csv.DictReader(f): + done.add((row["stem"], row["shape"])) + mode = "a" if done else "w" + print(f"stems={len(stems)} shapes={len(SHAPES)} total={total} resume={len(done)}", flush=True) + + t0 = time.time(); n = len(done) + with open(CSV_OUT, mode, newline="") as fh: + w = csv.DictWriter(fh, fieldnames=FIELDS) + if mode == "w": + w.writeheader() + for stem in stems: + shapes_todo = [(M, N, K) for (M, N, K) in SHAPES + if (stem, f"{M}x{N}x{K}") not in done] + if not shapes_todo: + continue + parts = stem.split("_") + dtype, layout = parts[0], parts[1] + pipeline = pipeline_of(stem) + oldte_built = (OLD_BIN_DIR / f"benchmark_gemm_universal_{stem}").exists() + + bridge = run_bridge_all(stem) + for (M, N, K) in shapes_todo: + shape = f"{M}x{N}x{K}" + bt = bridge.get(shape) + ot = run_oldte(stem, M, N, K) + if bt is not None and ot not in (None, 0): + gap = (bt - ot) / ot * 100.0 + else: + gap = float("nan") + w.writerow(dict( + stem=stem, pipeline=pipeline, dtype=dtype, layout=layout, shape=shape, + bridge_tflops=f"{bt:.4f}" if bt is not None else "nan", + old_tflops=f"{ot:.4f}" if ot is not None else "nan", + gap_pct=f"{gap:.4f}" if gap == gap else "nan", + bridge_verified="None", oldte_built=str(oldte_built))) + fh.flush() + n += 1 + el = time.time() - t0 + rate = (n - len(done)) / el if el > 0 else 0 + eta = (total - n) / rate / 3600 if rate > 0 else 0 + print(f"[{n}/{total}] {stem[:48]:48} rate={rate:.1f}/s ETA={eta:.1f}h", flush=True) + print(f"DONE rows={n} -> {CSV_OUT}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/dispatcher/parity_diag/regression/ab_same_harness.py b/dispatcher/parity_diag/regression/ab_same_harness.py new file mode 100644 index 0000000000..9191071744 --- /dev/null +++ b/dispatcher/parity_diag/regression/ab_same_harness.py @@ -0,0 +1,310 @@ +#!/usr/bin/env python3 +"""Apples-to-apples GEMM A/B: bridge kernel vs old-TE kernel, ONE harness. + +Why this exists +--------------- +The earlier sweep (allsweep6144rcrfp16.py) compared the bridge's dispatcher +measurement against old TE's *standalone benchmark binary* +(benchmark_gemm_universal_). That comparison is NOT apples-to-apples: +the device kernel is byte-identical, yet old TE's standalone binary reports +~18-20% lower TFLOPS at e.g. 1024^3 / compv4. rocprof shows the identical +kernel genuinely runs longer in that process -- ~+8% cycles plus a lower +sustained SCLK -- a power/clock + execution-environment artifact of that +binary, NOT a bridge speedup, compiler difference, or kernel difference. +(See diagnose.md sec.4.) + +This harness removes the artifact: it builds the OLD-TE kernel into a .so from +old TE's own generated header and runs BOTH the bridge kernel and the old-TE +kernel through the SAME worker (run_one_gemm_kernel.py). Measured this way the +gap collapses to ~1%, which is the honest result. + +The old-TE generated-header directory is derived per stem as +``///`` (e.g. fp16/rcr, bf16/crr), so a single +run covers every dtype/layout. Set OLD_TE_GEN to pin one explicit leaf dir for +all stems (legacy behavior); set OLD_TE_GEN_BASE to relocate the base. + +Usage: + python3 ab_same_harness.py # default kernel list + shapes + python3 ab_same_harness.py [...] # explicit stems + python3 ab_same_harness.py --stems-file F [--csv OUT] # sweep a stems file +""" +import argparse +import csv +import json +import os +import statistics +import subprocess +import sys +from pathlib import Path + +# composablekernel root: .../composablekernel/dispatcher/parity_diag/regression/ +ROOT = Path(__file__).resolve().parents[3] +DISP = ROOT / "dispatcher" +GEN = DISP / "build" / "generated_kernels" +SRC = DISP / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp" +STATIC = DISP / "build" / "libck_tile_dispatcher.a" +BR_SO_DIR = DISP / "build" / "examples" +WORKER = ROOT / "tile_engine/ops/gemm/run_one_gemm_kernel.py" +# Base dir of old-TE generated single-kernel headers; the per-stem leaf +# (/) is appended in old_gen_dir(). Points at a sibling +# develop-parity worktree under the rocm-libraries root by default. +OLD_GEN_BASE = Path(os.environ.get( + "OLD_TE_GEN_BASE", + str(ROOT.parents[1] / ".claude/worktrees/develop-parity" + "/projects/composablekernel/build/tile_engine/ops/gemm/gemm_universal"), +)) +# Legacy explicit override: when set, this exact leaf dir is used for ALL stems. +OLD_GEN_PIN = os.environ.get("OLD_TE_GEN") +OUT = DISP / "parity_diag" / "regression" / "_ab_same_harness_build" +ARCH = os.environ.get("GFX_ARCH", "gfx942") +DEVICE = os.environ.get("PARITY_DEVICE", "0") +REPEATS = int(os.environ.get("AB_REPEATS", "3")) + +SHAPES = [(512, 512, 512), (1024, 1024, 1024), (2048, 2048, 2048), + (1024, 512, 256), (4096, 4096, 4096)] + +DEFAULT_STEMS = [ + "fp16_rcr_compv4_default_intrawave_False_False_False_False_64x128x64_2x2x1_32x32x16", + "fp16_rcr_compv4_cshuffle_intrawave_False_False_False_False_64x128x64_1x4x1_32x32x16", + "fp16_rcr_compv4_default_intrawave_False_False_False_False_128x128x64_4x1x1_32x32x16", +] + +PYPATH = os.pathsep.join([str(DISP / "python"), str(ROOT / "tile_engine/ops/gemm")]) + + +def old_gen_dir(stem: str) -> Path: + """Old-TE header dir for a stem: // (or the pinned dir). + + Stems are named ``__...`` (e.g. fp16_rcr_..., bf16_crr_...), + which is exactly the develop-parity gen-tree layout, so the leaf is derived + from the stem itself -- no per-layout hardcoding. + """ + if OLD_GEN_PIN: + return Path(OLD_GEN_PIN) + parts = stem.split("_") + dtype, layout = parts[0], parts[1] + return OLD_GEN_BASE / dtype / layout + + +def build_old_so(stem: str) -> Path | None: + """Compile old TE's generated kernel header into a bridge-loadable .so. + + Cached: if the .so already exists it is reused, so a parallel --build-only + pre-pass (CPU-bound hipcc) can be separated from the serial GPU measurement. + """ + hdr = old_gen_dir(stem) / f"gemm_universal_single_{stem}.hpp" + if not hdr.exists(): + return None + OUT.mkdir(parents=True, exist_ok=True) + obj = OUT / f"{stem}.o" + lib = OUT / f"libold_{stem}.so" + if lib.exists(): + return lib + common = [ + "-fPIC", "-O3", + f"-I{DISP / 'include'}", f"-I{ROOT / 'include'}", f"-I{ROOT}", f"-I{GEN}", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", f"-include{hdr}", "-D__HIP_PLATFORM_AMD__", + f"--offload-arch={ARCH}", f'-DGFX_ARCH="{ARCH}"', + # Match the bridge build's AMDGPU codegen flags (gemm_utils.py + # _build_compile_jobs / _TILE_ENGINE_CODEGEN_FLAGS), which are also what + # Tile Engine's own CMake passes. Without these the old-TE side is built + # with a *different* instruction schedule (notably -enable-post-misched + # defaults back on) and runs ~10-40% faster than real old-TE, making the + # bridge look regressed when it is actually at parity. Build BOTH sides + # identically so the A/B measures the kernel, not a flag asymmetry. + "-mllvm", "-enable-noalias-to-md-conversion=0", + "-mllvm", "--lsr-drop-solution=1", + "-mllvm", "-enable-post-misched=0", + "-mllvm", "-amdgpu-early-inline-all=true", + "-mllvm", "-amdgpu-function-calls=false", + "-fno-offload-uniform-block", + "-Wno-undefined-func-template", "-Wno-float-equal", + ] + cc = subprocess.run(["/opt/rocm/bin/hipcc", "-c", *common, str(SRC), "-o", str(obj)], + capture_output=True) + if cc.returncode != 0: + return None + ln = subprocess.run(["/opt/rocm/bin/hipcc", "-shared", "-fPIC", + f"--offload-arch={ARCH}", "--hip-link", + str(obj), str(STATIC), "-o", str(lib)], capture_output=True) + return lib if ln.returncode == 0 else None + + +def meas(so: Path, M: int, N: int, K: int) -> float | None: + """Median TFLOPS over REPEATS worker calls (each call does its own + warmup=50/repeat=100 internally). Median, not max, to match the sweep + methodology and stay robust to the occasional clock-warmup outlier.""" + if not so or not Path(so).exists(): + return None + payload = json.dumps({"so_path": str(so), "problem": {"M": M, "N": N, "K": K}, + "kernel_name": "x"}) + env = os.environ.copy() + env["HIP_VISIBLE_DEVICES"] = DEVICE + env["GEMM_PYPATH"] = PYPATH + env["LD_LIBRARY_PATH"] = "/opt/rocm/lib:" + env.get("LD_LIBRARY_PATH", "") + samples = [] + for _ in range(REPEATS): + p = subprocess.run([sys.executable, str(WORKER)], input=payload.encode(), + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, env=env) + for line in p.stdout.decode().splitlines(): + try: + d = json.loads(line) + except json.JSONDecodeError: + continue + if d.get("ok"): + samples.append(d["tflops"]) + return statistics.median(samples) if samples else None + + +def meas_all(so: Path) -> dict: + """Median TFLOPS per shape from REPEATS *batched* worker calls. + + One worker call measures ALL shapes (5x fewer python+numpy+CDLL startups + than per-shape meas()), which is the throughput lever for a full sweep on a + single GPU. Returns {shape_str: tflops|None}.""" + out = {f"{M}x{N}x{K}": None for (M, N, K) in SHAPES} + if not so or not Path(so).exists(): + return out + items = [{"so_path": str(so), "problem": {"M": M, "N": N, "K": K}, + "kernel_name": "x"} for (M, N, K) in SHAPES] + payload = json.dumps({"items": items, "verify": False}) + env = os.environ.copy() + env["HIP_VISIBLE_DEVICES"] = DEVICE + env["GEMM_PYPATH"] = PYPATH + env["LD_LIBRARY_PATH"] = "/opt/rocm/lib:" + env.get("LD_LIBRARY_PATH", "") + samples = {s: [] for s in out} + for _ in range(REPEATS): + p = subprocess.run([sys.executable, str(WORKER)], input=payload.encode(), + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + env=env, timeout=900) + for line in p.stdout.decode().splitlines(): + try: + d = json.loads(line) + except json.JSONDecodeError: + continue + idx = d.get("idx") + if isinstance(idx, int) and 0 <= idx < len(SHAPES) and d.get("ok"): + M, N, K = SHAPES[idx] + samples[f"{M}x{N}x{K}"].append(d["tflops"]) + for s, xs in samples.items(): + if xs: + out[s] = statistics.median(xs) + return out + + +def pipeline_of(stem: str) -> str: + for p in ("compv3", "compv4", "mem"): + if f"_{p}_" in stem: + return p + return "other" + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("stems", nargs="*", help="kernel stems to A/B") + ap.add_argument("--stems-file", help="file with one stem per line") + ap.add_argument("--csv", help="write results to CSV (resume-aware)") + ap.add_argument("--build-only", action="store_true", + help="parallel-compile old-TE .so for all stems, then exit " + "(CPU pre-pass; GPU measurement reuses the cache)") + ap.add_argument("--jobs", type=int, default=min(os.cpu_count() or 8, 16), + help="parallel compile jobs for --build-only") + args = ap.parse_args() + + stems = list(args.stems) + if args.stems_file: + stems += [l.strip() for l in Path(args.stems_file).read_text().splitlines() + if l.strip()] + stems = stems or DEFAULT_STEMS + + # Parallel CPU pre-compile of every old-TE .so (no GPU touched). + if args.build_only: + from concurrent.futures import ProcessPoolExecutor, as_completed + ok = miss = fail = 0 + print(f"build-only: {len(stems)} stems, jobs={args.jobs}", flush=True) + with ProcessPoolExecutor(max_workers=args.jobs) as ex: + futs = {ex.submit(build_old_so, s): s for s in stems} + for i, fut in enumerate(as_completed(futs), 1): + try: + r = fut.result() + except Exception: + r = None + s = futs[fut] + if r is None: + # distinguish "no header" from "compile failed" + if (old_gen_dir(s) / f"gemm_universal_single_{s}.hpp").exists(): + fail += 1 + else: + miss += 1 + else: + ok += 1 + if i % 100 == 0: + print(f" [{i}/{len(stems)}] ok={ok} no_header={miss} fail={fail}", + flush=True) + print(f"build-only DONE: ok={ok} no_header={miss} fail={fail}", flush=True) + return + + # CSV sweep mode: same columns as the (now-corrected) sweep, resume-aware. + if args.csv: + fields = ["stem", "pipeline", "dtype", "layout", "shape", + "bridge_tflops", "old_tflops", "gap_pct", "oldte_built"] + out = Path(args.csv) + done = set() + if out.exists(): + with open(out) as f: + for row in csv.DictReader(f): + done.add((row["stem"], row["shape"])) + mode = "a" if done else "w" + print(f"stems={len(stems)} shapes={len(SHAPES)} resume={len(done)} -> {out}", + flush=True) + with open(out, mode, newline="") as fh: + w = csv.DictWriter(fh, fieldnames=fields) + if mode == "w": + w.writeheader() + for stem in stems: + todo = [(M, N, K) for (M, N, K) in SHAPES + if (stem, f"{M}x{N}x{K}") not in done] + if not todo: + continue + parts = stem.split("_") + dtype, layout = parts[0], parts[1] + old_so = build_old_so(stem) + br_so = BR_SO_DIR / f"libgemm_{stem}.so" + # Batched: one worker call per side covers all shapes. + bridge = meas_all(br_so) + old = meas_all(old_so) if old_so else {} + for (M, N, K) in todo: + shape = f"{M}x{N}x{K}" + b = bridge.get(shape) + o = old.get(shape) + gap = (b - o) / o * 100 if (b and o) else float("nan") + w.writerow(dict( + stem=stem, pipeline=pipeline_of(stem), dtype=dtype, + layout=layout, shape=shape, + bridge_tflops=f"{b:.4f}" if b is not None else "nan", + old_tflops=f"{o:.4f}" if o is not None else "nan", + gap_pct=f"{gap:.4f}" if gap == gap else "nan", + oldte_built=str(old_so is not None))) + fh.flush() + print(f" done {stem[:60]}", flush=True) + print(f"DONE -> {out}", flush=True) + return + + # Pretty-print mode. + print(f"{'shape':>14} {'bridge':>9} {'oldTE':>9} {'gap%':>7} kernel") + for stem in stems: + old_so = build_old_so(stem) + br_so = BR_SO_DIR / f"libgemm_{stem}.so" + if old_so is None: + print(f" [skip: no old-TE header] {stem}") + continue + for (M, N, K) in SHAPES: + b = meas(br_so, M, N, K) + o = meas(old_so, M, N, K) + gap = (b - o) / o * 100 if (b and o) else float("nan") + print(f"{f'{M}x{N}x{K}':>14} {b or float('nan'):9.2f} " + f"{o or float('nan'):9.2f} {gap:7.2f} {stem[:40]}") + + +if __name__ == "__main__": + main() diff --git a/dispatcher/python/ctypes_utils.py b/dispatcher/python/ctypes_utils.py index cc94ede685..56ea9ea5a8 100644 --- a/dispatcher/python/ctypes_utils.py +++ b/dispatcher/python/ctypes_utils.py @@ -1389,7 +1389,7 @@ class KernelConfig: gfx_arch: str = "gfx942" # GEMM variant (affects arch filter validation) - # "standard", "preshuffle", or "multi_d" + # "standard", "preshuffle", "multi_d", or "stream_k" variant: str = "standard" @property diff --git a/dispatcher/python/gemm_utils.py b/dispatcher/python/gemm_utils.py index 26e68c061b..52417e9a99 100644 --- a/dispatcher/python/gemm_utils.py +++ b/dispatcher/python/gemm_utils.py @@ -187,6 +187,9 @@ class GemmKernelConfig: gfx_arch: str = "gfx942" variant: str = "standard" + # Stream-K reduction strategy: "atomic" (default), "linear", or "tree". + # Only meaningful when variant == "stream_k". + reduction_strategy: str = "atomic" # ------------------------------------------------------------------ # # Derived string fragments @@ -231,8 +234,12 @@ class GemmKernelConfig: ) if self.variant == "preshuffle": name += "_preshuffle" - elif self.variant == "streamk": + elif self.variant == "stream_k": name += "_streamk" + # Atomic keeps the bare "_streamk" suffix (original parity); linear + # and tree are disambiguated, matching KernelNaming.generate. + if self.reduction_strategy != "atomic": + name += f"_{self.reduction_strategy}" elif self.variant == "grouped": name += "_grouped" return name @@ -247,7 +254,7 @@ class GemmKernelConfig: triple ``warp_*`` and the MFMA triple ``warp_tile_*``. We translate from dispatcher semantics here so the mapping cannot drift. """ - return { + cfg = { "tile_config": { "tile_m": [self.tile_m], "tile_n": [self.tile_n], @@ -271,6 +278,11 @@ class GemmKernelConfig: "persistent": [self.persistent], }, } + # Pin the single reduction strategy so stream-K codegen emits exactly this + # kernel (the generator otherwise expands all strategies in its default). + if self.variant == "stream_k": + cfg["streamk_config"] = {"reduction_strategy": [self.reduction_strategy]} + return cfg def to_dict(self) -> Dict[str, Any]: return { @@ -985,10 +997,16 @@ def _tile_engine_codegen_flags() -> Tuple[str, ...]: def _ctypes_source_name(config: GemmKernelConfig) -> str: """Pick the ctypes ABI source for a config's variant. - The grouped kernel has a multi-problem launch signature that the - single-problem ``gemm_ctypes_lib.cpp`` cannot express, so grouped configs - compile against the dedicated ``grouped_gemm_ctypes_lib.cpp``. + Variants whose launch ABI differs from the single-problem + ``dispatcher_run_gemm`` path need their own lib: + * stream_k keeps the single-problem C ABI (single A/B/C, M/N/K) but its + lib builds a ``StreamKHostArgs`` and calls ``SelectedKernel::launch`` + directly instead of routing through the registry. + * grouped has a multi-problem launch signature the single-problem + ``gemm_ctypes_lib.cpp`` cannot express. """ + if config.variant == "stream_k": + return "streamk_gemm_ctypes_lib.cpp" if config.variant == "grouped": return "grouped_gemm_ctypes_lib.cpp" return "gemm_ctypes_lib.cpp" @@ -1007,6 +1025,29 @@ def _build_compile_jobs( lib_path = build_dir / "examples" / f"lib{config.name}.so" obj_file = lib_path.with_suffix(".o") + # The Stream-K path skips the cmake build that would normally create this + # directory, so ensure it exists before hipcc writes the object/.so here. + lib_path.parent.mkdir(parents=True, exist_ok=True) + + # Per-variant AMDGPU codegen flags. The regular path matches Tile Engine's + # gemm_universal build via _tile_engine_codegen_flags(). Stream-K must instead + # match TE's gemm_streamk build EXACTLY for a fair A/B: -enable-post-misched=0 + # is applied unconditionally (not persistent-gated) and it does NOT use + # -enable-noalias-to-md-conversion=0. + is_streamk = getattr(config, "variant", "") == "stream_k" + variant_flags = ( + [ + "-std=c++20", + "-fno-offload-uniform-block", + "-mllvm", "--lsr-drop-solution=1", + "-mllvm", "-enable-post-misched=0", + "-mllvm", "-amdgpu-early-inline-all=true", + "-mllvm", "-amdgpu-function-calls=false", + "--offload-compress", + ] + if is_streamk + else list(_tile_engine_codegen_flags()) + ) compile_cmd = [ _resolve_hipcc(), @@ -1022,13 +1063,13 @@ def _build_compile_jobs( "-D__HIP_PLATFORM_AMD__", f"--offload-arch={config.gfx_arch}", f'-DGFX_ARCH="{config.gfx_arch}"', - # Match Tile Engine's AMDGPU codegen flags exactly (see + # Match Tile Engine's AMDGPU codegen flags exactly (see variant_flags / # _tile_engine_codegen_flags). Without them the kernel is compiled with # different inlining/register allocation, which changes occupancy; # persistent kernels size their grid by occupancy # (UniversalGemmKernel::MaxOccupancyGridSize = #CUs x occupancy), so a # mismatch shows up as large perf gaps vs Tile Engine on persistent tiles. - *_tile_engine_codegen_flags(), + *variant_flags, "-Wno-undefined-func-template", "-Wno-float-equal", str(ctypes_source), @@ -1042,7 +1083,10 @@ def _build_compile_jobs( f"--offload-arch={config.gfx_arch}", "--hip-link", str(obj_file), - str(static_lib), + # The Stream-K ctypes lib launches the force-included kernel directly and + # references no registry/dispatcher symbols, so its .so does not need the + # dispatcher static lib. The regular path still links it. + *([] if is_streamk else [str(static_lib)]), "-o", str(lib_path), ] @@ -1086,7 +1130,11 @@ def setup_multiple_gemm_dispatchers( ctypes_dir = _cu.get_dispatcher_root() / "bindings" / "ctypes" needed_sources = {ctypes_dir / _ctypes_source_name(c) for c in configs} missing = [str(p) for p in needed_sources if not p.exists()] - if not static_lib.exists() or missing: + # Stream-K .so links only the force-included kernel (no registry/dispatcher + # symbols), so it does not need the dispatcher static lib; only the regular + # path requires it. + streamk_build = all(c.variant == "stream_k" for c in configs) + if (not streamk_build and not static_lib.exists()) or missing: raise FileNotFoundError( "Missing static lib or ctypes source required for compilation:\n" f" {static_lib}\n " + "\n ".join(missing) + "\n" @@ -1249,6 +1297,14 @@ def expand_sweep( pad_ks = _expand_values(tr.get("pad_k"), [False]) persistents = _expand_values(tr.get("persistent"), [False]) + # Stream-K only: sweep reduction strategies (atomic/linear/tree). Other + # variants keep a single dummy value so the product is unaffected. + if variant == "stream_k": + sk = cfg.get("streamk_config", {}) + reductions = _expand_values(sk.get("reduction_strategy"), ["atomic"]) + else: + reductions = ["atomic"] + la, lb, lc = layout[0], layout[1], layout[2] configs: List[GemmKernelConfig] = [] @@ -1270,6 +1326,7 @@ def expand_sweep( pn, pk, persist, + red, ) in itertools.product( tile_ms, tile_ns, @@ -1287,6 +1344,7 @@ def expand_sweep( pad_ns, pad_ks, persistents, + reductions, ): c = GemmKernelConfig( dtype_a=dtype, @@ -1314,6 +1372,7 @@ def expand_sweep( persistent=bool(persist), gfx_arch=arch, variant=variant, + reduction_strategy=red, ) if c.name in seen: continue diff --git a/dispatcher/tests/test_streamk_gemm_utils.py b/dispatcher/tests/test_streamk_gemm_utils.py new file mode 100644 index 0000000000..4d698ff09d --- /dev/null +++ b/dispatcher/tests/test_streamk_gemm_utils.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU-only unit tests for the stream-K surface of python/gemm_utils.py. + +The stream-K bridge adds a ``reduction_strategy`` to GemmKernelConfig and a +dedicated ctypes source. These tests lock in the two pieces of pure host-side +logic that must stay byte-exact with the codegen and the build: + + * ``GemmKernelConfig.name`` -- the suffix rules mirror + unified_gemm_codegen.py::KernelNaming.generate. Atomic keeps the bare + ``_streamk`` (original parity name); linear/tree are disambiguated as + ``_streamk_``. If this drifts, the runtime registry lookup key + misses the kernel baked into the generated header. + * ``_ctypes_source_name`` -- stream-K launches StreamKHostArgs directly + (registry-bypass) so it needs its own bridge .cpp; every other variant + shares gemm_ctypes_lib.cpp. + +No GPU is touched. Run: python3 -m pytest tests/test_streamk_gemm_utils.py -v +""" + +import sys +import unittest +from pathlib import Path + +SCRIPT_DIR = Path(__file__).parent.resolve() +DISPATCHER_DIR = SCRIPT_DIR.parent +sys.path.insert(0, str(DISPATCHER_DIR / "python")) + +from gemm_utils import ( # noqa: E402 + GemmKernelConfig, + _ctypes_source_name, + _dtype_from_kernel_name, + _layout_from_kernel_name, +) + + +class TestStreamKNaming(unittest.TestCase): + """Stream-K variant naming and reduction-strategy plumbing.""" + + def _cfg(self, variant="stream_k", reduction_strategy="atomic"): + return GemmKernelConfig(variant=variant, reduction_strategy=reduction_strategy) + + def test_atomic_keeps_bare_streamk_suffix(self): + name = self._cfg(reduction_strategy="atomic").name + self.assertTrue(name.endswith("_streamk")) + self.assertNotIn("_streamk_atomic", name) + + def test_linear_and_tree_are_disambiguated(self): + self.assertTrue( + self._cfg(reduction_strategy="linear").name.endswith("_streamk_linear") + ) + self.assertTrue( + self._cfg(reduction_strategy="tree").name.endswith("_streamk_tree") + ) + + def test_standard_has_no_streamk_suffix(self): + self.assertNotIn("streamk", self._cfg(variant="standard").name) + + def test_streamk_name_still_roundtrips_dtype_and_layout(self): + # The variant suffix must not disturb the dtype/layout tokens the runner + # parses back out of the compiled .so name. + for red in ("atomic", "linear", "tree"): + cfg = GemmKernelConfig( + dtype_a="bf16", + dtype_b="bf16", + dtype_c="bf16", + layout_a="col", + layout_b="col", + layout_c="row", + variant="stream_k", + reduction_strategy=red, + ) + self.assertEqual(_dtype_from_kernel_name(cfg.name), "bf16") + self.assertEqual(_layout_from_kernel_name(cfg.name), "ccr") + + def test_codegen_json_pins_reduction_only_for_streamk(self): + sk = self._cfg(reduction_strategy="tree").to_codegen_json() + self.assertEqual(sk["streamk_config"], {"reduction_strategy": ["tree"]}) + # Non-stream-K configs must not emit a streamk_config block. + self.assertNotIn( + "streamk_config", + GemmKernelConfig(variant="standard").to_codegen_json(), + ) + + +class TestCtypesSourceRouting(unittest.TestCase): + """Each variant routes to the ctypes bridge .cpp matching its launch ABI.""" + + def test_streamk_gets_dedicated_source(self): + cfg = GemmKernelConfig(variant="stream_k") + self.assertEqual(_ctypes_source_name(cfg), "streamk_gemm_ctypes_lib.cpp") + + def test_standard_uses_default_source(self): + self.assertEqual( + _ctypes_source_name(GemmKernelConfig(variant="standard")), + "gemm_ctypes_lib.cpp", + ) + self.assertEqual( + _ctypes_source_name(GemmKernelConfig(variant="preshuffle")), + "gemm_ctypes_lib.cpp", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tile_engine/ops/gemm/gemm_streamk/configs/default_config.json b/tile_engine/ops/gemm/gemm_streamk/configs/default_config.json new file mode 100644 index 0000000000..59f9a888a4 --- /dev/null +++ b/tile_engine/ops/gemm/gemm_streamk/configs/default_config.json @@ -0,0 +1,89 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 128 + ] + }, + "tile_n": { + "values": [ + 128 + ] + }, + "tile_k": { + "values": [ + 32, + 64 + ] + }, + "warp_m": { + "values": [ + 2 + ] + }, + "warp_n": { + "values": [ + 2 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 32 + ] + }, + "warp_tile_n": { + "values": [ + 32 + ] + }, + "warp_tile_k": { + "values": [ + 16 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3", + "compv4" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "cshuffle" + ] + }, + "pad_m": { + "values": [ + true + ] + }, + "pad_n": { + "values": [ + true + ] + }, + "pad_k": { + "values": [ + true + ] + }, + "persistent": { + "values": [ + false + ] + } + }, + "k_block_per_cu": 1 +} diff --git a/tile_engine/ops/gemm/run_one_streamk_gemm_kernel.py b/tile_engine/ops/gemm/run_one_streamk_gemm_kernel.py new file mode 100644 index 0000000000..2b9450b5af --- /dev/null +++ b/tile_engine/ops/gemm/run_one_streamk_gemm_kernel.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +"""Worker script for running Stream-K GEMM kernels in an isolated subprocess. + +Stream-K is a single-problem GEMM with the same C ABI as regular GEMM, so this +worker is identical in shape to ``run_one_gemm_kernel.py`` and reuses +``GemmProblem`` / ``GpuGemmRunner`` unchanged -- the Stream-K specifics live +entirely inside the force-included kernel and ``streamk_gemm_ctypes_lib.cpp``. + +- Receives kernel config + problem via stdin as JSON +- Loads the .so library ONLY inside this subprocess +- Outputs timing results as JSON to stdout (one line per kernel, flushed) +- A GPU fault kills only this process; the parent driver can continue + +Input JSON format: + Single: {"so_path": "...", "problem": {"M":.., "N":.., "K":..}, "kernel_name": "..."} + Batch: {"items": [{"so_path": "...", "problem": {...}, "kernel_name": "..."}, ...]} + +Optional top-level keys ``verify`` (bool) and ``verify_tol`` (float) enable an +fp32 numpy reference check; when set, each OK result also carries ``verified`` +and ``max_rel``. Stream-K's Atomic reduction does multiple fp16 atomic-adds (one +per K-split partial), so it is inherently noisier than a single fp32->fp16 store; +the default gate tolerance (2e-2) is loose enough to pass while still catching +gross errors. + +Output JSON format (one line per kernel): + {"idx": 0, "ok": true, "ms": 0.123, "tflops": 456.7, "non_zero": 1, "kernel": "..."} + {"idx": 0, "ok": true, ..., "verified": true, "max_rel": 1.1e-3} # with --verify + {"idx": 1, "ok": false, "error": "...", "kernel": "..."} +""" + +import json +import os +import sys + +# Add dispatcher python paths from environment (os.pathsep-separated). +gemm_pypath = os.environ.get("GEMM_PYPATH", "") +if gemm_pypath: + for p in gemm_pypath.split(os.pathsep): + if p and p not in sys.path: + sys.path.insert(0, p) + +from gemm_utils import ( # noqa: E402 + GemmProblem, + GpuGemmRunner, + _dtype_from_kernel_name, + _fp32_to_bf16_u16, + _bf16_u16_to_fp32, + _fp32_to_fp8_u8, + _fp8_u8_to_fp32, + _fp32_to_bf8_u8, + _bf8_u8_to_fp32, +) +import numpy as np # noqa: E402 + + +def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2): + """Run a single kernel and emit its result as one JSON line. + + When ``verify`` is set, the kernel output is checked against an fp32 numpy + reference (``A @ B``) using the global relative metric + ``max|out - ref| / max|ref|``; the emitted ``verified`` field then reflects + correctness, not just liveness (``non_zero``). + """ + try: + problem = GemmProblem.from_dict(prob_dict) + + np.random.seed(42) + A = (np.random.randn(problem.M, problem.K) * 0.1).astype(np.float32) + B = (np.random.randn(problem.K, problem.N) * 0.1).astype(np.float32) + + # CRITICAL: load the library ONLY inside this subprocess. The runner reads + # dtype + layout off the kernel name and arranges/encodes A/B accordingly. + runner = GpuGemmRunner(lib_path=so_path) + result = runner.run(A, B, problem) + + if result.success: + non_zero = ( + int(np.count_nonzero(result.output)) + if result.output is not None + else 0 + ) + out = { + "idx": idx, + "ok": True, + "ms": result.time_ms, + "tflops": result.tflops, + "non_zero": non_zero, + "kernel": kernel_name, + } + if verify: + # Reference uses the SAME quantized inputs the device sees, per the + # kernel's dtype (bf16/fp8/bf8 bit-quantization vs fp16), so the + # metric isolates compute error from input quantization. The dtype + # comes from the kernel name and the quantizers are the same module + # helpers GpuGemmRunner uses to build the device buffers, so host + # and device see identical inputs. + kdt = _dtype_from_kernel_name(runner.kernel_name) + if kdt == "bf16": + Aq = _bf16_u16_to_fp32(_fp32_to_bf16_u16(A)) + Bq = _bf16_u16_to_fp32(_fp32_to_bf16_u16(B)) + elif kdt == "fp8": + Aq = _fp8_u8_to_fp32(_fp32_to_fp8_u8(A)) + Bq = _fp8_u8_to_fp32(_fp32_to_fp8_u8(B)) + elif kdt == "bf8": + Aq = _bf8_u8_to_fp32(_fp32_to_bf8_u8(A)) + Bq = _bf8_u8_to_fp32(_fp32_to_bf8_u8(B)) + else: # fp16 + Aq = A.astype(np.float16).astype(np.float32) + Bq = B.astype(np.float16).astype(np.float32) + ref = Aq @ Bq + got = result.output.astype(np.float32) + denom = float(np.max(np.abs(ref))) or 1.0 + max_rel = float(np.max(np.abs(got - ref)) / denom) + out["max_rel"] = max_rel + out["verified"] = bool(max_rel <= verify_tol) + print(json.dumps(out), flush=True) + else: + print( + json.dumps( + { + "idx": idx, + "ok": False, + "error": f"kernel returned status {result.status}", + "kernel": kernel_name, + } + ), + flush=True, + ) + + except Exception as e: + print( + json.dumps( + {"idx": idx, "ok": False, "error": str(e), "kernel": kernel_name} + ), + flush=True, + ) + + +def main(): + """Read JSON from stdin, run kernel(s), output results.""" + try: + d = json.loads(sys.stdin.buffer.read()) + except Exception as e: + print( + json.dumps({"idx": 0, "ok": False, "error": f"JSON parse error: {e}"}), + flush=True, + ) + sys.exit(1) + + verify = bool(d.get("verify", False)) + verify_tol = float(d.get("verify_tol", 2e-2)) + + if "items" in d: + for i, item in enumerate(d["items"]): + _run_one( + i, + item["so_path"], + item["problem"], + item.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + else: + _run_one( + 0, + d["so_path"], + d["problem"], + d.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + + +if __name__ == "__main__": + main() diff --git a/tile_engine/ops/gemm/streamk_gemm_full_benchmark.py b/tile_engine/ops/gemm/streamk_gemm_full_benchmark.py new file mode 100644 index 0000000000..84289922c4 --- /dev/null +++ b/tile_engine/ops/gemm/streamk_gemm_full_benchmark.py @@ -0,0 +1,492 @@ +#!/usr/bin/env python3 +"""Full Stream-K GEMM benchmark sweep driven through the Dispatcher bridge. + +Same 3-phase architecture as ``gemm_full_benchmark.py`` -- Stream-K is a +single-problem GEMM (one A/B/C, one M/N/K) with the *same* C ABI, so the only +bridge difference is ``variant="stream_k"`` threaded into ``expand_sweep`` (which +makes the dispatcher codegen the Stream-K launch and ``gemm_utils`` select +``streamk_gemm_ctypes_lib.cpp``). The GPU worker reuses ``GemmProblem`` / +``GpuGemmRunner`` unchanged. + + Phase 1: Compile all kernels (parallel, returns .so paths only -- no GPU) + Phase 2: Load problems (M, N, K shapes) + Phase 3: Benchmark via subprocess isolation, distributed across all visible + GPUs (one device-pinned worker per GPU, batched, fault-isolated) + +Like the regular bridge driver, Phase 3 fans the (kernel x problem) work out +across every visible GPU in parallel: each device runs its own stream of +disposable worker subprocesses pinned with ``HIP_VISIBLE_DEVICES``, so an N-GPU +box benchmarks roughly N times faster while keeping per-batch fault isolation. + +Examples: + # Default config (gemm_streamk/configs/default_config.json), all visible GPUs: + python streamk_gemm_full_benchmark.py + + # Explicit config on 4 GPUs with correctness checking: + python streamk_gemm_full_benchmark.py gemm_streamk/configs/default_config.json \ + --devices 4 --verify --csv streamk_gemm_results.csv +""" + +import argparse +import csv +import json +import os +import queue +import re +import subprocess +import sys +import threading +import time +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_DISPATCHER_ROOT = _THIS_DIR.parents[2] / "dispatcher" +sys.path.insert(0, str(_DISPATCHER_ROOT / "python")) +sys.path.insert(0, str(_THIS_DIR)) + +from gemm_utils import setup_multiple_gemm_dispatchers, expand_sweep # noqa: E402 + +# Stream-K is a single variant; its sweep configs live in gemm_streamk/configs/. +DEFAULT_CONFIG = _THIS_DIR / "gemm_streamk" / "configs" / "default_config.json" + +# Default problem set: squares plus a large-K skinny shape -- Stream-K's sweet +# spot is few output tiles with a long K reduction. Tiny problems (e.g. 257^3) +# have too few tiles to partition across CUs and the kernel reports them as +# unsupported (status -2), which the bridge surfaces gracefully. +DEFAULT_PROBLEMS = [ + {"M": 1024, "N": 1024, "K": 1024}, + {"M": 2048, "N": 2048, "K": 2048}, + {"M": 4096, "N": 4096, "K": 4096}, + {"M": 512, "N": 512, "K": 8192}, +] + +# Bridge surface for Stream-K. The dispatcher host path +# (streamk_gemm_ctypes_lib.cpp) derives strides from the kernel's layouts and the +# worker (run_one_streamk_gemm_kernel.py) reads dtype/layout off the kernel name, +# so all 4 A/B/C layouts are supported. dtypes cover fp16 + bf16 + fp8 + bf8: the +# bridge runner encodes fp16 natively, bf16 via bit-truncation, and fp8/bf8 via +# ml_dtypes in the gfx942 FNUZ formats (e4m3fnuz / e5m2fnuz), which accumulate +# into fp16. int8 is left out: it is blocked at the ck_tile engine level, not the +# bridge -- the int8 kernel codegens but fails to COMPILE for every reduction +# strategy (atomic/linear/tree). warp_gemm_dispatcher has no +# Dispatcher specialization for the streamk CompV3 +# path, so WarpGemm resolves to `int` and the BlockUniversalGemmAsBsCr +# WarpGemm::kM/kN static_asserts fail; this matches PR #8094 leaving int8 out. +SUPPORTED_DTYPES = ("fp16", "bf16", "fp8", "bf8") +SUPPORTED_LAYOUTS = ("rcr", "rrr", "ccr", "crr") + + +def detect_devices(): + """Return a list of visible GPU id strings (best-effort).""" + env = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get( + "CUDA_VISIBLE_DEVICES" + ) + if env: + ids = [d.strip() for d in env.split(",") if d.strip() != ""] + if ids: + return ids + try: + out = subprocess.check_output( + ["rocm-smi", "--showid"], stderr=subprocess.DEVNULL, text=True + ) + ids = sorted(set(re.findall(r"GPU\[(\d+)\]", out)), key=int) + if ids: + return ids + except Exception: + pass + try: + out = subprocess.check_output( + ["amd-smi", "list"], stderr=subprocess.DEVNULL, text=True + ) + ids = re.findall(r"^GPU:\s*(\d+)", out, re.MULTILINE) + if ids: + return ids + except Exception: + pass + return ["0"] + + +def resolve_devices(spec): + """Resolve --devices into a concrete list of device id strings. + + spec is None (auto: all visible), an int count, or a comma-list of ids. + A bare digit is a *count*, not an id; to target one specific id use the + comma form, e.g. "5,". + """ + detected = detect_devices() + if spec is None: + return detected + spec = str(spec).strip() + if "," in spec: + return [s.strip() for s in spec.split(",") if s.strip() != ""] + if spec.isdigit(): + n = int(spec) + if n <= 0: + return detected + return detected[:n] if len(detected) >= n else [str(i) for i in range(n)] + return [spec] + + +def resolve_configs(args): + """Resolve positional configs -> concrete list of config paths.""" + if args.configs: + return args.configs + return [str(DEFAULT_CONFIG)] + + +def load_problems(path): + if not path: + return DEFAULT_PROBLEMS + with open(path) as f: + data = json.load(f) + # Accept either a bare list or {"problems": [...]}. + return data["problems"] if isinstance(data, dict) else data + + +def _run_batch_on_device(device_id, unit, args, worker_path, base_env): + """Run one (problem, kernel-batch) unit in a device-pinned subprocess. + + Returns (rows, lines, n_fail) where rows are dicts ready for the CSV writer, + lines are formatted strings to print, and n_fail counts failures. + """ + prob_idx, prob_dict, batch = unit + M, N, K = prob_dict["M"], prob_dict["N"], prob_dict["K"] + + items = [ + {"so_path": str(lib), "problem": prob_dict, "kernel_name": cfg.name} + for _, cfg, lib in batch + ] + payload = json.dumps( + {"items": items, "verify": args.verify, "verify_tol": args.verify_tol} + ) + + env = base_env.copy() + env["HIP_VISIBLE_DEVICES"] = str(device_id) + + rows, lines, n_fail = [], [], 0 + proc = None + try: + proc = subprocess.Popen( + [sys.executable, str(worker_path)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + env=env, + ) + stdout_bytes, _ = proc.communicate( + input=payload.encode("utf-8"), + timeout=args.kernel_timeout * len(batch), + ) + + reported = set() + for line in stdout_bytes.decode("utf-8").strip().split("\n"): + if not line: + continue + try: + result = json.loads(line) + except json.JSONDecodeError: + lines.append(f" [gpu{device_id}] Warning: bad result line: {line[:50]}") + n_fail += 1 + continue + bidx = result.get("idx", 0) + _, cfg, _ = batch[bidx] + reported.add(bidx) + if result.get("ok", False): + status = "OK" if result.get("non_zero", 0) > 0 else "ZERO" + mismatch = False + if args.verify and "verified" in result: + if result["verified"]: + status = "VERIFY" + else: + status = "MISMATCH" + mismatch = True + extra = f" rel={result['max_rel']:.2e}" if "max_rel" in result else "" + lines.append( + f" [gpu{device_id}] {cfg.name:<58} {result['ms']:>10.3f} " + f"{result['tflops']:>10.2f} {status:>8}{extra}" + ) + rows.append( + { + "kernel": cfg.name, + "problem_idx": prob_idx, + "M": M, + "N": N, + "K": K, + "device": device_id, + "latency_ms": result["ms"], + "tflops": result["tflops"], + "non_zero": result.get("non_zero", 0), + "max_rel": result.get("max_rel", ""), + "verified": result.get("verified", ""), + } + ) + if mismatch: + n_fail += 1 + else: + lines.append(f" [gpu{device_id}] {cfg.name:<58} FAILED") + lines.append(f" Error: {result.get('error', 'unknown')[:100]}") + n_fail += 1 + + missing = set(range(len(batch))) - reported + if missing or proc.returncode != 0: + if proc.returncode != 0: + lines.append(f" [gpu{device_id}] worker exited code {proc.returncode}") + for idx in sorted(missing): + _, cfg, _ = batch[idx] + lines.append(f" [gpu{device_id}] {cfg.name:<58} MISSING (crash)") + n_fail += len(missing) + + except subprocess.TimeoutExpired: + lines.append(f" [gpu{device_id}] batch timeout ({len(batch)} kernels)") + try: + proc.kill() + proc.communicate(timeout=5) + except Exception: + pass + n_fail += len(batch) + except Exception as e: + lines.append(f" [gpu{device_id}] batch error: {e}") + try: + if proc and proc.poll() is None: + proc.kill() + except Exception: + pass + n_fail += len(batch) + + return rows, lines, n_fail + + +def main(): + parser = argparse.ArgumentParser( + description="Stream-K GEMM Benchmark Sweep (via Dispatcher)" + ) + parser.add_argument( + "configs", + nargs="*", + help="TE sweep config JSON files (default: gemm_streamk/configs/default_config.json)", + ) + parser.add_argument("--arch", default="gfx942") + parser.add_argument( + "--dtype", + default="fp16", + choices=SUPPORTED_DTYPES, + help=f"Input dtype (supported: {', '.join(SUPPORTED_DTYPES)})", + ) + parser.add_argument( + "--layout", + default="rcr", + choices=SUPPORTED_LAYOUTS, + help=f"A/B/C layout (supported: {', '.join(SUPPORTED_LAYOUTS)})", + ) + parser.add_argument("--problems", default=None, help="JSON file of M,N,K problems") + parser.add_argument("--csv", type=str, default="streamk_gemm_results.csv") + parser.add_argument("--workers", type=int, default=8, help="Parallel build workers") + parser.add_argument( + "--devices", + default=None, + help="GPUs to use: int count (e.g. 4) or comma-list of ids (e.g. 0,2,5); " + "for one specific id use the comma form (e.g. 5,) since a bare digit is " + "a count; default auto-detects all visible", + ) + parser.add_argument( + "--batch-size", + type=int, + default=20, + help="Kernels per subprocess (overhead vs fault isolation)", + ) + parser.add_argument( + "--kernel-timeout", type=int, default=30, help="Per-kernel timeout (s)" + ) + parser.add_argument( + "--max-kernels", type=int, default=0, help="Limit to first N kernels (0=all)" + ) + parser.add_argument( + "--verify", + action="store_true", + help="Check each kernel's output against an fp32 numpy reference " + "(global max|out-ref|/max|ref|); a mismatch counts as a failure", + ) + parser.add_argument( + "--verify-tol", + type=float, + default=2e-2, + help="Relative tolerance for --verify (default 2e-2; Stream-K's Atomic " + "reduction is noisier than regular GEMM but stays well under this)", + ) + args = parser.parse_args() + + config_paths = resolve_configs(args) + devices = resolve_devices(args.devices) + + # ======================================================================== + # Phase 1: Compile kernels (parallel, no GPU) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 1: Compile Stream-K kernels") + print(f"{'=' * 80}") + print(f" Configs: {', '.join(config_paths)}") + + all_configs = [] + for cfg_path in config_paths: + all_configs.extend( + expand_sweep( + cfg_path, + args.arch, + dtype=args.dtype, + layout=args.layout, + variant="stream_k", + ) + ) + + if args.max_kernels > 0: + all_configs = all_configs[: args.max_kernels] + + print(f" Expanded configs: {len(all_configs)}") + print(f" Build workers: {args.workers}") + + t0 = time.perf_counter() + # CRITICAL: returns Path objects only, does NOT load any .so. + lib_paths = setup_multiple_gemm_dispatchers( + all_configs, verbose=True, max_workers=args.workers + ) + build_time = time.perf_counter() - t0 + + built_kernels = [ + (cfg, lib) for cfg, lib in zip(all_configs, lib_paths) if lib is not None + ] + + # Dedupe by .so path (distinct configs can map to the same physical kernel). + seen_libs = set() + unique_kernels = [] + duplicate_count = 0 + for cfg, lib in built_kernels: + lib_key = str(lib.resolve()) + if lib_key not in seen_libs: + seen_libs.add(lib_key) + unique_kernels.append((cfg, lib)) + else: + duplicate_count += 1 + built_kernels = unique_kernels + + print( + f"\n Built {len(all_configs)} configs -> {len(built_kernels)} unique kernels " + f"({duplicate_count} duplicates filtered) in {build_time:.0f}s" + ) + + if not built_kernels: + print(" ERROR: No kernels built successfully") + return 1 + + # ======================================================================== + # Phase 2: Load problems + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 2: Load test problems") + print(f"{'=' * 80}") + + problems = load_problems(args.problems) + print(f" Problems: {len(problems)}") + print( + f" Total measurements: {len(built_kernels)} x {len(problems)} = " + f"{len(built_kernels) * len(problems)}" + ) + + # ======================================================================== + # Phase 3: Benchmark across all visible GPUs (subprocess isolation, batched) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 3: Benchmark (multi-GPU, subprocess isolation, batched)") + print(f"{'=' * 80}") + print(f" Devices: {len(devices)} -> {', '.join(devices)}") + print(f" Batch size: {args.batch_size} kernels per subprocess") + print(f" Timeout: {args.kernel_timeout}s per kernel\n") + + csv_path = Path(args.csv) + csv_fields = [ + "kernel", + "problem_idx", + "M", + "N", + "K", + "device", + "latency_ms", + "tflops", + "non_zero", + "max_rel", + "verified", + ] + csv_file = open(csv_path, "w", newline="") + writer = csv.DictWriter(csv_file, fieldnames=csv_fields) + writer.writeheader() + + worker_path = _THIS_DIR / "run_one_streamk_gemm_kernel.py" + base_env = os.environ.copy() + base_env["GEMM_PYPATH"] = os.pathsep.join( + [str(_DISPATCHER_ROOT / "python"), str(_THIS_DIR)] + ) + + # Build a single work queue of (prob_idx, prob_dict, kernel-batch) units and + # fan them out across device-pinned worker threads. + work_q = queue.Queue() + for prob_idx, prob in enumerate(problems): + prob_dict = {"M": int(prob["M"]), "N": int(prob["N"]), "K": int(prob["K"])} + for start in range(0, len(built_kernels), args.batch_size): + end = min(start + args.batch_size, len(built_kernels)) + batch = [ + (start + j, cfg, lib) + for j, (cfg, lib) in enumerate(built_kernels[start:end]) + ] + work_q.put((prob_idx, prob_dict, batch)) + + io_lock = threading.Lock() + stats = {"measurements": 0, "failures": 0} + bench_t0 = time.perf_counter() + + def device_thread(device_id): + while True: + try: + unit = work_q.get_nowait() + except queue.Empty: + return + rows, lines, n_fail = _run_batch_on_device( + device_id, unit, args, worker_path, base_env + ) + with io_lock: + for ln in lines: + print(ln) + for row in rows: + writer.writerow(row) + csv_file.flush() + stats["measurements"] += len(rows) + stats["failures"] += n_fail + work_q.task_done() + + threads = [ + threading.Thread(target=device_thread, args=(d,), daemon=True) for d in devices + ] + for t in threads: + t.start() + for t in threads: + t.join() + + bench_time = time.perf_counter() - bench_t0 + csv_file.close() + + # ======================================================================== + # Summary + # ======================================================================== + print(f"\n{'=' * 80}") + print("BENCHMARK COMPLETE") + print(f"{'=' * 80}") + print(f" Build time: {build_time:.0f}s") + print(f" Benchmark time: {bench_time:.0f}s") + print(f" Total time: {build_time + bench_time:.0f}s") + print(f" Devices used: {len(devices)}") + print(f" Successful measurements: {stats['measurements']}") + print(f" Failed measurements: {stats['failures']}") + print(f" Output: {csv_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main())