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

[rocm-libraries] ROCm/rocm-libraries#9000 (commit 9faa8de)

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

> Re-opened from #8130 with a policy-compliant branch name
(`users/muozturk/ck-tile/dispatcher-te-bridge-grouped-gemm`). Supersedes
#8130.

## What this PR does

Routes the **grouped_gemm** variant through the Tile Engine (TE) →
Dispatcher **bridge**: TE only generates configs and benchmarks; the
Dispatcher owns codegen, build, and runtime. This is the grouped
counterpart of the regular-GEMM bridge (#8123/#8479), the fp8/bf8/int8
bridge (#8887), and the Stream-K bridge (#8136).

**This PR now also contains the grouped Dispatcher codegen** that
previously lived in #8075 — that PR has been **closed in favor of this
one** to keep the grouped codegen in a single place (it was otherwise
duplicated across both).

## Why grouped needs special handling

Grouped GEMM is **multi-problem**: one launch runs a *list* of `(M, N,
K)` sub-problems with arrays of A/B/C device pointers.

1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`)
cannot express a list of problems.
2. The generated registry wrapper (`generated_tile_backend.hpp::run()`)
hard-codes the single-problem launch and won't compile against a grouped
`SelectedKernel`.

So the grouped path **bypasses the registry**: a dedicated ctypes lib
calls the generated `SelectedKernel::launch(descs, stream)` directly and
reports the name from the compile-time `KERNEL_NAME` macro.

## Changes

**Codegen (absorbed from #8075)**
- `codegen/arch_filter.py` — `GEMM_GROUPED` operator tile constraints.
- `codegen/unified_gemm_codegen.py` — `GemmVariant.GROUPED`, the grouped
launch generator (DeviceMem internal workspace via `MakeKargs`,
persistent/non-persistent grid), `grouped` in `--variants`.
- `examples/gemm/cpp/02_grouped_gemm_driver.cpp` — standalone,
layout/dtype-generic grouped driver with per-group reference
verification.
- `codegen/README.md` + `examples/gemm/cpp/README.md` — grouped
sections.

**Bridge**
- `bindings/ctypes/grouped_gemm_ctypes_lib.cpp` — multi-problem,
registry-bypass C ABI; per-group device alloc/copy; strides derived from
the compile-time `ALayout/BLayout/CLayout`; warmup/repeat timing matched
to Old-TE (`CK_TILE_BENCH_WARMUP/REPEAT`).
- `python/gemm_utils.py` — `GroupedGemmProblem`/`GroupedGemmResult`,
`GpuGroupedGemmRunner`, `run_grouped`, fp16/bf16/fp8(E4M3 FNUZ)/bf8(E5M2
FNUZ) codecs, output-dtype-aware C buffer.
- `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` +
`run_one_grouped_gemm_kernel.py` — TE driver + worker for the parity
sweep.
- `bindings/ctypes/GROUPED_GEMM_BRIDGE.md` — design README.

## Coverage (= Old-TE grouped runnable set on develop)

| Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) |
|---|---|---|---|---|
| rcr / rrr / ccr / crr | ✓ | ✓ | ✓ | ✓ |

C is always row-major. `int8` (rejected by the TE grouped builder) and
`fp32`/`fp64` (no MFMA warp tiles) are excluded on both sides.

## Parity vs Old-TE (MI300X / gfx942)

Apples-to-apples (same warmup=50/repeat=100 both sides, A/B interleaved,
single GPU, both engines rebuilt fresh, stale-`.so` guard, matched
compile flags):

- **Correctness: 64/64 PASS.**
- **Performance: 64/64 within ±15%.**
- The 5 small-shape (1024³ fp8/bf8) rows that initially read >15% were
proven by `rocprof` to be a **measurement-harness artifact** (Old-TE's
JSON `latency(ms)` rounded to 2 decimals → 30–50% TFLOPS swing on ~0.02
ms kernels), **not** a kernel/codegen difference — bridge and Old-TE
launch byte-identical kernels (same grid/VGPR/SGPR, duration ≤3.22%);
full-precision re-measure collapses all 5 to <3%.

## Notes

- Targets `develop`. Depends on #8997 (fp16/bf16 bridge) and #8998
(fp8/bf8/int8 bridge) merging to `develop` first; until then this PR's
diff also shows their content, after which it reduces to the
grouped-only files.
- Supersedes #8075 (closed).
This commit is contained in:
Muhammed Emin Ozturk
2026-07-16 02:55:42 +00:00
committed by assistant-librarian[bot]
parent a6028c883b
commit 6648115aed
12 changed files with 1654 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
dispatcher/bindings/README.md
View File

@@ -8,6 +8,7 @@ This directory contains language bindings for the CK Tile Dispatcher.
bindings/
|---- ctypes/ # Python ctypes bindings (C API)
| |---- gemm_ctypes_lib.cpp # GEMM dispatcher C API
| |---- grouped_gemm_ctypes_lib.cpp # Grouped (multi-problem) GEMM bridge C API -- see GROUPED_GEMM_BRIDGE.md
| |---- conv_ctypes_lib.cpp # Grouped conv dispatcher C API (fwd + bwd_data)
| |---- conv_bwdw_ctypes_lib.cpp # Grouped conv backward weight C API (separate library)
| |---- fmha_ctypes_lib.cpp # FMHA dispatcher C API (fwd + bwd)

121
dispatcher/bindings/ctypes/GROUPED_GEMM_BRIDGE.md Normal file
View File

@@ -0,0 +1,121 @@
<!--
Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
SPDX-License-Identifier: MIT
-->
# Grouped GEMM: Tile Engine -> Dispatcher Bridge
This document describes the **grouped_gemm** variant of the Tile Engine (TE) ->
Dispatcher bridge (PR #8130). It is the grouped counterpart of the regular-GEMM
bridge (#8123/#8479), the fp8/bf8/int8 bridge (#8887), and the Stream-K bridge
(#8136).
## What the bridge is
In the bridge model the **Dispatcher is the single source of truth** for
codegen, build, and runtime; **Tile Engine only generates configs and
benchmarks** them. TE no longer carries its own kernel-instance build path —
it shells out to the dispatcher codegen and runs the resulting kernel.
For most variants the dispatcher runs a kernel through its registry/backend.
Grouped GEMM cannot use that path (see below), so the grouped bridge takes the
same approach as Stream-K: a dedicated ctypes library that **bypasses the
registry** and calls the generated `SelectedKernel::launch(...)` directly.
## Why grouped needs special handling
Grouped GEMM is **multi-problem**: a single launch runs a *list* of `(M, N, K)`
sub-problems, each with its own A/B/C device pointers. Two consequences:
1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`) cannot
express a list of problems.
2. The generated registry wrapper (`generated_tile_backend.hpp::run()`)
hard-codes the single-problem `SelectedKernel::launch(GemmHostArgs, ...)`
signature and will not compile against a grouped `SelectedKernel`.
So the grouped kernel header exposes a different launch signature
```cpp
static float launch(const std::vector<ck_tile::GroupedGemmHostArgs<>>& descs,
const stream_config& stream);
```
and the grouped ctypes lib force-includes one generated kernel header
(`-include ..._grouped.hpp` with `CK_TILE_SINGLE_KERNEL_INCLUDE`), calls that
`launch` directly, and reports the kernel name from the compile-time
`KERNEL_NAME` macro.
## Components
| Layer | File | Role |
|---|---|---|
| Codegen | `dispatcher/codegen/unified_gemm_codegen.py` | `GemmVariant.GROUPED`; `_launch_function_grouped` (DeviceMem internal workspace, `MakeKargs`, persistent/non-persistent grid). Kept in lockstep with PR #8075. |
| Codegen | `dispatcher/codegen/arch_filter.py` | `GEMM_GROUPED` operator tile constraints. |
| C API | `dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp` | Multi-problem ABI; per-group device alloc/copy; layout-derived strides; warmup/repeat timing. |
| Python | `dispatcher/python/gemm_utils.py` | `GroupedGemmProblem` / `GroupedGemmResult`, `GpuGroupedGemmRunner`, `run_grouped`, `build_grouped`, dtype/layout codecs. |
| Python | `dispatcher/python/ctypes_utils.py` | Threads the `grouped` variant into the codegen `--variants` flag. |
| TE driver | `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` | Generates configs, builds `.so`s in parallel, benchmarks in disposable workers. |
| TE worker | `tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py` | Runs one grouped kernel; dtype/layout-aware operand generation. |
## C ABI
```c
int dispatcher_init(void); // lightweight no-op (no registry)
int dispatcher_run_grouped_gemm(
int group_count,
const int64_t* Ms, // [group_count]
const int64_t* Ns, // [group_count]
const int64_t* Ks, // [group_count]
const void** A_ptrs, // host A buffers, one per group
const void** B_ptrs, // host B buffers, one per group
void** C_ptrs, // host C out buffers, one per group
float* time_ms); // out: average kernel time
// returns 0 ok, -1 HIP/throw, -2 arguments unsupported by the kernel
```
The lib `hipMalloc`s A/B/C per group, copies A and B host->device, memsets C,
builds `std::vector<ck_tile::GroupedGemmHostArgs<>>` with **strides derived from
the compile-time `ALayout`/`BLayout`/`CLayout`** of the `-include`d header
(`std::is_same_v<…, RowMajor>`), launches once, then copies each C back. The ABI
is `void*` + element-size, so it is dtype-agnostic; the Python runner owns the
numpy codecs.
## Coverage
The bridge runnable set is exactly the Old-TE grouped_gemm runnable set on
`develop` — no more, no less:
| Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) |
|---|---|---|---|---|
| rcr | ✓ | ✓ | ✓ | ✓ |
| rrr | ✓ | ✓ | ✓ | ✓ |
| ccr | ✓ | ✓ | ✓ | ✓ |
| crr | ✓ | ✓ | ✓ | ✓ |
- **Matrix C is always row-major** (grouped builder constraint), so the layout
string varies A/B only.
- **Excluded:** `int8` (rejected by the TE grouped builder), `fp32`/`fp64`
(no MFMA warp tiles). These are excluded on both sides.
- fp8/bf8 use the **FNUZ** encoding on gfx942 (matches the regular #8887 path);
the Python codecs require `ml_dtypes`.
## Building and running
Generate + build one grouped `.so` and run the A/B parity sweep vs Old-TE:
```bash
# Codegen smoke (no GPU): one variant/dtype/layout
python3 dispatcher/codegen/unified_gemm_codegen.py \
--output-dir /tmp/grp --datatype bf16 --layout ccr \
--variants grouped --config dispatcher/codegen/default_config.json
# Full TE-driven parity sweep (build + benchmark)
python3 tile_engine/ops/gemm/grouped_gemm_full_benchmark.py <config.json> \
--arch gfx942 --dtype fp16 --layout rcr --csv grouped_results.csv
```
Timing knobs `CK_TILE_BENCH_WARMUP` (default 50) and `CK_TILE_BENCH_REPEAT`
(default 100) are honored by **both** the grouped ctypes lib and the registry
backend, so bridge-vs-Old-TE A/B comparisons stay matched. For fair parity keep
`flush_cache=false`, `rotating_count=1`, run on a single GPU, and re-measure any
`|gap|>15%` outlier standalone.

295
dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp Normal file
View File

@@ -0,0 +1,295 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Grouped GEMM Dispatcher ctypes Library
*
* Provides C API for Python ctypes integration for the GROUPED GEMM variant.
* Kernel header included via -include at compile time.
*
* The grouped kernel has a genuinely different ABI from regular GEMM: it takes a
* LIST of (M,N,K) sub-problems plus arrays of A/B/C device pointers, and its
* generated launch() builds the per-group arg workspace internally:
*
* static float launch(const std::vector<ck_tile::GroupedGemmHostArgs<>>& descs,
* const stream_config& stream);
*
* The single-problem dispatcher run path (g_dispatcher->run / GemmHostArgs) cannot
* express this, and the generated_tile_backend wrapper hard-codes the single-problem
* launch signature, so this lib calls SelectedKernel::launch(descs, stream) directly
* and reports the kernel name from the compile-time KERNEL_NAME macro instead of the
* registry.
*
* Usage from Python:
* lib = ctypes.CDLL("libdispatcher_grouped_gemm.so")
* lib.dispatcher_init()
* lib.dispatcher_run_grouped_gemm(...)
*/
#include <hip/hip_runtime.h>
#include <array>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <string>
#include <type_traits>
#include <vector>
// 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::GroupedGemmHostArgs 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. Mirrors generated_tile_backend.hpp so
// both bridge sides honor the same CK_TILE_BENCH_* env vars.
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);
}
extern "C" {
/**
* Initialize the grouped GEMM library.
*
* The grouped 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 grouped GEMM on GPU by launching the force-included kernel directly.
*
* For each group: hipMalloc A/B/C, copy A and B host->device, memset C, then build
* a std::vector<ck_tile::GroupedGemmHostArgs<>> with strides derived from the
* compile-time ALayout/BLayout/CLayout of the -include'd kernel header (k_batch=1)
* and launch. After the launch the per-group C buffers are copied back to the
* caller's host buffers.
*
* Layout contract: A is MxK, B is KxN, C is MxN; leading dimensions follow each
* operand's row/col-major layout (CLayout is always RowMajor for grouped).
*
* Returns: 0 on success, -1 on HIP error / generic throw, -2 if the kernel reports
* the arguments are unsupported.
*/
int dispatcher_run_grouped_gemm(int group_count,
const int64_t* Ms,
const int64_t* Ns,
const int64_t* Ks,
const void** A_ptrs,
const void** B_ptrs,
void** C_ptrs,
float* time_ms)
{
if(!g_initialized || group_count <= 0 || !Ms || !Ns || !Ks || !A_ptrs || !B_ptrs || !C_ptrs)
{
return -1;
}
std::vector<ADataType*> A_dev(group_count, nullptr);
std::vector<BDataType*> B_dev(group_count, nullptr);
std::vector<CDataType*> C_dev(group_count, nullptr);
auto cleanup_gpu_mem = [&]() {
for(int g = 0; g < group_count; ++g)
{
if(A_dev[g])
(void)hipFree(A_dev[g]);
if(B_dev[g])
(void)hipFree(B_dev[g]);
if(C_dev[g])
(void)hipFree(C_dev[g]);
}
};
std::vector<ck_tile::GroupedGemmHostArgs<>> descs;
descs.reserve(group_count);
for(int g = 0; g < group_count; ++g)
{
const int64_t M = Ms[g];
const int64_t N = Ns[g];
const int64_t K = Ks[g];
if(M <= 0 || N <= 0 || K <= 0 || !A_ptrs[g] || !B_ptrs[g] || !C_ptrs[g])
{
cleanup_gpu_mem();
return -1;
}
if(hipMalloc(&A_dev[g], M * K * sizeof(ADataType)) != hipSuccess)
{
cleanup_gpu_mem();
return -1;
}
if(hipMalloc(&B_dev[g], K * N * sizeof(BDataType)) != hipSuccess)
{
cleanup_gpu_mem();
return -1;
}
if(hipMalloc(&C_dev[g], M * N * sizeof(CDataType)) != hipSuccess)
{
cleanup_gpu_mem();
return -1;
}
if(hipMemcpy(A_dev[g], A_ptrs[g], M * K * sizeof(ADataType), hipMemcpyHostToDevice) !=
hipSuccess)
{
cleanup_gpu_mem();
return -1;
}
if(hipMemcpy(B_dev[g], B_ptrs[g], K * N * sizeof(BDataType), hipMemcpyHostToDevice) !=
hipSuccess)
{
cleanup_gpu_mem();
return -1;
}
if(hipMemset(C_dev[g], 0, M * N * sizeof(CDataType)) != hipSuccess)
{
cleanup_gpu_mem();
return -1;
}
// Derive leading dimensions from the compile-time layouts the kernel was
// generated with (ALayout/BLayout/CLayout from the -include'd header),
// matching Old-TE gemm_validation_utils.get_abc_layouts:
// stride_A = ALayout row-major ? K : M
// stride_B = BLayout row-major ? N : K
// stride_E = CLayout row-major ? N : M (CLayout is always RowMajor for grouped)
using RowMajor = ck_tile::tensor_layout::gemm::RowMajor;
const auto stride_A = std::is_same_v<ALayout, RowMajor> ? static_cast<ck_tile::index_t>(K)
: static_cast<ck_tile::index_t>(M);
const auto stride_B = std::is_same_v<BLayout, RowMajor> ? static_cast<ck_tile::index_t>(N)
: static_cast<ck_tile::index_t>(K);
const auto stride_E = std::is_same_v<CLayout, RowMajor> ? static_cast<ck_tile::index_t>(N)
: static_cast<ck_tile::index_t>(M);
// k_batch=1 for numeric parity.
descs.emplace_back(static_cast<const void*>(A_dev[g]),
static_cast<const void*>(B_dev[g]),
std::array<const void*, 0>{},
static_cast<void*>(C_dev[g]),
/*k_batch=*/1,
static_cast<ck_tile::index_t>(M),
static_cast<ck_tile::index_t>(N),
static_cast<ck_tile::index_t>(K),
stride_A,
stride_B,
std::array<ck_tile::index_t, 0>{},
stride_E);
}
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_ = false;
stream_cfg.rotating_count_ = 1;
float exec_time = 0.0f;
try
{
exec_time = SelectedKernel::launch(descs, 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;
}
// Copy each group's result back to host.
for(int g = 0; g < group_count; ++g)
{
const int64_t M = Ms[g];
const int64_t N = Ns[g];
if(hipMemcpy(C_ptrs[g], C_dev[g], 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 grouped .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 grouped 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"

25
dispatcher/codegen/README.md
View File

@@ -77,7 +77,7 @@ results = codegen.generate_all()
| `--datatype` | `fp16`, `bf16`, `fp32`, `int8` | Data type |
| `--layout` | `rcr`, `rrr`, `crr`, `ccr` | Matrix layouts |
| `--gpu-target` | `gfx942`, `gfx90a`, `gfx950` | Target GPU |
| `--variants` | `standard`, `preshuffle`, `multi_d` | Kernel variants |
| `--variants` | `standard`, `preshuffle`, `multi_d`, `grouped` | Kernel variants |
| `--preselected` | `fp16_rcr_essential`, etc. | Predefined kernel set |
### Layout Notation
@@ -98,6 +98,28 @@ Element-wise fusion: `C = op(A x B + D0 + D1 + ...)`
Supported ops: `PassThrough`, `MultiDAdd`, `Relu`, `Gelu`, `Sigmoid`, `Tanh`
### Grouped
Batched GEMM over a list of independently-shaped groups in a single launch
(`ck_tile::GroupedGemmKernel`). Brings the dispatcher to parity with the Tile Engine
`grouped_gemm` op. The per-group argument vector is built with `MakeKargs`, copied to an
internally-allocated `DeviceMem` workspace, and the device pointer + group count are passed
to the kernel (the dispatcher workspace idiom — no external `kargs_ptr`).
- Datatypes: `fp16`, `bf16`, `fp8`, `bf8` (matches the Tile Engine grouped runnable set;
`fp8`/`bf8` accumulate in `fp32` and emit an `fp16` C output).
- Layouts: `rcr`, `rrr`, `ccr`, `crr` (C is always row-major).
```bash
python3 unified_gemm_codegen.py \
--datatype fp16 \
--layout rcr \
--variants grouped \
--gpu-target gfx942 \
--output-dir generated_kernels
```
Build and run end-to-end with [`examples/gemm/cpp/02_grouped_gemm_driver.cpp`](../examples/gemm/cpp/README.md).
## Output Structure
```
@@ -105,6 +127,7 @@ generated_kernels/
|---- gemm_fp16_rcr_compv4_..._128x128x32_....hpp # GEMM kernels
|---- gemm_fp16_rcr_compv4_..._preshuffle.hpp
|---- gemm_fp16_rcr_compv4_..._multid_Relu_d1.hpp
|---- gemm_fp16_rcr_compv3_..._128x128x64_..._grouped.hpp # Grouped GEMM kernels
|---- grouped_conv_fwd_fp16_nhwgc_..._128x128x32_....hpp # Grouped conv kernels
+---- ...
```

2
dispatcher/codegen/arch_filter.py
View File

@@ -50,6 +50,7 @@ class OperatorType(Enum):
GEMM = "gemm"
GEMM_PRESHUFFLE = "gemm_preshuffle"
GEMM_MULTI_D = "gemm_multi_d"
GEMM_GROUPED = "gemm_grouped"
GEMM_STREAMK = "gemm_streamk"
CONV_FWD = "conv_fwd"
CONV_BWD_DATA = "conv_bwd_data"
@@ -86,6 +87,7 @@ OPERATOR_TILE_CONSTRAINTS = {
"tile_n_alignment": 16,
"tile_k_alignment": 8,
},
OperatorType.GEMM_GROUPED: {
# NOTE: these are copied from plain GEMM and only gate tile *shape* validity.
# They do NOT express Stream-K's real feasibility requirement -- that a problem
# has enough output tiles to partition K-work across the CUs. That gate is

92
dispatcher/codegen/unified_gemm_codegen.py
View File

@@ -202,6 +202,7 @@ class GemmVariant(Enum):
STANDARD = "standard"
PRESHUFFLE = "preshuffle"
MULTI_D = "multi_d"
GROUPED = "grouped"
# Stream-K. COVERAGE LIMITATION: the dispatcher does NOT yet emit the full
# Old-TE Stream-K tile surface. The kernels generated here are driven by the
# tile list passed to this codegen, which is narrower than tile_engine's:
@@ -367,6 +368,8 @@ class KernelNaming:
name += "_preshuffle"
elif config.variant == GemmVariant.MULTI_D:
name += f"_multid_{config.elementwise_op}_d{config.num_d_tensors}"
elif config.variant == GemmVariant.GROUPED:
name += "_grouped"
elif config.variant == GemmVariant.STREAM_K:
name += "_streamk"
# Atomic keeps the bare "_streamk" suffix for name parity with the
@@ -420,6 +423,15 @@ class CKTileKernelGenerator:
includes += """
#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp"
#include "ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp"
"""
if config.variant == GemmVariant.GROUPED:
includes += """
#include <vector>
#include <hip/hip_runtime.h>
#include "ck_tile/host/device_memory.hpp"
#include "ck_tile/host/hip_check_error.hpp"
#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp"
"""
if config.preshuffle:
@@ -604,6 +616,9 @@ using CLayout = {ns_name}::CLayout;
#define GEMM_KEY_TRANSPOSE_C 0
#define GEMM_KEY_GROUPED 0
#define GEMM_KEY_SPLIT_K 1
using ALayout = {ns_name}::ALayout;
using BLayout = {ns_name}::BLayout;
using CLayout = {ns_name}::CLayout;
#endif // CK_TILE_SINGLE_KERNEL_INCLUDE
"""
@@ -629,6 +644,8 @@ using CLayout = {ns_name}::CLayout;
"""Generate launch function"""
if config.variant == GemmVariant.MULTI_D:
return self._launch_function_multi_d(config)
if config.variant == GemmVariant.GROUPED:
return self._launch_function_grouped(config)
if config.variant == GemmVariant.STREAM_K:
return self._launch_function_streamk(config)
if config.preshuffle:
@@ -683,6 +700,69 @@ using CLayout = {ns_name}::CLayout;
return ave_time;
}}"""
def _launch_function_grouped(self, config: KernelConfig) -> str:
"""Generate launch function for grouped GEMM.
Follows the dispatcher's workspace idiom (see grouped_conv stream-K launch in
unified_grouped_conv_codegen.py): signature is (args, stream); the device
workspace is allocated internally via DeviceMem rather than passed in. The
grouped kernel's per-group arg vector is built with MakeKargs, copied to the
workspace, and the device pointer + group count are passed to the kernel.
"""
persistent = config.trait.persistent
grid_expr = (
"GemmKernel::MaxOccupancyGridSize(stream)"
if persistent
else "dim3(kargs.empty() ? 0 : kargs.back().block_end, 1, 1)"
)
return f"""
static float launch(const std::vector<ck_tile::GroupedGemmHostArgs<>>& gemm_descs,
const stream_config& stream) {{
if(gemm_descs.empty()) return 0.0f;
float ave_time{{0}};
constexpr auto scheduler = {self.tm.SCHEDULER_TO_CK[config.trait.scheduler]};
using UniversalGemmProblem = UniversalGemmPipelineProblem<
ADataType, BDataType, AccDataType, TileShape,
TileGemmUniversalTraits<kPadM, kPadN, kPadK, DoubleSmemBuffer,
ALayout, BLayout, CLayout, TransposeC,
UseStructuredSparsity, UsePersistentKernel,
NumWaveGroups, Preshuffle>,
scheduler>;
using GemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}<UniversalGemmProblem>;
{self._epilogue_code(config)}
using GemmKernel = ck_tile::GroupedGemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
auto kargs = GemmKernel::MakeKargs(gemm_descs);
if(!GemmKernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Arguments not supported for grouped gemm kernel");
}}
// Workspace allocated internally (dispatcher idiom, mirrors grouped_conv stream-K).
const std::size_t ws_size = kargs.size() * sizeof(ck_tile::GemmTransKernelArg<>);
ck_tile::DeviceMem workspace_dev(ws_size);
HIP_CHECK_ERROR(hipMemcpyWithStream(workspace_dev.GetDeviceBuffer(),
kargs.data(),
ws_size,
hipMemcpyHostToDevice,
stream.stream_id_));
const dim3 grids = {grid_expr};
const dim3 blocks = GemmKernel::BlockSize();
constexpr int kBlockPerCu = {config.k_block_per_cu};
ave_time = launch_kernel(stream,
make_kernel<kBlockPerCu>(GemmKernel{{}}, grids, blocks, 0,
cast_pointer_to_constant_address_space(workspace_dev.GetDeviceBuffer()),
kargs.size()));
return ave_time;
}}"""
def _launch_function_preshuffle(self, config: KernelConfig) -> str:
"""Generate launch function for preshuffle GEMM (weight preshuffle variant)
@@ -985,7 +1065,7 @@ using CLayout = {ns_name}::CLayout;
tuple<>, CLayout, element_wise::PassThrough,
TilePartitioner::MPerBlock, TilePartitioner::NPerBlock,
WarpPerBlock_M, WarpPerBlock_N, WarpTileM, WarpTileN, WarpTileK,
TransposeC, NumWaveGroups>;
TransposeC, NumWaveGroups, false, 1, 1, DoubleSmemBuffer>;
using GemmEpilogue = CShuffleEpilogue<EpilogueProblem>;"""
else:
return """
@@ -1327,7 +1407,7 @@ class UnifiedGemmCodegen:
"""Get all configurations for a variant
Args:
variant: GEMM variant (STANDARD, PRESHUFFLE, MULTI_D)
variant: GEMM variant (STANDARD, PRESHUFFLE, MULTI_D, GROUPED)
Returns:
List of valid kernel configurations for the variant
@@ -1428,6 +1508,11 @@ class UnifiedGemmCodegen:
)
)
elif variant == GemmVariant.GROUPED:
# Grouped GEMM uses the same tile/trait configs as STANDARD —
# the only difference is the kernel type (GroupedGemmKernel vs GemmKernel)
configs.append(KernelConfig(tile=tile, trait=trait, variant=variant))
return configs
def _get_tile_configs(self) -> List[TileConfig]:
@@ -1534,6 +1619,7 @@ class UnifiedGemmCodegen:
GemmVariant.STANDARD: OperatorType.GEMM,
GemmVariant.PRESHUFFLE: OperatorType.GEMM_PRESHUFFLE,
GemmVariant.MULTI_D: OperatorType.GEMM_MULTI_D,
GemmVariant.GROUPED: OperatorType.GEMM_GROUPED,
GemmVariant.STREAM_K: OperatorType.GEMM_STREAMK,
}
operator = variant_to_operator.get(variant, OperatorType.GEMM)
@@ -1784,7 +1870,7 @@ def main():
parser.add_argument(
"--variants",
nargs="+",
choices=["standard", "preshuffle", "multi_d", "stream_k"],
choices=["standard", "preshuffle", "multi_d", "stream_k" ,"grouped"],
default=["standard"],
help="Variants to generate",
)

192
dispatcher/examples/gemm/cpp/02_grouped_gemm_driver.cpp Normal file
View File

@@ -0,0 +1,192 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Minimal standalone grouped-GEMM driver (dispatcher way).
*
* Grouped GEMM cannot ride the standard dispatcher.run(A,B,C,problem) path:
* that backend hardcodes a single GemmHostArgs. Instead, this driver includes a
* single generated grouped kernel header (CK_TILE_SINGLE_KERNEL_INCLUDE) and
* calls SelectedKernel::launch(descs, stream) directly with a vector of
* descriptors -- the same 2-arg signature the dispatcher generates (workspace is
* allocated INSIDE launch()). It builds per-group tensors, runs, and verifies
* each group against ck_tile::reference_gemm.
*
* Build (single-kernel include style):
* hipcc -std=c++17 --offload-arch=gfx942 \
* -DCK_TILE_SINGLE_KERNEL_INCLUDE \
* -I <ck>/include -I <generated_dir> \
* -include <generated_dir>/<one>_grouped.hpp \
* 02_grouped_gemm_driver.cpp -o grouped_gemm_driver
*/
#include <hip/hip_runtime.h>
#include <algorithm>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/ops/gemm.hpp"
// The generated grouped kernel header is injected on the command line with
// -include and -DCK_TILE_SINGLE_KERNEL_INCLUDE. It exports into the global
// namespace: SelectedKernel, ADataType, BDataType, CDataType, AccDataType,
// ALayout, BLayout, CLayout, and KERNEL_NAME.
template <typename Layout>
static constexpr inline auto is_row_major(Layout)
{
return ck_tile::bool_constant<
std::is_same_v<ck_tile::remove_cvref_t<Layout>, ck_tile::tensor_layout::gemm::RowMajor>>{};
}
static std::vector<int> parse_csv_ints(const std::string& s)
{
std::vector<int> out;
std::string cur;
for(char c : s)
{
if(c == ',')
{
if(!cur.empty())
{
out.push_back(std::stoi(cur));
cur.clear();
}
}
else
cur.push_back(c);
}
if(!cur.empty())
out.push_back(std::stoi(cur));
return out;
}
static std::string get_opt(int argc, char** argv, const std::string& key, const std::string& def)
{
for(int i = 1; i < argc - 1; ++i)
if(key == argv[i])
return argv[i + 1];
return def;
}
int main(int argc, char** argv)
{
const int group_count = std::stoi(get_opt(argc, argv, "--groups", "8"));
const int kbatch = std::stoi(get_opt(argc, argv, "--kbatch", "1"));
const int warmup = std::stoi(get_opt(argc, argv, "--warmup", "10"));
const int repeat = std::stoi(get_opt(argc, argv, "--repeat", "50"));
const bool validate = get_opt(argc, argv, "--validate", "1") != "0";
std::vector<int> Ms = parse_csv_ints(get_opt(argc, argv, "--Ms", ""));
std::vector<int> Ns = parse_csv_ints(get_opt(argc, argv, "--Ns", ""));
std::vector<int> Ks = parse_csv_ints(get_opt(argc, argv, "--Ks", ""));
const int dm = std::stoi(get_opt(argc, argv, "--m", "256"));
const int dn = std::stoi(get_opt(argc, argv, "--n", "256"));
const int dk = std::stoi(get_opt(argc, argv, "--k", "512"));
auto sz = static_cast<std::size_t>(group_count);
if(Ms.size() != sz || Ns.size() != sz || Ks.size() != sz)
{
Ms.assign(group_count, dm);
Ns.assign(group_count, dn);
Ks.assign(group_count, dk);
}
std::cout << "Kernel: " << KERNEL_NAME << "\n";
std::cout << "groups=" << group_count << " kbatch=" << kbatch << "\n";
std::vector<ck_tile::HostTensor<ADataType>> a_host, b_host;
std::vector<ck_tile::HostTensor<CDataType>> c_host;
std::vector<std::unique_ptr<ck_tile::DeviceMem>> a_dev, b_dev, c_dev;
std::vector<ck_tile::index_t> sA(group_count), sB(group_count), sC(group_count);
std::vector<ck_tile::GroupedGemmHostArgs<>> descs;
descs.reserve(group_count);
for(int i = 0; i < group_count; ++i)
{
const ck_tile::index_t M = Ms[i], N = Ns[i], K = Ks[i];
sA[i] = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{}));
sB[i] = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{}));
sC[i] = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{}));
a_host.push_back(ck_tile::HostTensor<ADataType>(
ck_tile::host_tensor_descriptor(M, K, sA[i], is_row_major(ALayout{}))));
b_host.push_back(ck_tile::HostTensor<BDataType>(
ck_tile::host_tensor_descriptor(K, N, sB[i], is_row_major(BLayout{}))));
c_host.push_back(ck_tile::HostTensor<CDataType>(
ck_tile::host_tensor_descriptor(M, N, sC[i], is_row_major(CLayout{}))));
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_host[i]);
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_host[i]);
c_host[i].SetZero();
a_dev.push_back(std::make_unique<ck_tile::DeviceMem>(a_host[i]));
b_dev.push_back(std::make_unique<ck_tile::DeviceMem>(b_host[i]));
c_dev.push_back(std::make_unique<ck_tile::DeviceMem>(c_host[i]));
c_dev[i]->SetZero();
descs.push_back(ck_tile::GroupedGemmHostArgs<>{a_dev[i]->GetDeviceBuffer(),
b_dev[i]->GetDeviceBuffer(),
{},
c_dev[i]->GetDeviceBuffer(),
kbatch,
M,
N,
K,
sA[i],
sB[i],
{},
sC[i]});
}
const ck_tile::stream_config s{nullptr, true, /*log=*/0, warmup, repeat};
float ave_time = SelectedKernel::launch(descs, s);
std::size_t flop = 0, bytes = 0;
for(int i = 0; i < group_count; ++i)
{
flop += std::size_t(2) * Ms[i] * Ns[i] * Ks[i];
bytes += sizeof(ADataType) * Ms[i] * Ks[i] + sizeof(BDataType) * Ks[i] * Ns[i] +
sizeof(CDataType) * Ms[i] * Ns[i];
}
const float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
const float gbps = static_cast<float>(bytes) / 1.E6 / ave_time;
std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, " << gbps
<< " GB/s\n";
for(int i = 0; i < group_count; ++i)
c_dev[i]->FromDevice(c_host[i].data());
bool pass = true;
if(validate)
{
for(int i = 0; i < group_count; ++i)
{
ck_tile::HostTensor<CDataType> ref(
ck_tile::host_tensor_descriptor(Ms[i], Ns[i], sC[i], is_row_major(CLayout{})));
ref.SetZero();
ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(
a_host[i], b_host[i], ref);
const float maxv = *std::max_element(ref.mData.begin(), ref.mData.end());
const auto rtol = ck_tile::get_relative_threshold<ADataType, CDataType, AccDataType>(
ck_tile::integer_divide_ceil(Ks[i], kbatch));
const auto atol = ck_tile::get_absolute_threshold<ADataType, CDataType, AccDataType>(
maxv / kbatch, ck_tile::integer_divide_ceil(Ks[i], kbatch));
bool ok =
ck_tile::check_err(c_host[i], ref, "group[" + std::to_string(i) + "]", rtol, atol);
pass &= ok;
}
std::cout << "Verification: " << (pass ? "PASS" : "FAIL") << "\n";
}
return pass ? 0 : 1;
}

