ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 17:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#9028 (commit 2d6a3d6)

Browse Source 
feat(ck-tile): stream-K GEMM TE to dispatcher bridge
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

> Re-opened from #8136 with a policy-compliant branch name. Supersedes
#8136.

## Summary

Routes the **stream_k** GEMM variant through the same Tile Engine (TE) →
Dispatcher bridge already landed for regular GEMM (#8123) and grouped
GEMM (#8130). The Dispatcher stays the single source of truth for
codegen/build/runtime; TE only produces configs and benchmarks.

## Design note

Stream-K is a single-problem GEMM with the **same C ABI** as regular
GEMM, so the Python runner (`GpuGemmRunner`/`GemmProblem`) and the GPU
worker are reused unchanged. The one twist: the registry path can't
compile against a Stream-K `SelectedKernel`, so the Stream-K ctypes lib
**bypasses the registry** and calls `SelectedKernel::launch(args,
stream)` directly — the same approach grouped GEMM uses. The generated
launch owns the reduction workspace internally (`DeviceMem`) and uses
the Atomic strategy.

## Changes

**New**
- `bindings/ctypes/streamk_gemm_ctypes_lib.cpp` — single-problem C ABI
(`dispatcher_run_gemm`), builds `StreamKHostArgs`, direct launch;
returns `0` / `-1` (HIP/throw) / `-2` (args unsupported).
- `tile_engine/ops/gemm/streamk_gemm_full_benchmark.py` +
`run_one_streamk_gemm_kernel.py` — 3-phase driver (expand → build →
subprocess-isolated benchmark) and disposable GPU worker.
- `tile_engine/ops/gemm/gemm_streamk/configs/default_config.json` —
small sweep config.

**Modified**
- `dispatcher/python/gemm_utils.py`, `ctypes_utils.py` — thread
`variant="stream_k"` through codegen/build and `.so` selection.

## Validation

fp16/rcr on gfx942/MI300X: numeric parity vs an fp32 numpy reference
(widened fp16-atomic tolerance) and a full driver run of **16/16 OK**
with end-to-end name parity; unsupported tiny shapes are reported
gracefully (`status -2`), not crashes. fp8/bf8/bf16 now supported via
`ml_dtypes` FNUZ codecs. Full tables and the bridge-vs-old-TE comparison
are in the comments.

## Test plan

- [x] codegen emits `*_streamk.hpp` with stem ==
`GemmKernelConfig(variant="stream_k").name`
- [x] build/link against `streamk_gemm_ctypes_lib.cpp`
- [x] numeric parity passes (fp16 atomic tolerance)
- [x] full driver run 16/16 OK, name parity end-to-end
- [x] unsupported shape → `status -2` handled gracefully

## Next

Land #8123, then this; afterwards remove the legacy
`tile_engine/ops/gemm_streamk/` machinery (Phase 4).
This commit is contained in:
Muhammed Emin Ozturk
2026-07-17 04:08:05 +00:00
committed by assistant-librarian[bot]
parent c401998b69
commit 58c1bcbe2e
9 changed files with 1696 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

290
dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp Normal file
View File

@@ -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 <hip/hip_runtime.h>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <string>
#include <type_traits>
// 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<int>(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<char>(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<const ADataType*>(A);
const BDataType* B_host = static_cast<const BDataType*>(B);
CDataType* C_host = static_cast<CDataType*>(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<ck_tile::index_t>(std::is_same_v<ALayout, RowMajor> ? K : M);
const ck_tile::index_t ldb =
static_cast<ck_tile::index_t>(std::is_same_v<BLayout, RowMajor> ? N : K);
const ck_tile::index_t ldc =
static_cast<ck_tile::index_t>(std::is_same_v<CLayout, RowMajor> ? N : M);
// k_batch is fixed to 1 inside StreamKHostArgs.
ck_tile::StreamKHostArgs args(static_cast<const void*>(A_dev),
static_cast<const void*>(B_dev),
static_cast<void*>(C_dev),
static_cast<ck_tile::index_t>(M),
static_cast<ck_tile::index_t>(N),
static_cast<ck_tile::index_t>(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<size_t>(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"

163
dispatcher/parity_diag/regression/ab_efficient_sweep.py Normal file
View File

@@ -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()

310
dispatcher/parity_diag/regression/ab_same_harness.py Normal file
View File

@@ -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_<stem>). 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
``<OLD_TE_GEN_BASE>/<dtype>/<layout>/`` (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 <stem> [<stem>...] # 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/<this>
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
# (<dtype>/<layout>) 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: <base>/<dtype>/<layout> (or the pinned dir).
Stems are named ``<dtype>_<layout>_...`` (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()

2
dispatcher/python/ctypes_utils.py
View File

@@ -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

77
dispatcher/python/gemm_utils.py
View File

@@ -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

108
dispatcher/tests/test_streamk_gemm_utils.py Normal file
View File

@@ -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_<strategy>``. 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()

89
tile_engine/ops/gemm/gemm_streamk/configs/default_config.json Normal file
View File

@@ -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
}

175
tile_engine/ops/gemm/run_one_streamk_gemm_kernel.py Normal file
View File

@@ -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()

492
tile_engine/ops/gemm/streamk_gemm_full_benchmark.py Normal file
View File

@@ -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<int8,int8,float,32,32,16,...> 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())