26
dispatcher/examples/gemm/cpp/README.md
View File

@@ -37,6 +37,7 @@ cd examples
| [04_heuristics.cpp](04_heuristics.cpp) | Heuristic-based kernel selection |
| [05_json_export.cpp](05_json_export.cpp) | Registry JSON export for external tools |
| [06_multi_registry.cpp](06_multi_registry.cpp) | Multiple registries with named kernel sets |
| [02_grouped_gemm_driver.cpp](02_grouped_gemm_driver.cpp) | Standalone grouped (batched) GEMM driver: builds per-group descriptors, launches `GroupedGemmKernel`, verifies each group |
## Example Details
@@ -113,6 +114,31 @@ Consolidated example combining performance benchmarking with correctness validat
- GPU reference validation using `ck_tile::reference_gemm_gpu`
- Configurable tolerances
### 02_grouped_gemm_driver.cpp - Grouped (Batched) GEMM
Standalone driver for the `grouped` codegen variant. One generated grouped kernel header is
injected on the command line (`-include <stem>_grouped.hpp -DCK_TILE_SINGLE_KERNEL_INCLUDE`);
the driver reads the kernel's `ADataType/BDataType/CDataType/ALayout/BLayout/CLayout` so it
works for any datatype/layout the kernel was generated for (no hardcoded `fp16`/`rcr`).
```bash
# Build against one generated grouped kernel
hipcc -std=c++17 --offload-arch=gfx942 -O3 -DCK_TILE_SINGLE_KERNEL_INCLUDE \
-I ../../../include -I ../../../../include \
-I generated_kernels \
-include gemm_fp16_rcr_compv3_..._grouped.hpp \
02_grouped_gemm_driver.cpp -o gemm_02_grouped
# Run: 8 groups, verify each group against the CPU reference
./gemm_02_grouped --groups 8 --Ms 3840 --Ns 4096 --Ks 2048 \
--warmup 50 --repeat 100 --validate 1
```
**Features:**
- Builds a vector of per-group `GroupedGemmHostArgs` descriptors (per-group M/N/K)
- Layout-driven leading dimensions via `get_default_stride(is_row_major(Layout))`
- Per-group correctness check using `ck_tile::reference_gemm`
- Reports per-group PASS/FAIL plus aggregate TFLOPS
### 04_heuristics.cpp - Heuristic Selection
Demonstrates custom kernel selection based on problem characteristics:

335
dispatcher/python/gemm_utils.py
View File

@@ -59,6 +59,88 @@ def _cap(flag: bool) -> str:
return "True" if flag else "False"
# ---------------------------------------------------------------------------
# Dtype codecs: map a bridge dtype token -> numpy dtype for host operands.
#
# fp16 maps to plain numpy; bf16/fp8/bf8 need ml_dtypes. fp8/bf8 use the FNUZ
# encodings (E4M3FNUZ / E5M2FNUZ) that the gfx942 MFMA path expects -- matching
# the regular bridge's fp8/bf8 codec (PR #8887). ml_dtypes is imported lazily so
# the fp16-only path keeps working where ml_dtypes is unavailable.
# ---------------------------------------------------------------------------
# Canonicalize common spellings to a single token.
_DTYPE_ALIASES = {
"fp16": "fp16",
"f16": "fp16",
"half": "fp16",
"float16": "fp16",
"bf16": "bf16",
"bfloat16": "bf16",
"fp8": "fp8",
"fp8_e4m3": "fp8",
"e4m3": "fp8",
"bf8": "bf8",
"fp8_e5m2": "bf8",
"e5m2": "bf8",
}
def numpy_dtype_for(dtype: str):
"""Return the numpy dtype object used for host operands of ``dtype``.
fp16 -> np.float16; bf16/fp8/bf8 require the ``ml_dtypes`` package (imported
lazily) and use FNUZ fp8 encodings for gfx942 parity.
"""
token = _DTYPE_ALIASES.get(str(dtype).lower())
if token is None:
raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}")
if token == "fp16":
return np.float16
try:
import ml_dtypes # noqa: WPS433 (lazy: optional dep)
except ImportError as exc: # pragma: no cover - env-dependent
raise RuntimeError(
f"dtype {dtype!r} requires the 'ml_dtypes' package (pip install ml_dtypes)"
) from exc
if token == "bf16":
return np.dtype(ml_dtypes.bfloat16)
if token == "fp8":
return np.dtype(ml_dtypes.float8_e4m3fnuz)
if token == "bf8":
return np.dtype(ml_dtypes.float8_e5m2fnuz)
raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}") # pragma: no cover
def output_dtype_for(dtype: str) -> str:
"""Return the bridge dtype token of a kernel's OUTPUT for input ``dtype``.
Mirrors ``codegen_common.CommonTypeMappings.get_output_dtype`` (fp8/bf8 ->
fp16, else identity): the generated grouped kernel emits an fp16 ``CDataType``
for fp8/bf8 inputs, so the host C buffer must be sized/typed by the OUTPUT
dtype, not the INPUT dtype. ``codegen_common`` lives on the dispatcher
``codegen`` dir which ctypes_utils already puts on ``sys.path``; import it
lazily so the fp16-only path has no extra dependency.
"""
token = _DTYPE_ALIASES.get(str(dtype).lower())
if token is None:
raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}")
try:
from codegen_common import CommonTypeMappings # noqa: WPS433 (lazy)
except ImportError: # pragma: no cover - fall back to the documented mapping
return "fp16" if token in ("fp8", "bf8") else token
return CommonTypeMappings.get_output_dtype(token)
def output_numpy_dtype_for(dtype: str):
"""Numpy dtype of a kernel's OUTPUT buffer for input ``dtype``.
Composition of :func:`output_dtype_for` + :func:`numpy_dtype_for`. For
fp8/bf8 this resolves to ``np.float16`` (2 bytes) because the kernel's
``CDataType`` is fp16; for fp16/bf16 it equals the input dtype.
"""
return numpy_dtype_for(output_dtype_for(dtype))
# ============================================================================
# The shared contract: GemmKernelConfig
# ============================================================================
@@ -151,6 +233,8 @@ class GemmKernelConfig:
name += "_preshuffle"
elif self.variant == "streamk":
name += "_streamk"
elif self.variant == "grouped":
name += "_grouped"
return name
# ------------------------------------------------------------------ #
@@ -264,6 +348,39 @@ class GemmProblem:
return cls(M=int(d["M"]), N=int(d["N"]), K=int(d["K"]))
@dataclass
class GroupedGemmProblem:
"""A grouped GEMM problem: a list of independent (M, N, K) sub-problems
all run by a single grouped kernel launch.
Each group g computes C_g[M_g x N_g] = A_g[M_g x K_g] @ B_g[K_g x N_g].
"""
groups: List[Tuple[int, int, int]]
@classmethod
def uniform(
cls, group_count: int, M: int, N: int, K: int
) -> "GroupedGemmProblem":
"""All groups share the same (M, N, K) shape."""
return cls(groups=[(int(M), int(N), int(K)) for _ in range(int(group_count))])
@property
def group_count(self) -> int:
return len(self.groups)
@property
def flops(self) -> float:
return sum(2.0 * m * n * k for (m, n, k) in self.groups)
def to_dict(self) -> Dict[str, Any]:
return {"groups": [[int(m), int(n), int(k)] for (m, n, k) in self.groups]}
@classmethod
def from_dict(cls, d: Dict[str, Any]) -> "GroupedGemmProblem":
return cls(groups=[(int(m), int(n), int(k)) for (m, n, k) in d["groups"]])
@dataclass
class GemmResult:
output: np.ndarray
@@ -277,6 +394,22 @@ class GemmResult:
return self.status == 0
@dataclass
class GroupedGemmResult:
"""Result of a grouped GEMM launch: one output per group plus aggregate
timing/throughput across the whole batch."""
outputs: List[np.ndarray]
time_ms: float
status: int
tflops: float
kernel_name: str
@property
def success(self) -> bool:
return self.status == 0
# ============================================================================
# ctypes ABI wrapper
# ============================================================================
@@ -294,6 +427,8 @@ class GemmDispatcherLib:
self._path = Path(so_path)
self._lib = ctypes.CDLL(str(self._path))
self._has_indexed = hasattr(self._lib, "dispatcher_get_kernel_name_at")
self._has_grouped = hasattr(self._lib, "dispatcher_run_grouped_gemm")
self._has_single = hasattr(self._lib, "dispatcher_run_gemm")
self._setup_functions()
def _setup_functions(self) -> None:
@@ -316,16 +451,32 @@ class GemmDispatcherLib:
]
lib.dispatcher_get_kernel_name_at.restype = ctypes.c_int
lib.dispatcher_run_gemm.argtypes = [
ctypes.c_void_p, # A (host)
ctypes.c_void_p, # B (host)
ctypes.c_void_p, # C (host)
ctypes.c_int64, # M
ctypes.c_int64, # N
ctypes.c_int64, # K
ctypes.POINTER(ctypes.c_float), # time_ms
]
lib.dispatcher_run_gemm.restype = ctypes.c_int
# Single-problem ABI (regular GEMM .so). Absent on grouped libs.
if self._has_single:
lib.dispatcher_run_gemm.argtypes = [
ctypes.c_void_p, # A (host)
ctypes.c_void_p, # B (host)
ctypes.c_void_p, # C (host)
ctypes.c_int64, # M
ctypes.c_int64, # N
ctypes.c_int64, # K
ctypes.POINTER(ctypes.c_float), # time_ms
]
lib.dispatcher_run_gemm.restype = ctypes.c_int
# Multi-problem ABI (grouped GEMM .so). Absent on regular libs.
if self._has_grouped:
lib.dispatcher_run_grouped_gemm.argtypes = [
ctypes.c_int, # group_count
ctypes.POINTER(ctypes.c_int64), # Ms[]
ctypes.POINTER(ctypes.c_int64), # Ns[]
ctypes.POINTER(ctypes.c_int64), # Ks[]
ctypes.POINTER(ctypes.c_void_p), # A_ptrs[]
ctypes.POINTER(ctypes.c_void_p), # B_ptrs[]
ctypes.POINTER(ctypes.c_void_p), # C_ptrs[]
ctypes.POINTER(ctypes.c_float), # time_ms
]
lib.dispatcher_run_grouped_gemm.restype = ctypes.c_int
lib.dispatcher_cleanup.argtypes = []
lib.dispatcher_cleanup.restype = None
@@ -371,6 +522,52 @@ class GemmDispatcherLib:
)
return status, time_ms.value
def run_grouped(
self,
A_list: List[np.ndarray],
B_list: List[np.ndarray],
C_list: List[np.ndarray],
Ms: List[int],
Ns: List[int],
Ks: List[int],
) -> Tuple[int, float]:
"""Launch the grouped kernel over a batch of (M, N, K) sub-problems.
Each A/B/C entry is a host numpy array already laid out (dtype + row/col
transpose) as the kernel expects for its compile-time layout; the caller
(GpuGroupedGemmRunner) does that per-dtype/per-layout packing. Pointers
are marshalled into ctypes pointer arrays.
"""
if not self._has_grouped:
raise RuntimeError(
f"{self._path} does not expose dispatcher_run_grouped_gemm"
)
g = len(A_list)
c_int64_arr = (ctypes.c_int64 * g)
c_void_arr = (ctypes.c_void_p * g)
ms = c_int64_arr(*[int(m) for m in Ms])
ns = c_int64_arr(*[int(n) for n in Ns])
ks = c_int64_arr(*[int(k) for k in Ks])
a_ptrs = c_void_arr(*[A.ctypes.data_as(ctypes.c_void_p) for A in A_list])
b_ptrs = c_void_arr(*[B.ctypes.data_as(ctypes.c_void_p) for B in B_list])
c_ptrs = c_void_arr(*[C.ctypes.data_as(ctypes.c_void_p) for C in C_list])
time_ms = ctypes.c_float(0.0)
status = self._lib.dispatcher_run_grouped_gemm(
g,
ms,
ns,
ks,
a_ptrs,
b_ptrs,
c_ptrs,
ctypes.byref(time_ms),
)
return status, time_ms.value
def cleanup(self) -> None:
self._lib.dispatcher_cleanup()
@@ -619,6 +816,94 @@ class GpuGemmRunner:
)
class GpuGroupedGemmRunner:
"""High-level runner for the GROUPED variant: construct from a grouped .so
path, call run(A_list, B_list, problem).
Like GpuGemmRunner, the ctypes ABI takes HOST pointers and manages GPU
memory internally (per group), so this runner only marshals the host operand
arrays. The runner is parameterized by ``(dtype, layout)`` (mirroring
``GpuGemmRunner``/``GemmProblem``): the A/B operands are cast to the per-dtype
INPUT numpy codec (fp16/bf16/fp8-E4M3FNUZ/bf8-E5M2FNUZ) and transposed per the
A/B/C layout so the contiguous host buffer matches the layout the kernel was
generated with (the ctypes lib derives strides from the same layouts).
The C/output buffer is sized/typed by the kernel's OUTPUT dtype, not the input
dtype: for fp8/bf8 inputs the generated kernel's ``CDataType`` is fp16, so the
host C buffer is fp16 (2 bytes) even though A/B are 1-byte fp8/bf8. Sizing C by
the input dtype would under-allocate by 2x and the ctypes copy-back would
overrun the host buffer (heap corruption). See :func:`output_numpy_dtype_for`.
"""
def __init__(self, lib_path: Path, dtype: str = "fp16", layout: str = "rcr"):
self.lib = GemmDispatcherLib(lib_path)
if not self.lib.initialize():
raise RuntimeError(
f"Failed to initialize grouped dispatcher .so: {lib_path}"
)
names = self.lib.kernel_names
self._kernel_name = names[0] if names else "unknown"
self._dtype = dtype
# A/B (input) codec vs C (output) codec: they differ for fp8/bf8
# (output is fp16), so keep them distinct to size the C buffer correctly.
self._np_dtype = numpy_dtype_for(dtype)
self._c_np_dtype = output_numpy_dtype_for(dtype)
if len(layout) != 3 or any(ch not in ("r", "c") for ch in layout):
raise ValueError(f"layout must be a 3-char r/c string, got {layout!r}")
self._layout = layout
@property
def kernel_name(self) -> str:
return self._kernel_name
def run(
self,
A_list: List[np.ndarray],
B_list: List[np.ndarray],
problem: GroupedGemmProblem,
) -> GroupedGemmResult:
groups = problem.groups
if len(A_list) != len(groups) or len(B_list) != len(groups):
raise ValueError(
"A_list/B_list length must match the number of groups "
f"({len(A_list)}/{len(B_list)} vs {len(groups)})"
)
Ms = [g[0] for g in groups]
Ns = [g[1] for g in groups]
Ks = [g[2] for g in groups]
la, lb, _lc = self._layout[0], self._layout[1], self._layout[2]
nd = self._np_dtype
c_nd = self._c_np_dtype # OUTPUT dtype (fp16 for fp8/bf8); see __init__.
A_h: List[np.ndarray] = []
B_h: List[np.ndarray] = []
C_h: List[np.ndarray] = []
for A, B, (M, N, _K) in zip(A_list, B_list, groups):
# A logically MxK, B logically KxN, C row-major MxN (CLayout is always
# RowMajor for grouped). Store each operand so its contiguous buffer
# matches its layout: row-major -> as-is, col-major -> transpose.
A_buf = A if la == "r" else A.T
B_buf = B if lb == "r" else B.T
A_h.append(np.ascontiguousarray(A_buf, dtype=nd))
B_h.append(np.ascontiguousarray(B_buf, dtype=nd))
# Size C by the kernel's CDataType (output dtype), NOT the input dtype:
# fp8/bf8 inputs produce fp16 output, so a 1-byte C would be overrun.
C_h.append(np.zeros((M, N), dtype=c_nd))
status, time_ms = self.lib.run_grouped(A_h, B_h, C_h, Ms, Ns, Ks)
tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0
return GroupedGemmResult(
outputs=C_h,
time_ms=time_ms,
status=status,
tflops=tflops,
kernel_name=self._kernel_name,
)
# ============================================================================
# Build API: codegen + hipcc -> .so paths (no GPU)
# ============================================================================
@@ -697,6 +982,18 @@ def _tile_engine_codegen_flags() -> Tuple[str, ...]:
return tuple(flags)
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``.
"""
if config.variant == "grouped":
return "grouped_gemm_ctypes_lib.cpp"
return "gemm_ctypes_lib.cpp"
def _build_compile_jobs(
config: GemmKernelConfig, header: Path
) -> Tuple[Dict[str, Any], Path]:
@@ -705,7 +1002,7 @@ def _build_compile_jobs(
ck_root = root.parent
build_dir = _cu.get_build_dir()
output_dir = _cu.get_generated_kernels_dir()
ctypes_source = root / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp"
ctypes_source = root / "bindings" / "ctypes" / _ctypes_source_name(config)
static_lib = build_dir / "libck_tile_dispatcher.a"
lib_path = build_dir / "examples" / f"lib{config.name}.so"
@@ -786,13 +1083,13 @@ def setup_multiple_gemm_dispatchers(
codegen_script = _cu.get_codegen_path()
output_dir = _cu.get_generated_kernels_dir()
static_lib = _cu.get_build_dir() / "libck_tile_dispatcher.a"
ctypes_source = (
_cu.get_dispatcher_root() / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp"
)
if not static_lib.exists() or not ctypes_source.exists():
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:
raise FileNotFoundError(
"Missing static lib or ctypes source required for compilation:\n"
f" {static_lib}\n {ctypes_source}\n"
f" {static_lib}\n " + "\n ".join(missing) + "\n"
"Build the dispatcher first (cmake + make)."
)
@@ -915,6 +1212,7 @@ def expand_sweep(
arch: str,
dtype: str = "fp16",
layout: str = "rcr",
variant: str = "standard",
) -> List[GemmKernelConfig]:
"""Expand a Tile Engine GEMM JSON sweep config into GemmKernelConfig list.
@@ -924,8 +1222,8 @@ def expand_sweep(
one GemmKernelConfig. Invalid combinations are dropped via the dispatcher's
own validator, and duplicates (by .name) are collapsed.
The signature is controlled by the `dtype` and `layout` arguments (defaults
to fp16 / rcr).
The operand signature (``dtype``, ``layout``) is applied to every emitted
GemmKernelConfig, so the same sweep expands across any supported dtype/layout.
"""
with open(config_path) as f:
cfg = json.load(f)
@@ -1015,6 +1313,7 @@ def expand_sweep(
pad_k=bool(pk),
persistent=bool(persist),
gfx_arch=arch,
variant=variant,
)
if c.name in seen:
continue

130
dispatcher/tests/test_grouped_gemm_codegen.py Normal file
View File

@@ -0,0 +1,130 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""CPU-only unit tests for the grouped GEMM variant of the TE -> Dispatcher bridge.
Grouped GEMM has a genuinely different ABI from regular GEMM (a list of
sub-problems, a launch(descs, stream) signature, and an internally-owned device
workspace), so it is generated by a dedicated path in the unified GEMM codegen
and launched by grouped_gemm_ctypes_lib.cpp. These tests lock in the host-side
contract of that path without touching a GPU:
* the GROUPED variant generates the grouped launch/kernel with fully-qualified
ck_tile:: types (so the header does not silently depend on a using-directive);
* the grouped ctypes bridge includes the standard headers it uses directly
rather than relying on transitive includes from the force-included kernel;
* the FNUZ decode table shared by the fp8/bf8 path is cached and read-only.
No GPU is touched -- codegen is pure string formatting and the helpers are pure.
Run: python3 -m pytest tests/test_grouped_gemm_codegen.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 / "codegen"))
sys.path.insert(0, str(DISPATCHER_DIR / "python"))
from codegen_common import TileConfig # noqa: E402
from unified_gemm_codegen import ( # noqa: E402
GemmVariant,
KernelConfig,
TraitConfig,
CKTileKernelGenerator,
)
from gemm_utils import _fnuz_decode_table # noqa: E402
GROUPED_CTYPES_LIB = (
DISPATCHER_DIR / "bindings" / "ctypes" / "grouped_gemm_ctypes_lib.cpp"
)
def _grouped_config() -> KernelConfig:
"""A minimal valid grouped-GEMM kernel config (fp16/rcr, cshuffle)."""
tile = TileConfig(
tile_m=128, tile_n=128, tile_k=32,
warp_m=2, warp_n=2, warp_k=1,
warp_tile_m=32, warp_tile_n=32, warp_tile_k=16,
)
trait = TraitConfig(
pipeline="compv3", epilogue="cshuffle", scheduler="intrawave",
pad_m=False, pad_n=False, pad_k=False, persistent=False,
)
return KernelConfig(tile=tile, trait=trait, variant=GemmVariant.GROUPED)
class TestGroupedVariant(unittest.TestCase):
def test_enum_value(self):
self.assertEqual(GemmVariant.GROUPED.value, "grouped")
class TestGroupedCodegen(unittest.TestCase):
"""The generated grouped header must be self-contained and correctly typed."""
def setUp(self):
gen = CKTileKernelGenerator("fp16", "rcr")
self.src = gen.generate(_grouped_config())
def test_generates_grouped_launch_and_kernel(self):
# The grouped path emits the (descs, stream) launch and the grouped kernel.
self.assertIn(
"launch(const std::vector<ck_tile::GroupedGemmHostArgs<>>&", self.src
)
self.assertIn("ck_tile::GroupedGemmKernel<", self.src)
def test_device_memory_is_ck_tile_qualified(self):
# DeviceMem lives in the ck_tile namespace; the workspace allocation must
# be fully qualified so the header does not depend on a using-directive.
self.assertIn("ck_tile::DeviceMem workspace_dev", self.src)
# And no bare, unqualified `DeviceMem <name>` declaration should remain.
import re
self.assertIsNone(
re.search(r"(?<!:)\bDeviceMem\s+workspace_dev", self.src),
"found an unqualified DeviceMem declaration in generated grouped header",
)
def test_grouped_includes_are_present(self):
# The grouped variant pulls in the headers its launch needs directly.
for inc in (
"ck_tile/host/device_memory.hpp",
"ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp",
):
self.assertIn(inc, self.src)
class TestGroupedCtypesLibIncludes(unittest.TestCase):
"""The grouped ctypes bridge must include what it uses, not lean on the
transitively-included kernel header (include order is otherwise fragile)."""
def test_includes_array_header(self):
text = GROUPED_CTYPES_LIB.read_text()
# std::array<..., 0>{} descriptors are built here; <array> must be direct.
self.assertIn("std::array", text)
self.assertIn("#include <array>", text)
class TestFnuzTableCaching(unittest.TestCase):
"""The 256-entry FNUZ decode table is pure per (exp_bits, mant_bits): it must
be cached and handed out read-only so callers cannot mutate the shared copy."""
def test_table_is_cached_same_object(self):
self.assertIs(_fnuz_decode_table(4, 3), _fnuz_decode_table(4, 3))
self.assertIs(_fnuz_decode_table(5, 2), _fnuz_decode_table(5, 2))
def test_table_is_read_only(self):
table = _fnuz_decode_table(4, 3)
self.assertFalse(table.flags.writeable)
with self.assertRaises(ValueError):
table[0] = 1.0
if __name__ == "__main__":
unittest.main()

327
tile_engine/ops/gemm/grouped_gemm_full_benchmark.py Normal file
View File

@@ -0,0 +1,327 @@
#!/usr/bin/env python3
"""Full GROUPED GEMM benchmark sweep driven through the Dispatcher bridge.
The grouped counterpart of gemm_full_benchmark.py. Same 3-phase architecture:
Phase 1: Compile all grouped kernels (parallel, returns .so paths only -- no GPU)
Phase 2: Load grouped problems (lists of (M, N, K) sub-problems)
Phase 3: Benchmark via subprocess isolation (serial GPU, batched)
Tile Engine generates NO binaries here: it expands its grouped sweep config into
shared ``GemmKernelConfig(variant="grouped")`` objects and hands them to the
dispatcher, which codegens + compiles each into a grouped .so (built against
``grouped_gemm_ctypes_lib.cpp``). Each kernel runs in a disposable worker
subprocess so a GPU fault takes down only one worker.
Usage:
python grouped_gemm_full_benchmark.py grouped_gemm/configs/default_config.json \
--arch gfx942 --csv grouped_gemm_results.csv
A grouped "problem" is a batch of sub-problems run by one launch. The default set
uses uniform shapes at two group counts; override with --problems pointing at a
JSON file of [{"groups": [[M,N,K], ...]}, ...].
"""
import argparse
import csv
import json
import os
import subprocess
import sys
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
# Default grouped problem set: uniform groups at two batch sizes (Phase 3 parity).
DEFAULT_PROBLEMS = [
{"groups": [[1024, 1024, 1024]] * 4},
{"groups": [[2048, 2048, 2048]] * 4},
{"groups": [[1024, 1024, 1024]] * 8},
{"groups": [[512, 1024, 2048], [1024, 512, 1024], [2048, 2048, 512], [256, 768, 1024]]},
]
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 _problem_dims(prob):
"""Aggregate (group_count, total_flops) for reporting."""
groups = prob["groups"]
flops = sum(2.0 * m * n * k for (m, n, k) in groups)
return len(groups), flops
def main():
parser = argparse.ArgumentParser(
description="Grouped GEMM Benchmark Sweep (via Dispatcher)"
)
parser.add_argument("configs", nargs="+", help="TE grouped sweep config JSON files")
parser.add_argument("--arch", default="gfx942")
parser.add_argument("--dtype", default="fp16")
parser.add_argument("--layout", default="rcr")
parser.add_argument(
"--problems", default=None, help="JSON file of grouped problems"
)
parser.add_argument("--csv", type=str, default="grouped_gemm_results.csv")
parser.add_argument("--workers", type=int, default=8, help="Parallel build workers")
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=60, help="Per-kernel timeout (s)"
)
parser.add_argument(
"--max-kernels", type=int, default=0, help="Limit to first N kernels (0=all)"
)
args = parser.parse_args()
# ========================================================================
# Phase 1: Compile kernels (parallel, no GPU)
# ========================================================================
print(f"\n{'=' * 80}")
print("Phase 1: Compile grouped kernels")
print(f"{'=' * 80}")
all_configs = []
for cfg_path in args.configs:
all_configs.extend(
expand_sweep(
cfg_path,
args.arch,
dtype=args.dtype,
layout=args.layout,
variant="grouped",
)
)
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 grouped 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 via subprocess (serial GPU, batched)
# ========================================================================
print(f"\n{'=' * 80}")
print("Phase 3: Benchmark (subprocess isolation, batched)")
print(f"{'=' * 80}")
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",
"group_count",
"total_flops",
"latency_ms",
"tflops",
"non_zero",
]
csv_file = open(csv_path, "w", newline="")
writer = csv.DictWriter(csv_file, fieldnames=csv_fields)
writer.writeheader()
worker_path = _THIS_DIR / "run_one_grouped_gemm_kernel.py"
worker_env = os.environ.copy()
worker_env["GEMM_PYPATH"] = os.pathsep.join(
[str(_DISPATCHER_ROOT / "python"), str(_THIS_DIR)]
)
total_measurements = 0
total_failures = 0
bench_t0 = time.perf_counter()
for prob_idx, prob in enumerate(problems):
group_count, total_flops = _problem_dims(prob)
print(
f"\nProblem [{prob_idx + 1}/{len(problems)}]: groups={group_count} "
f"flops={total_flops / 1e9:.1f}G ({len(built_kernels)} kernels)"
)
print(f" {'Kernel':<60} {'Time(ms)':>10} {'TFLOPS':>10} {'Status':>10}")
print(f" {'-' * 95}")
prob_dict = {"groups": [list(g) for g in prob["groups"]]}
for batch_start in range(0, len(built_kernels), args.batch_size):
batch_end = min(batch_start + args.batch_size, len(built_kernels))
batch = built_kernels[batch_start:batch_end]
items = [
{
"so_path": str(lib_path),
"problem": prob_dict,
"kernel_name": cfg.name,
"dtype": args.dtype,
"layout": args.layout,
}
for cfg, lib_path in batch
]
payload = json.dumps({"items": items})
try:
proc = subprocess.Popen(
[sys.executable, str(worker_path)],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
env=worker_env,
)
timeout_total = args.kernel_timeout * len(batch)
stdout_bytes, _ = proc.communicate(
input=payload.encode("utf-8"), timeout=timeout_total
)
reported_indices = set()
for line in stdout_bytes.decode("utf-8").strip().split("\n"):
if not line:
continue
try:
result = json.loads(line)
batch_idx = result.get("idx")
if not isinstance(batch_idx, int) or not (
0 <= batch_idx < len(batch)
):
print(
f" [warn] worker returned out-of-range idx "
f"{batch_idx!r}; skipping line"
)
continue
cfg, lib_path = batch[batch_idx]
reported_indices.add(batch_idx)
if result.get("ok", False):
status = "OK" if result.get("non_zero", 0) > 0 else "ZERO"
print(
f" {cfg.name:<60} {result['ms']:>10.3f} "
f"{result['tflops']:>10.2f} {status:>10}"
)
writer.writerow(
{
"kernel": cfg.name,
"problem_idx": prob_idx,
"group_count": group_count,
"total_flops": total_flops,
"latency_ms": result["ms"],
"tflops": result["tflops"],
"non_zero": result.get("non_zero", 0),
}
)
csv_file.flush()
total_measurements += 1
else:
print(f" {cfg.name:<60} FAILED")
print(f" Error: {result.get('error', 'unknown')[:100]}")
total_failures += 1
except json.JSONDecodeError:
print(f" Warning: could not parse result line: {line[:50]}")
total_failures += 1
missing_indices = set(range(len(batch))) - reported_indices
if missing_indices or proc.returncode != 0:
if proc.returncode != 0:
print(f" Worker exited with code {proc.returncode}")
for idx in sorted(missing_indices):
cfg, _ = batch[idx]
print(f" {cfg.name:<60} MISSING (worker crash)")
total_failures += len(missing_indices)
except subprocess.TimeoutExpired:
print(f" Batch timeout ({len(batch)} kernels)")
try:
proc.kill()
proc.communicate(timeout=5)
except Exception:
pass
total_failures += len(batch)
except Exception as e:
print(f" Batch error: {e}")
try:
if proc and proc.poll() is None:
proc.kill()
except Exception:
pass
total_failures += len(batch)
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" Successful measurements: {total_measurements}")
print(f" Failed measurements: {total_failures}")
print(f" Output: {csv_path}")
return 0
if __name__ == "__main__":
sys.exit(main())

130
tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py Normal file
View File

@@ -0,0 +1,130 @@
#!/usr/bin/env python3
"""Worker script for running GROUPED GEMM kernels in an isolated subprocess.
Mirrors run_one_gemm_kernel.py, but for the grouped variant:
- Receives a grouped kernel config + grouped 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
A grouped problem is a LIST of (M, N, K) sub-problems run by one grouped launch.
Input JSON format (dtype/layout default to fp16/rcr if absent):
Single: {"so_path": "...", "problem": {"groups": [[M,N,K], ...]},
"kernel_name": "...", "dtype": "fp16", "layout": "rcr"}
Batch: {"items": [{"so_path": "...", "problem": {...}, "kernel_name": "...",
"dtype": "...", "layout": "..."}, ...]}
Output JSON format (one line per kernel):
{"idx": 0, "ok": true, "ms": 0.123, "tflops": 456.7, "non_zero": 1,
"group_count": 8, "kernel": "..."}
{"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 GroupedGemmProblem, GpuGroupedGemmRunner # noqa: E402
import numpy as np # noqa: E402
def _run_one(idx, so_path, prob_dict, kernel_name, dtype="fp16", layout="rcr"):
"""Run a single grouped kernel and emit its result as one JSON line."""
try:
problem = GroupedGemmProblem.from_dict(prob_dict)
# Operands are generated as fp32-ish floats; the runner casts them to the
# per-dtype codec (fp16/bf16/fp8/bf8) and applies layout transposes.
np.random.seed(42)
A_list = []
B_list = []
for (M, N, K) in problem.groups:
A_list.append((np.random.randn(M, K) * 0.1).astype(np.float32))
B_list.append((np.random.randn(K, N) * 0.1).astype(np.float32))
# CRITICAL: load the library ONLY inside this subprocess.
runner = GpuGroupedGemmRunner(lib_path=so_path, dtype=dtype, layout=layout)
result = runner.run(A_list, B_list, problem)
if result.success:
non_zero = sum(
int(np.count_nonzero(o)) for o in result.outputs if o is not None
)
print(
json.dumps(
{
"idx": idx,
"ok": True,
"ms": result.time_ms,
"tflops": result.tflops,
"non_zero": 1 if non_zero > 0 else 0,
"group_count": problem.group_count,
"kernel": kernel_name,
}
),
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 grouped 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)
if "items" in d:
for i, item in enumerate(d["items"]):
_run_one(
i,
item["so_path"],
item["problem"],
item.get("kernel_name", "unknown"),
item.get("dtype", "fp16"),
item.get("layout", "rcr"),
)
else:
_run_one(
0,
d["so_path"],
d["problem"],
d.get("kernel_name", "unknown"),
d.get("dtype", "fp16"),
d.get("layout", "rcr"),
)
if __name__ == "__main__":
main()