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#8985 (commit 3d4cbef)

Browse Source 
feat(ck-tile): add stream_k variant to GEMM Dispatcher
 codegen (#8985)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

> Supersedes #8094 (closed when its branch was renamed to a
policy-compliant path). Same commits, same head SHA.

## Motivation

This is the next slice of the Tile Engine → Dispatcher consolidation,
following the same pattern as the grouped_gemm PR (#8075). It adds the
**stream-K** GEMM variant to the unified GEMM codegen, implemented **the
dispatcher way** (workspace owned internally via `DeviceMem`, clean
`launch(args, stream)` signature), and proves numeric + performance
parity against Tile Engine.

Branch is based on `develop` and contains **only** the stream-K work (no
grouped_gemm commits).

## Technical Details

- **`codegen/arch_filter.py`** — added `OperatorType.GEMM_STREAMK` and
its tile constraints.
- **`codegen/unified_gemm_codegen.py`**:
- Added `GemmVariant.STREAM_K`, made it reachable from the CLI
(`--variants stream_k`), wired naming (`_streamk` suffix), includes, and
the variant→operator map.
- New `_launch_function_streamk`: builds a single `StreamKHostArgs`,
`MakeKernelArgs` → `GetWorkSpaceSize` → allocate `DeviceMem` workspace
**internally** + `SetZero` → `SetWorkSpacePointer` →
`IsSupportedArgument` check → `make_kernel` via
`launch_kernel_time_mask` with an Atomic-reduction preprocess that zeros
C between timed iterations. No external `kargs_ptr` (not the Tile Engine
way).
- Exported `A/B/CLayout` in the `CK_TILE_SINGLE_KERNEL_INCLUDE` block so
a single-kernel driver is layout-generic.
- Restricted stream_k configs to the `cshuffle` epilogue (only one the
kernel supports).
- **`examples/gemm/cpp/03_streamk_gemm_driver.cpp`** (NEW) — minimal
standalone driver: `-include`s one generated stream-K header, builds a
single A/B/C tensor, calls `SelectedKernel::launch(args, stream)`,
verifies against `ck_tile::reference_gemm`, prints TFLOPS/GB/s.

The generated GPU kernel (`StreamKKernel<StreamKTilePartitioner,
GemmPipeline, GemmEpilogue>`) is identical to TE's; only host-side
workspace ownership differs (internal `DeviceMem` vs TE's external
pointer). Numerics match.

## Test Plan

- **Config:**
`fp16_rcr_compv3_cshuffle_intrawave_..._128x128x64_2x2x1_32x32x16`
(atomic reduction; exists identically in TE and the dispatcher).
- **Shape:** `M=3840, N=4096, K=2048`, `warmup=10`, `repeat=50`, MI300X
(gfx942), ROCm 7.1.1.
- Run the `03_streamk_gemm_driver` and verify against
`ck_tile::reference_gemm`; compare latency/TFLOPS/GB/s against the
matching Tile Engine config.

> Methodology note: TE's benchmark forces `repeat=1, warmup=0` whenever
`verify=1` (the atomic kernel accumulates into C, so it can only verify
a single run). A `verify=1` invocation therefore reports a single cold
iteration (~0.30 ms), which is **not** a representative perf number. The
table below uses TE `verify=0` (so warmup/repeat are honored) for the
perf row and a separate TE `verify=1` run for correctness. The
dispatcher driver times (warmup=10/repeat=50) and verifies in the same
run because it re-zeros C between timed iterations via the masked
preprocess.

## Test Result

Performance + numerical verification (Dispatcher vs Tile Engine):

| | latency (ms) | TFLOPS | GB/s | verify |
|---|---|---|---|---|
| **Tile Engine** (warmup=10, repeat=50) | 0.24 | 266.7 | 264.8 |
correct |
| **Dispatcher** (warmup=10, repeat=50) | 0.242 | 266.1 | 264.2 | PASS |
| **Δ** | ~0% | ~0% | ~0% | identical |

## Next

- Once signed off, delete `tile_engine/ops/gemm_streamk/`.
- Continue toward a first-class `dispatcher` GEMM interface folder
(roadmap step 5).
This commit is contained in:
Muhammed Emin Ozturk
2026-07-15 16:12:14 +00:00
committed by assistant-librarian[bot]
parent 8c5870f962
commit 5d3380aa30
15 changed files with 1770 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
dispatcher/README.md
View File

@@ -4,6 +4,8 @@ A unified kernel dispatch system for AMD GPUs with C++ and Python frontends, sup
**Validated Platform:** AMD Instinct MI300 series (gfx942)
> **Stream-K GEMM:** see [STREAMK.md](STREAMK.md) for how to generate, build, run, and
> test the Stream-K deep-core path (atomic/linear/tree reductions).
---

230
dispatcher/STREAMK.md Normal file
View File

@@ -0,0 +1,230 @@
# Stream-K GEMM (Dispatcher Deep-Core Path)
Stream-K is a single GEMM that splits the **K** dimension across compute units (CUs)
and reduces the partial results, instead of giving each CU a whole output tile. It
keeps every CU busy on shapes where a classic data-parallel tiling would leave some
idle (tall-skinny / large-K problems), at the cost of a reduction step.
This document explains how to **generate**, **build**, **run**, and **test** the
Stream-K kernels through the CK Tile dispatcher.
> **Validated platform:** AMD Instinct MI300X (gfx942). See [Known limitations](#known-limitations)
> for gfx950 (MI350) status.
---
## Why Stream-K needs its own path
A plain GEMM rides `Dispatcher::run(A, B, C, problem)`. Stream-K cannot use that
signature unchanged: it needs a **reduction workspace** and a **reduction strategy**,
so its host args type (`ck_tile::StreamKHostArgs`) is ABI-incompatible with the
regular `GemmHostArgs`. The deep-core path makes Stream-K ride the registry anyway:
```
codegen (unified_gemm_codegen.py)
-> generated Stream-K kernel + dispatcher wrapper
-> Registry::register_kernel(GeneratedStreamKKernelInstance)
-> Dispatcher::select_kernel(Problem.streamk + reduction_strategy)
-> GeneratedStreamKKernelInstance::run() (Dispatcher owns the workspace)
-> SelectedKernel::launch(StreamKHostArgs, cfg, workspace)
```
### Reduction strategies
The reduction strategy is a **compile-time** property, so each strategy is a
*distinct kernel*. The registry holds all three side by side and the dispatcher
selects by `Problem::reduction_strategy`:
| Strategy | Workspace | Identifier suffix | Notes |
|---|---|---|---|
| `atomic` | none | `_streamk` | partials accumulate directly into C via atomics |
| `linear` | yes | `_streamk_linear` | partials reduced through a device workspace, in order |
| `tree` | yes | `_streamk_tree` | tree reduction through a device workspace |
### Supported datatypes / layouts
- **Datatypes:** `fp16`, `bf16`, `fp8`, `bf8`. (`fp32`/`fp64` have no MFMA warp tiles;
`int8` Stream-K is out of scope for this path.)
- **Layouts:** `rcr`, `rrr`, `ccr`, `crr` — A/B in either order, **C is row-major**
(the atomic C-reset relies on it).
---
## Prerequisites
A full ROCm toolchain with HIP headers (`hip/hip_runtime.h`) and `hipcc`. Bare SLURM
compute nodes on the cluster often ship an incomplete ROCm, so build inside the CK
ROCm container, e.g.:
```bash
# on a GPU node (pyxis/enroot), mounting your home:
srun --jobid=<JOBID> --overlap \
--container-image=/cluster/images/ck/ck_rocm7.1.1_therock_<date>.sqsh \
--container-mounts=$HOME:$HOME \
bash -lc '<commands below>'
```
---
> All commands below are run from the dispatcher root
> (`projects/composablekernel/dispatcher`).
## 1. Generate a Stream-K kernel
The codegen emits all three reduction-strategy headers from one tile config:
```bash
python3 codegen/unified_gemm_codegen.py \
--datatype fp16 --layout rcr \
--gpu-target gfx942 \
--variants stream_k \
--tile-config-json '{
"tile_config": {"tile_m":[128],"tile_n":[128],"tile_k":[64],
"warp_m":[2],"warp_n":[2],"warp_k":[1],
"warp_tile_m":[32],"warp_tile_n":[32],"warp_tile_k":[16],
"block_size":[256]},
"trait_config": {"pipeline":["compv3"],"epilogue":["cshuffle"],"scheduler":["intrawave"],
"pad_m":[false],"pad_n":[false],"pad_k":[false],"persistent":[false]},
"streamk_config": {"reduction_strategy":["atomic","linear","tree"]}
}' \
--output-dir ./gen_fp16_rcr
```
This produces, per strategy, a header named:
```
gemm_<dtype>_<layout>_compv3_cshuffle_intrawave_<padM>_<padN>_<padK>_<persistent>_<TILE>_<variant>.hpp
# variant ∈ { streamk, streamk_linear, streamk_tree }
```
Each header force-includes into the global namespace: `SelectedKernel`,
`ADataType/BDataType/CDataType/AccDataType`, `ALayout/BLayout/CLayout`, `KERNEL_NAME`.
Omit `--tile-config-json` to generate the full arch-filtered tile set instead of a
single config. Use `--show-arch-info` to print what a target GPU supports.
---
## 2a. Run via the standalone driver (`03_streamk_gemm_driver.cpp`)
Calls `SelectedKernel::launch()` **directly** (bypasses the dispatcher). Use this for
apple-to-apple performance measurement against Tile Engine.
```bash
HDR=gen_fp16_rcr/gemm_fp16_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16_streamk.hpp
hipcc -std=c++17 --offload-arch=gfx942 -O3 \
-DCK_TILE_SINGLE_KERNEL_INCLUDE \
-I ../include -I gen_fp16_rcr \
-include "$HDR" \
examples/gemm/cpp/03_streamk_gemm_driver.cpp -o streamk_gemm_driver
# performance (cold cache, TE-matched defaults):
./streamk_gemm_driver --m 4096 --n 4096 --k 4096 --validate 0
# correctness (single cold shot so C matches the reference):
./streamk_gemm_driver --m 4096 --n 4096 --k 4096 --validate 1
```
| Option | Default | Meaning |
|---|---|---|
| `--m/--n/--k` | 3840/4096/2048 | GEMM dims |
| `--warmup` | 50 | warmup iterations (timing) |
| `--repeat` | 100 | timed iterations |
| `--validate` | 1 | verify vs `reference_gemm`; forces 1 cold shot, no rotation |
| `--timer` | 1 | use the GPU timer |
| `--flush_cache` | 1 | flush L2 each iter (cold measurement, like Tile Engine) |
| `--rotating_count` | 1000 | rotating input copies to defeat cache (Tile Engine default) |
> **Methodology:** leaving the cache warm over-reports TFlops and is the entire
> source of spurious "dispatcher vs Tile Engine" perf gaps. Always measure perf with
> the cold-cache defaults (`--validate 0`); run correctness separately (`--validate 1`).
---
## 2b. Run via the registry/dispatcher (`04_streamk_registry_driver.cpp`)
Exercises the **full deep-core path**: registers the kernel, lets the dispatcher
select it by `Problem::reduction_strategy`, runs it (dispatcher owns the workspace),
and verifies vs the reference with a **split-K-aware tolerance**.
```bash
HDR=gen_fp16_rcr/gemm_fp16_rcr_compv3_cshuffle_intrawave_False_False_False_False_128x128x64_2x2x1_32x32x16_streamk.hpp
# core objects (once, no force-include):
hipcc -std=c++17 --offload-arch=gfx942 -O3 -I ../include -I include -c src/dispatcher.cpp -o dispatcher.o
hipcc -std=c++17 --offload-arch=gfx942 -O3 -I ../include -I include -c src/registry.cpp -o registry.o
# driver (force-include one strategy's header):
hipcc -std=c++17 --offload-arch=gfx942 -O3 \
-DCK_TILE_SINGLE_KERNEL_INCLUDE -DGFX_ARCH='"gfx942"' \
-I ../include -I include -I gen_fp16_rcr -include "$HDR" \
-c examples/gemm/cpp/04_streamk_registry_driver.cpp -o drv04.o
hipcc --offload-arch=gfx942 drv04.o dispatcher.o registry.o -o streamk_registry_driver
./streamk_registry_driver --m 3840 --n 4096 --k 2048 --strategy atomic --validate 1
```
| Option | Default | Meaning |
|---|---|---|
| `--m/--n/--k` | 3840/4096/2048 | GEMM dims |
| `--strategy` | atomic | `atomic` / `linear` / `tree` (must match the force-included header) |
| `--validate` | 1 | verify vs `reference_gemm` (split-K-aware rtol/atol) |
> The registry `run()` path is a functional dispatch path; its `Perf:` line is a
> cold-but-**non-rotated** measurement, **not** the calibrated apple-to-apple surface.
> Use the `03` driver (`--validate 0`) for Tile-Engine-comparable numbers.
---
## 3. Test (CTest)
The deep-core path is guarded by `test_streamk_registry.py`, which generates, builds,
dispatches, and verifies every `datatype × layout × strategy` against two shapes
(the default plus a small-M/large-K shape that stresses the split-K tolerance). It
**SKIPs** (exit 77) when no GPU or `hipcc` is present.
```bash
# directly:
python3 tests/test_streamk_registry.py --arch gfx942
python3 tests/test_streamk_registry.py --arch gfx942 --datatypes fp16,bf16 --layouts rcr,ccr
# via ctest (from your dispatcher build dir):
ctest -R dispatcher_test_streamk_registry --output-on-failure
```
---
## Verification tolerance (why Stream-K is special)
Stream-K reduces `kbatch` partial products into each output element, so the
accumulation error is larger than a single-pass GEMM. The drivers use the same
split-K-aware tolerance as Tile Engine (`calculate_rtol_atol`): `kbatch` is taken
from the kernel's own tile partitioner, and the tolerance is
`max(per-split threshold, split-K-reduction threshold)`. Using the plain
`get_relative/absolute_threshold(K)` here spuriously FAILs correct atomic results on
small-M/N, large-K shapes.
---
## Known limitations
- **gfx950 (MI350) fp8/bf8 not validated.** On CDNA4 the fp8/bf8 host reference/codec
hits an FNUZ-vs-OCP format mismatch; those combos currently fail verification. fp16
and bf16 are fine on gfx950. Validate/gate before enabling fp8/bf8 there.
- **Tile coverage is narrower than Tile Engine.** The dispatcher emits fewer Stream-K
tiles than TE (e.g. fp16 `rcr` TE=180 vs DISP=73). Numeric+perf parity is validated
per matched tile config, not over the whole TE tile surface. See the coverage note
at the `STREAM_K` variant in `codegen/unified_gemm_codegen.py`.
---
## File map
| Path | Role |
|---|---|
| `codegen/unified_gemm_codegen.py` | generates Stream-K kernels + dispatcher wrappers (`--variants stream_k`) |
| `include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp` | `GeneratedStreamKKernelInstance` (registry/workspace/launch glue) |
| `include/ck_tile/dispatcher/kernel_key.hpp` | registry key carrying `streamk` + `reduction_strategy` |
| `examples/gemm/cpp/03_streamk_gemm_driver.cpp` | standalone driver (direct `launch`, perf surface) |
| `examples/gemm/cpp/04_streamk_registry_driver.cpp` | deep-core driver (Registry → Dispatcher → verify) |
| `tests/test_streamk_registry.py` | CTest `dispatcher_test_streamk_registry` |

15
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_STREAMK = "gemm_streamk"
CONV_FWD = "conv_fwd"
CONV_BWD_DATA = "conv_bwd_data"
CONV_BWD_WEIGHT = "conv_bwd_weight"
@@ -85,6 +86,20 @@ OPERATOR_TILE_CONSTRAINTS = {
"tile_n_alignment": 16,
"tile_k_alignment": 8,
},
# 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
# runtime (StreamKKernel::IsSupportedArgument / the backend supports() check),
# which lets the dispatcher fall back to a non-Stream-K kernel for too-small
# problems instead of rejecting them at codegen time.
OperatorType.GEMM_STREAMK: {
"min_tile_m": 16,
"min_tile_n": 16,
"min_tile_k": 8,
"tile_m_alignment": 16,
"tile_n_alignment": 16,
"tile_k_alignment": 8,
},
OperatorType.CONV_FWD: {
"min_tile_m": 1, # N dimension can be 1
"min_tile_n": 16, # K (output channels) should be reasonable

271
dispatcher/codegen/unified_gemm_codegen.py
View File

@@ -202,6 +202,21 @@ class GemmVariant(Enum):
STANDARD = "standard"
PRESHUFFLE = "preshuffle"
MULTI_D = "multi_d"
# 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:
# measured per layout, e.g. fp16/bf16 rcr TE=180 vs DISP=73 tiles (124 TE-only,
# 17 DISP-only); ccr TE=144 vs DISP=73; fp8/bf8 closer (rcr TE=296 vs DISP=232)
# but still short. TE-vs-DISP numeric+perf parity is therefore validated
# per matched tile config, NOT over the whole TE tile space -- "functional
# equivalence" should be read with that scope. Closing the gap means feeding
# the missing TE tiles into the tile list (the codegen handles them); the
# divergent DISP-only tiles are configs TE does not enumerate at all.
# NOTE: this limitation is inherent only to driving the codegen standalone.
# When the bridge is implemented on top of this codegen, the tile list is
# supplied by Tile-Engine directly, so the emitted Stream-K surface matches
# the full Old-TE tile space by construction and the gap closes.
STREAM_K = "stream_k"
# TileConfig imported from codegen_common
@@ -227,6 +242,9 @@ class KernelConfig:
elementwise_op: str = "PassThrough"
num_d_tensors: int = 0
d_layout: str = "r" # Layout for D tensors (r=row, c=col) - same for all D tensors
# Stream-K reduction strategy: "atomic" (partials atomic-add into C),
# "linear", or "tree" (partials accumulate through a device workspace).
reduction_strategy: str = "atomic"
# Fixed parameters
block_size: int = 256
@@ -289,6 +307,11 @@ class KernelConfig:
parts.append(f"nd{self.num_d_tensors}")
parts.append(f"dly_{self.d_layout}")
# Stream-K variant: reduction strategy distinguishes otherwise-identical
# kernels (each strategy is a separate compiled binary).
if self.variant == GemmVariant.STREAM_K:
parts.append(f"redux_{self.reduction_strategy}")
# Occupancy parameters (only if non-default)
if self.num_wave_groups != 1:
parts.append(f"wg{self.num_wave_groups}")
@@ -344,6 +367,12 @@ 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.STREAM_K:
name += "_streamk"
# Atomic keeps the bare "_streamk" suffix for name parity with the
# original single-strategy bridge; linear/tree are disambiguated.
if config.reduction_strategy != "atomic":
name += f"_{config.reduction_strategy}"
return name
@@ -397,6 +426,15 @@ class CKTileKernelGenerator:
includes += """
#include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp"
#include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp"
"""
if config.variant == GemmVariant.STREAM_K:
includes += """
#include <functional>
#include <hip/hip_runtime.h>
#include "ck_tile/host/device_memory.hpp"
#include "ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp"
#include "ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp"
"""
return includes
@@ -526,6 +564,9 @@ using ADataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.datatype]};
using BDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.datatype]};
using CDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[output_dtype]};
using AccDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[acc_dtype]};
using ALayout = {ns_name}::ALayout;
using BLayout = {ns_name}::BLayout;
using CLayout = {ns_name}::CLayout;
// KernelKey field descriptors for the force-included kernel.
// The ctypes library builds the registry KernelKey from these so the
@@ -588,6 +629,8 @@ using AccDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[acc_dtype]};
"""Generate launch function"""
if config.variant == GemmVariant.MULTI_D:
return self._launch_function_multi_d(config)
if config.variant == GemmVariant.STREAM_K:
return self._launch_function_streamk(config)
if config.preshuffle:
return self._launch_function_preshuffle(config)
return self._launch_function_standard(config)
@@ -768,6 +811,162 @@ using AccDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[acc_dtype]};
return launch(multi_d_args, stream);
}}"""
def _launch_function_streamk(self, config: KernelConfig) -> str:
"""Generate launch function for Stream-K GEMM (the dispatcher way).
Stream-K is a single GEMM that splits the K dimension across CUs and
reduces partial results through a device workspace. Unlike Tile Engine
(which takes an external workspace pointer), the dispatcher allocates the
workspace INTERNALLY via DeviceMem inside launch(args, stream).
The reduction strategy is taken from the config (atomic/linear/tree).
Atomic: partial tiles atomic-add into C, so C is zeroed before every
kernel invocation. Linear/Tree: partials accumulate through the device
workspace, which is zeroed instead. Both are handled by the preprocess
callback passed to launch_kernel_time_mask.
"""
reduction_ck = {
"atomic": "Atomic",
"linear": "Linear",
"tree": "Tree",
}[config.reduction_strategy]
# The Atomic strategy zeroes C with a row-major hipMemset2D (pitch =
# stride_E rows of N elems). A column-major C would be zeroed incorrectly
# and atomic accumulation would then corrupt results, so fail loudly at
# compile time rather than silently. Linear/Tree zero the workspace, not C,
# so they carry no such requirement.
c_rowmajor_assert = (
"""
static_assert(
std::is_same_v<ck_tile::remove_cvref_t<CLayout>,
ck_tile::tensor_layout::gemm::RowMajor>,
"Stream-K Atomic reduction requires a row-major C: the hipMemset2D C-reset "
"assumes row-major layout and would zero a column-major C incorrectly.");
"""
if config.reduction_strategy == "atomic"
else ""
)
return f"""{c_rowmajor_assert}
// ---- Stream-K kernel type, hoisted to struct scope so the workspace API
// ---- (GetWorkSpaceSize + external-workspace launch) can reuse the same type. ----
static constexpr auto SkScheduler = {self.tm.SCHEDULER_TO_CK[config.trait.scheduler]};
static constexpr auto SkReductionStrategy = ck_tile::StreamKReductionStrategy::{reduction_ck};
static constexpr int SkBlockPerCu = {config.k_block_per_cu};
using SkGemmUniversalTraits = TileGemmUniversalTraits<kPadM, kPadN, kPadK, DoubleSmemBuffer,
ALayout, BLayout, CLayout, TransposeC,
UseStructuredSparsity, UsePersistentKernel,
NumWaveGroups, Preshuffle>;
using SkUniversalGemmProblem = UniversalGemmPipelineProblem<
ADataType, BDataType, AccDataType, TileShape, SkGemmUniversalTraits, SkScheduler>;
using SkGemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}<SkUniversalGemmProblem>;
{self._epilogue_code(config)}
using SkStreamKTilePartitioner =
ck_tile::StreamKTilePartitioner<TileShape, SkReductionStrategy, UsePersistentKernel>;
using StreamKGemmKernel =
ck_tile::StreamKKernel<SkStreamKTilePartitioner, SkGemmPipeline, GemmEpilogue>;
// Device workspace (bytes) this kernel needs for `args`. 0 for Atomic;
// >0 for Linear/Tree. The Dispatcher uses this to size the buffer it owns.
static std::size_t GetWorkSpaceSize(const ck_tile::StreamKHostArgs& args) {{
auto kargs = StreamKGemmKernel::MakeKernelArgs(args);
return StreamKGemmKernel::GetWorkSpaceSize(kargs);
}}
// Whether the kernel can actually partition this problem (enough tiles across
// CUs). Lets the dispatcher's supports() reject too-small problems and fall
// back to a non-Stream-K kernel instead of throwing at launch.
static bool IsSupported(const ck_tile::StreamKHostArgs& args) {{
return StreamKGemmKernel::IsSupportedArgument(StreamKGemmKernel::MakeKernelArgs(args));
}}
// Internal-workspace launch: allocates a fresh DeviceMem on every call.
// Kept unchanged for the bridge ctypes lib and the standalone 03 driver.
static float launch(const ck_tile::StreamKHostArgs& args, const stream_config& stream) {{
auto kargs = StreamKGemmKernel::MakeKernelArgs(args);
const auto ws_size = StreamKGemmKernel::GetWorkSpaceSize(kargs);
ck_tile::DeviceMem workspace_dev(ws_size);
workspace_dev.SetZero();
StreamKGemmKernel::SetWorkSpacePointer(kargs, workspace_dev.GetDeviceBuffer());
if (!StreamKGemmKernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Arguments not supported for stream-k kernel!");
}}
const dim3 grids = StreamKGemmKernel::GridSize(kargs.tile_partitioner);
const dim3 blocks = StreamKGemmKernel::BlockSize();
// Atomic reduction accumulates into C, so reset buffers before each run.
auto reset_data_buffers = [&]() {{
if constexpr (SkReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic) {{
// Stride-aware: CLayout is row-major with stride_E elems/row, so a
// padded C is zeroed correctly (not just the contiguous M*N case).
if(hipMemset2DAsync(args.e_ptr,
args.stride_E * sizeof(CDataType),
0,
args.N * sizeof(CDataType),
args.M,
stream.stream_id_) != hipSuccess) {{
throw std::runtime_error(
"stream-k: hipMemset2DAsync failed to reset C between iterations");
}}
}} else {{
workspace_dev.SetZero();
}}
}};
std::function<void()> preprocess = reset_data_buffers;
float ave_time = launch_kernel_time_mask(stream, preprocess,
make_kernel<SkBlockPerCu>(StreamKGemmKernel{{}}, grids, blocks, 0, kargs));
return ave_time;
}}
// External-workspace launch (PR-D): the Dispatcher owns and reuses the
// reduction buffer and passes it in. `workspace` may be null for Atomic
// (size 0). The per-iteration reset stays here because it needs CDataType
// and the reduction strategy, which the dtype-erased Dispatcher lacks.
static float launch(const ck_tile::StreamKHostArgs& args, const stream_config& stream,
void* workspace) {{
auto kargs = StreamKGemmKernel::MakeKernelArgs(args);
const auto ws_size = StreamKGemmKernel::GetWorkSpaceSize(kargs);
if (workspace != nullptr) {{
StreamKGemmKernel::SetWorkSpacePointer(kargs, workspace);
}}
if (!StreamKGemmKernel::IsSupportedArgument(kargs)) {{
throw std::runtime_error("Arguments not supported for stream-k kernel!");
}}
const dim3 grids = StreamKGemmKernel::GridSize(kargs.tile_partitioner);
const dim3 blocks = StreamKGemmKernel::BlockSize();
auto reset_data_buffers = [&]() {{
if constexpr (SkReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic) {{
// Stride-aware: CLayout is row-major with stride_E elems/row, so a
// padded C is zeroed correctly (not just the contiguous M*N case).
if(hipMemset2DAsync(args.e_ptr,
args.stride_E * sizeof(CDataType),
0,
args.N * sizeof(CDataType),
args.M,
stream.stream_id_) != hipSuccess) {{
throw std::runtime_error(
"stream-k: hipMemset2DAsync failed to reset C between iterations");
}}
}} else {{
if(hipMemsetAsync(workspace, 0, ws_size, stream.stream_id_) != hipSuccess) {{
throw std::runtime_error(
"stream-k: hipMemsetAsync failed to reset reduction workspace");
}}
}}
}};
std::function<void()> preprocess = reset_data_buffers;
float ave_time = launch_kernel_time_mask(stream, preprocess,
make_kernel<SkBlockPerCu>(StreamKGemmKernel{{}}, grids, blocks, 0, kargs));
return ave_time;
}}"""
def _epilogue_code(self, config: KernelConfig) -> str:
"""Generate epilogue code"""
if config.variant == GemmVariant.MULTI_D:
@@ -820,12 +1019,49 @@ class DispatcherWrapperGenerator:
acc_dtype = self.tm.get_acc_dtype(self.datatype)
rel_path = kernel_path.relative_to(output_dir)
# Stream-K kernels need the Stream-K backend (StreamKHostArgs launch) and
# the SK key fields, so the registry can tell atomic/linear/tree apart and
# the right launch path compiles. All other variants use the regular backend.
is_streamk = config.variant == GemmVariant.STREAM_K
backend_inc = (
"generated_tile_backend_streamk.hpp"
if is_streamk
else "generated_kernel_backend.hpp"
)
sk_fields = ""
if is_streamk:
rs = {"atomic": "Atomic", "linear": "Linear", "tree": "Tree"}[
config.reduction_strategy
]
ws = str(config.reduction_strategy != "atomic").lower()
sk_fields = f"""
key.algorithm.pad_m = {str(config.trait.pad_m).lower()};
key.algorithm.pad_n = {str(config.trait.pad_n).lower()};
key.algorithm.pad_k = {str(config.trait.pad_k).lower()};
key.algorithm.streamk = true;
key.algorithm.reduction_strategy = ::ck_tile::dispatcher::ReductionStrategy::{rs};
key.algorithm.workspace = {ws};"""
if is_streamk:
ret_stmt = (
"return backends::create_generated_streamk_kernel<KernelStruct, "
"KernelStruct::ADataType, KernelStruct::BDataType, "
"KernelStruct::CDataType, KernelStruct::AccDataType>"
f'(key, "{kernel_name}");'
)
else:
ret_stmt = (
"return std::make_shared<backends::GeneratedKernelInstance<KernelStruct>>"
f'(key, "{kernel_name}");'
)
return f"""// SPDX-License-Identifier: MIT
// Auto-generated dispatcher wrapper
#pragma once
#include "ck_tile/dispatcher.hpp"
#include "ck_tile/dispatcher/backends/generated_kernel_backend.hpp"
#include "ck_tile/dispatcher/backends/{backend_inc}"
#include "{rel_path}"
namespace ck_tile {{
@@ -876,11 +1112,11 @@ inline KernelInstancePtr make_{kernel_name}(const std::string& gfx_arch = "gfx94
key.algorithm.persistent = {str(config.trait.persistent).lower()};
key.algorithm.preshuffle = {str(config.preshuffle).lower()};
key.algorithm.transpose_c = false;
key.algorithm.num_wave_groups = {config.num_wave_groups};
key.algorithm.num_wave_groups = {config.num_wave_groups};{sk_fields}
key.gfx_arch = gfx_arch;
return std::make_shared<backends::GeneratedKernelInstance<KernelStruct>>(key, "{kernel_name}");
{ret_stmt}
}}
}}}}}}
@@ -985,6 +1221,10 @@ class UnifiedGemmCodegen:
"elementwise_ops": ["MultiDAdd", "MultiDMultiply"],
"num_d_tensors": [1, 2],
},
"streamk_config": {
# Each reduction strategy compiles to a separate kernel binary.
"reduction_strategy": ["atomic", "linear", "tree"],
},
}
def generate_all(self, parallel: bool = True) -> Dict:
@@ -1119,6 +1359,24 @@ class UnifiedGemmCodegen:
continue
configs.append(KernelConfig(tile=tile, trait=trait, variant=variant))
elif variant == GemmVariant.STREAM_K:
# Stream-K reuses the standard trait space but requires the cshuffle
# epilogue (the only epilogue the stream-K kernel supports). Each
# reduction strategy (atomic/linear/tree) is a distinct compiled
# kernel, so we expand one config per requested strategy.
if trait.epilogue == "cshuffle":
streamk_cfg = self.config.get("streamk_config", {})
strategies = streamk_cfg.get("reduction_strategy", ["atomic"])
for reduction_strategy in strategies:
configs.append(
KernelConfig(
tile=tile,
trait=trait,
variant=variant,
reduction_strategy=reduction_strategy,
)
)
elif variant == GemmVariant.PRESHUFFLE:
# Preshuffle needs specific pipeline (preshufflev2) and scheduler (default)
# Skip configs that don't use preshuffle-compatible traits
@@ -1276,6 +1534,7 @@ class UnifiedGemmCodegen:
GemmVariant.STANDARD: OperatorType.GEMM,
GemmVariant.PRESHUFFLE: OperatorType.GEMM_PRESHUFFLE,
GemmVariant.MULTI_D: OperatorType.GEMM_MULTI_D,
GemmVariant.STREAM_K: OperatorType.GEMM_STREAMK,
}
operator = variant_to_operator.get(variant, OperatorType.GEMM)
@@ -1525,7 +1784,7 @@ def main():
parser.add_argument(
"--variants",
nargs="+",
choices=["standard", "preshuffle", "multi_d"],
choices=["standard", "preshuffle", "multi_d", "stream_k"],
default=["standard"],
help="Variants to generate",
)

148
dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp Normal file
View File

@@ -0,0 +1,148 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Minimal standalone stream-K GEMM driver (dispatcher way).
*
* Stream-K is a SINGLE GEMM that splits the K dimension across CUs and reduces
* the partial results through a device workspace. Like grouped GEMM it cannot
* ride the standard dispatcher.run(A,B,C,problem) path, so this driver includes
* a single generated stream-K kernel header (CK_TILE_SINGLE_KERNEL_INCLUDE) and
* calls SelectedKernel::launch(args, stream) directly with a single
* StreamKHostArgs -- the same 2-arg signature the dispatcher generates (the
* workspace is allocated INSIDE launch() via DeviceMem). It builds one A/B/C
* tensor, runs, and verifies against ck_tile::reference_gemm.
*
* Build (single-kernel include style):
* hipcc -std=c++17 --offload-arch=gfx942 -O3 \
* -DCK_TILE_SINGLE_KERNEL_INCLUDE \
* -I <ck>/include -I <generated_dir> \
* -include <generated_dir>/<one>_streamk.hpp \
* 03_streamk_gemm_driver.cpp -o streamk_gemm_driver
*/
#include <hip/hip_runtime.h>
#include <algorithm>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/ops/gemm.hpp"
#include "streamk_driver_common.hpp"
// The generated stream-K 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.
int main(int argc, char** argv)
{
const ck_tile::index_t M = std::stoll(get_opt(argc, argv, "--m", "3840"));
const ck_tile::index_t N = std::stoll(get_opt(argc, argv, "--n", "4096"));
const ck_tile::index_t K = std::stoll(get_opt(argc, argv, "--k", "2048"));
int warmup = std::stoi(get_opt(argc, argv, "--warmup", "50"));
int repeat = std::stoi(get_opt(argc, argv, "--repeat", "100"));
const bool validate = get_opt(argc, argv, "--validate", "1") != "0";
// Apple-to-apple with tile_engine: time the kernel with the SAME methodology the
// tile_engine benchmark uses (gemm_streamk_profiler.hpp) -- gpu timer and a
// cold-cache measurement that flushes the cache and rotates input buffers each
// iteration. tile_engine defaults: timer=true, flush_cache=true, rotating_count=1000.
// Without these the driver measured a warm-cache best case and over-reported TFlops,
// which is the entire source of the dispatcher-vs-TE "performance gap".
const bool gpu_timer = get_opt(argc, argv, "--timer", "1") != "0";
bool flush_cache = get_opt(argc, argv, "--flush_cache", "1") != "0";
int rotating_count = std::stoi(get_opt(argc, argv, "--rotating_count", "1000"));
// Verification reads C back and compares against the reference for the known A/B.
// Rotating buffers and multi-repeat rotate/accumulate the output, so the C left on
// the device would not correspond to the reference inputs. tile_engine handles this
// with repeat_once_if_verify(); we mirror it -- a validating run times a single cold
// shot. Run a separate --validate 0 pass to collect apple-to-apple perf numbers.
if(validate)
{
warmup = 0;
repeat = 1;
flush_cache = false;
rotating_count = 1;
}
std::cout << "Kernel: " << KERNEL_NAME << "\n";
std::cout << "M=" << M << " N=" << N << " K=" << K << "\n";
const ck_tile::index_t sA = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{}));
const ck_tile::index_t sB = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{}));
const ck_tile::index_t sC = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{}));
ck_tile::HostTensor<ADataType> a_host(
ck_tile::host_tensor_descriptor(M, K, sA, is_row_major(ALayout{})));
ck_tile::HostTensor<BDataType> b_host(
ck_tile::host_tensor_descriptor(K, N, sB, is_row_major(BLayout{})));
ck_tile::HostTensor<CDataType> c_host(
ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{})));
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_host);
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_host);
c_host.SetZero();
ck_tile::DeviceMem a_dev(a_host);
ck_tile::DeviceMem b_dev(b_host);
ck_tile::DeviceMem c_dev(c_host);
c_dev.SetZero();
ck_tile::StreamKHostArgs args{a_dev.GetDeviceBuffer(),
b_dev.GetDeviceBuffer(),
c_dev.GetDeviceBuffer(),
M,
N,
K,
sA,
sB,
sC};
const ck_tile::stream_config s{
nullptr, true, /*log=*/0, warmup, repeat, gpu_timer, flush_cache, rotating_count};
float ave_time = SelectedKernel::launch(args, s);
const std::size_t flop = std::size_t(2) * M * N * K;
const std::size_t bytes =
sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
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";
c_dev.FromDevice(c_host.data());
bool pass = true;
if(validate)
{
ck_tile::HostTensor<CDataType> ref(
ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{})));
ref.SetZero();
ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(a_host, b_host, ref);
const float maxv = *std::max_element(ref.mData.begin(), ref.mData.end());
// num_wgs_per_tile is the number of workgroups reducing into a single
// output tile (Stream-K has no fixed split-k), taken from the kernel's
// own tile partitioner so the driver and tile_engine agree on the split
// factor. streamk_tolerance() then widens the verify tolerance for the
// split-K accumulation error (see streamk_driver_common.hpp).
using ComputeType =
std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
auto kargs = SelectedKernel::StreamKGemmKernel::MakeKernelArgs(args);
const ck_tile::index_t num_wgs_per_tile =
std::max<ck_tile::index_t>(1, kargs.tile_partitioner.estimate_num_wgs_per_tile());
const auto tol =
streamk_tolerance<ComputeType, CDataType, AccDataType>(K, num_wgs_per_tile, maxv);
pass = ck_tile::check_err(c_host, ref, "streamk", tol.rtol, tol.atol);
std::cout << "Verification: " << (pass ? "PASS" : "FAIL") << "\n";
}
return pass ? 0 : 1;
}

251
dispatcher/examples/gemm/cpp/04_streamk_registry_driver.cpp Normal file
View File

@@ -0,0 +1,251 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Stream-K GEMM driver through the Registry + Dispatcher (deep-core path).
*
* Unlike 03_streamk_gemm_driver.cpp (which calls SelectedKernel::launch()
* DIRECTLY, bypassing the dispatcher), this driver proves the full deep-core
* path that PR-A..PR-C built:
*
* Registry::register_kernel(GeneratedStreamKKernelInstance)
* -> Dispatcher::run(Problem.stream_k(Atomic))
* -> Dispatcher::select_first_fit -> SK instance.supports()
* -> GeneratedStreamKKernelInstance::run -> SelectedKernel::launch()
*
* It registers ONE generated Stream-K kernel (force-included via
* -include / -DCK_TILE_SINGLE_KERNEL_INCLUDE), selects it through the registry
* by Problem::reduction_strategy, runs it, and verifies vs reference_gemm.
*
* Build (single-kernel include style):
* hipcc -std=c++17 --offload-arch=gfx942 -O3 \
* -DCK_TILE_SINGLE_KERNEL_INCLUDE \
* -I <ck>/include -I <ck>/dispatcher/include -I <generated_dir> \
* -include <generated_dir>/<one>_streamk.hpp \
* 04_streamk_registry_driver.cpp -o streamk_registry_driver
*/
#include <hip/hip_runtime.h>
#include <algorithm>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <string>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/ops/gemm.hpp"
#include "ck_tile/dispatcher/dispatcher.hpp"
#include "ck_tile/dispatcher/registry.hpp"
#include "ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp"
#include "streamk_driver_common.hpp"
// The generated stream-K 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.
#ifndef GFX_ARCH
#define GFX_ARCH "gfx942"
#endif
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::backends;
using Priority = ck_tile::dispatcher::Registry::Priority;
// CLI parsing, layout/dtype tags, and the Stream-K verification tolerance are
// shared with the standalone 03 driver via streamk_driver_common.hpp
// (is_row_major, get_opt, dtype_enum_of, layout_tag_of, streamk_tolerance).
// Build the KernelKey for the force-included Stream-K kernel. Only the Stream-K
// axis (streamk + reduction_strategy) governs selection; the remaining fields
// are populated for a faithful encode_identifier()/registry entry.
static KernelKey make_streamk_key(ReductionStrategy strategy)
{
KernelKey key;
key.signature.dtype_a = dtype_enum_of<ADataType>();
key.signature.dtype_b = dtype_enum_of<BDataType>();
key.signature.dtype_c = dtype_enum_of<CDataType>();
key.signature.dtype_acc = dtype_enum_of<AccDataType>();
key.signature.layout_a = layout_tag_of<ALayout>();
key.signature.layout_b = layout_tag_of<BLayout>();
key.signature.layout_c = layout_tag_of<CLayout>();
key.signature.transpose_a = false;
key.signature.transpose_b = false;
key.signature.grouped = false;
// Stream-K performs its own K-dimension partitioning through the tile
// partitioner, so classic split-k is always 1 here. A value > 1 would
// describe a two-level K split the Stream-K kernel does not implement.
key.signature.split_k = 1;
key.signature.elementwise_op = "PassThrough";
key.signature.num_d_tensors = 0;
key.signature.structured_sparsity = false;
// Derive algorithm metadata from the generated kernel's own static traits so
// the registry identifier describes the kernel that was actually built,
// instead of assuming one tile/wave config. (Selection keys only on the
// Stream-K axis below, but a faithful identifier matters for logging and any
// future key-based lookup.)
key.algorithm.tile_shape = {
SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK};
key.algorithm.warp_tile_shape = {static_cast<std::uint8_t>(SelectedKernel::WarpTileM),
static_cast<std::uint8_t>(SelectedKernel::WarpTileN),
static_cast<std::uint8_t>(SelectedKernel::WarpTileK)};
key.algorithm.wave_shape = {static_cast<std::uint8_t>(SelectedKernel::WarpPerBlock_M),
static_cast<std::uint8_t>(SelectedKernel::WarpPerBlock_N),
static_cast<std::uint8_t>(SelectedKernel::WarpPerBlock_K)};
// Pipeline (CompV3) and scheduler (Intrawave) are baked into the generated
// kernel's type, not exposed as standalone enum values, and are not part of
// the Stream-K selection axis -- they stay at the codegen defaults.
key.algorithm.pipeline = Pipeline::CompV3;
key.algorithm.scheduler = Scheduler::Intrawave;
key.algorithm.epilogue = Epilogue::CShuffle;
key.algorithm.block_size = SelectedKernel::BlockSize;
key.algorithm.double_buffer = SelectedKernel::DoubleSmemBuffer;
key.algorithm.persistent = SelectedKernel::UsePersistentKernel;
key.algorithm.preshuffle = SelectedKernel::Preshuffle;
key.algorithm.transpose_c = SelectedKernel::TransposeC;
key.algorithm.num_wave_groups = SelectedKernel::NumWaveGroups;
key.algorithm.pad_m = SelectedKernel::kPadM;
key.algorithm.pad_n = SelectedKernel::kPadN;
key.algorithm.pad_k = SelectedKernel::kPadK;
// The Stream-K selection axis (the whole point of this path).
key.algorithm.streamk = true;
key.algorithm.reduction_strategy = strategy;
key.algorithm.workspace = (strategy != ReductionStrategy::Atomic);
key.gfx_arch = GFX_ARCH;
return key;
}
static ReductionStrategy parse_strategy(const std::string& s)
{
if(s == "linear")
return ReductionStrategy::Linear;
if(s == "tree")
return ReductionStrategy::Tree;
return ReductionStrategy::Atomic;
}
int main(int argc, char** argv)
{
const ck_tile::index_t M = std::stoll(get_opt(argc, argv, "--m", "3840"));
const ck_tile::index_t N = std::stoll(get_opt(argc, argv, "--n", "4096"));
const ck_tile::index_t K = std::stoll(get_opt(argc, argv, "--k", "2048"));
const bool validate = get_opt(argc, argv, "--validate", "1") != "0";
const ReductionStrategy strategy = parse_strategy(get_opt(argc, argv, "--strategy", "atomic"));
std::cout << "Kernel: " << KERNEL_NAME << "\n";
std::cout << "M=" << M << " N=" << N << " K=" << K << " strategy=" << to_string(strategy)
<< "\n";
// --- Register the kernel into the global registry ---------------------------
KernelKey key = make_streamk_key(strategy);
auto kernel = create_generated_streamk_kernel<SelectedKernel,
ADataType,
BDataType,
CDataType,
AccDataType>(key, KERNEL_NAME);
Registry::instance().clear();
Registry::instance().register_kernel(kernel, Priority::High);
std::cout << "Registered kernels: " << Registry::instance().size()
<< " identifier=" << key.encode_identifier() << "\n";
// --- Build the problem requesting THIS Stream-K strategy --------------------
Problem problem(M, N, K);
problem.streamk = true;
problem.reduction_strategy = strategy;
Dispatcher dispatcher;
auto selected = dispatcher.select_kernel(problem);
if(!selected)
{
std::cout << "Dispatcher selected NO kernel for the Stream-K problem -> FAIL\n";
return 1;
}
std::cout << "Dispatcher selected: " << selected->get_name() << "\n";
// --- Tensors (rcr) ---------------------------------------------------------
const ck_tile::index_t sA = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{}));
const ck_tile::index_t sB = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{}));
const ck_tile::index_t sC = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{}));
ck_tile::HostTensor<ADataType> a_host(
ck_tile::host_tensor_descriptor(M, K, sA, is_row_major(ALayout{})));
ck_tile::HostTensor<BDataType> b_host(
ck_tile::host_tensor_descriptor(K, N, sB, is_row_major(BLayout{})));
ck_tile::HostTensor<CDataType> c_host(
ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{})));
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_host);
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_host);
c_host.SetZero();
ck_tile::DeviceMem a_dev(a_host);
ck_tile::DeviceMem b_dev(b_host);
ck_tile::DeviceMem c_dev(c_host);
c_dev.SetZero();
// --- Run through the dispatcher (registry -> Dispatcher::run -> SK backend) -
float ave_time = 0.f;
try
{
ave_time = dispatcher.run(
a_dev.GetDeviceBuffer(), b_dev.GetDeviceBuffer(), c_dev.GetDeviceBuffer(), problem);
}
catch(const std::exception& e)
{
std::cout << "Dispatcher::run threw: " << e.what() << " -> FAIL\n";
return 1;
}
const std::size_t flop = std::size_t(2) * M * N * K;
const std::size_t bytes =
sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
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";
c_dev.FromDevice(c_host.data());
bool pass = true;
if(validate)
{
ck_tile::HostTensor<CDataType> ref(
ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{})));
ref.SetZero();
ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(a_host, b_host, ref);
const float maxv = *std::max_element(ref.mData.begin(), ref.mData.end());
// num_wgs_per_tile is the number of workgroups reducing into a single
// output tile (Stream-K has no fixed split-k), taken from the kernel's
// own tile partitioner so this driver and tile_engine agree on the split
// factor. streamk_tolerance() then widens the verify tolerance for the
// split-K accumulation error (see streamk_driver_common.hpp).
ck_tile::StreamKHostArgs sk_args{a_dev.GetDeviceBuffer(),
b_dev.GetDeviceBuffer(),
c_dev.GetDeviceBuffer(),
M,
N,
K,
sA,
sB,
sC};
using ComputeType =
std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
auto kargs = SelectedKernel::StreamKGemmKernel::MakeKernelArgs(sk_args);
const ck_tile::index_t num_wgs_per_tile =
std::max<ck_tile::index_t>(1, kargs.tile_partitioner.estimate_num_wgs_per_tile());
const auto tol =
streamk_tolerance<ComputeType, CDataType, AccDataType>(K, num_wgs_per_tile, maxv);
pass = ck_tile::check_err(c_host, ref, "streamk_registry", tol.rtol, tol.atol);
std::cout << "Verification: " << (pass ? "PASS" : "FAIL") << "\n";
}
return pass ? 0 : 1;
}

101
dispatcher/examples/gemm/cpp/streamk_driver_common.hpp Normal file
View File

@@ -0,0 +1,101 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
/**
* Shared helpers for the Stream-K GEMM example drivers (03 standalone and 04
* registry). Kept in one place so the two drivers do not duplicate CLI parsing,
* layout/dtype tags, and the Stream-K verification tolerance.
*/
#include <algorithm>
#include <string>
#include <type_traits>
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/ops/gemm.hpp"
#include "ck_tile/dispatcher/kernel_key.hpp"
template <typename Layout>
constexpr 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>>{};
}
inline 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;
}
// Map a ck_tile element type to the dispatcher's DataType enum so the registry
// key reflects the kernel that was actually generated (fp16/bf16/fp8/bf8/...),
// instead of assuming fp16. Keeps the registry identifier and selection correct
// across every datatype the codegen emits.
template <typename T>
constexpr ck_tile::dispatcher::DataType dtype_enum_of()
{
using U = ck_tile::remove_cvref_t<T>;
if constexpr(std::is_same_v<U, ck_tile::fp16_t>)
return ck_tile::dispatcher::DataType::FP16;
else if constexpr(std::is_same_v<U, ck_tile::bf16_t>)
return ck_tile::dispatcher::DataType::BF16;
else if constexpr(std::is_same_v<U, ck_tile::fp8_t>)
return ck_tile::dispatcher::DataType::FP8;
else if constexpr(std::is_same_v<U, ck_tile::bf8_t>)
return ck_tile::dispatcher::DataType::BF8;
else if constexpr(std::is_same_v<U, ck_tile::int8_t>)
return ck_tile::dispatcher::DataType::INT8;
else if constexpr(std::is_same_v<U, float>)
return ck_tile::dispatcher::DataType::FP32;
else
return ck_tile::dispatcher::DataType::UNKNOWN;
}
template <typename Layout>
constexpr ck_tile::dispatcher::LayoutTag layout_tag_of()
{
return std::is_same_v<ck_tile::remove_cvref_t<Layout>, ck_tile::tensor_layout::gemm::RowMajor>
? ck_tile::dispatcher::LayoutTag::RowMajor
: ck_tile::dispatcher::LayoutTag::ColMajor;
}
struct StreamKTolerance
{
double rtol;
double atol;
};
// Stream-K verification tolerance. Stream-K splits K across CUs and reduces the
// partials; atomic reduction accumulates them directly into low-precision C, so
// the tolerance must account for the split-K accumulation error -- exactly as
// tile_engine's calculate_rtol_atol does. The plain single-pass
// get_relative/absolute_threshold(K) under-estimates the error and would
// spuriously FAIL correct atomic results on small-M/N, large-K shapes.
//
// `num_wgs_per_tile` is the number of workgroups reducing into a single output
// tile (Stream-K has no fixed split-k), taken from the kernel's own tile
// partitioner so the driver and tile_engine agree on the split factor.
template <typename ComputeType, typename CDataType, typename AccDataType>
inline StreamKTolerance
streamk_tolerance(ck_tile::index_t K, ck_tile::index_t num_wgs_per_tile, float maxv)
{
const ck_tile::index_t k_per_split = ck_tile::integer_divide_ceil(K, num_wgs_per_tile);
// single-pass (per-split) tolerance
const double rtol_base =
ck_tile::get_relative_threshold<ComputeType, CDataType, AccDataType>(k_per_split);
const double atol_base = ck_tile::get_absolute_threshold<ComputeType, CDataType, AccDataType>(
maxv / num_wgs_per_tile, k_per_split);
// error contributed by reducing num_wgs_per_tile partials in low-precision C
const double rtol_split_k =
ck_tile::get_relative_threshold<CDataType, CDataType, CDataType>(num_wgs_per_tile);
const double atol_split_k =
ck_tile::get_absolute_threshold<CDataType, CDataType, CDataType>(maxv, num_wgs_per_tile);
return {std::max(rtol_base, rtol_split_k), std::max(atol_base, atol_split_k)};
}

257
dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp Normal file
View File

@@ -0,0 +1,257 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include "ck_tile/dispatcher/kernel_instance.hpp"
#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/ops/gemm.hpp"
#include "ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp"
#include "ck_tile/ops/common/streamk_common.hpp"
#include <hip/hip_runtime.h>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>
namespace ck_tile {
namespace dispatcher {
namespace backends {
// Lock the dispatcher's ReductionStrategy (defined in kernel_key.hpp, which is
// deliberately kept ck_tile-free -- same policy as the void* workspace in
// dispatcher.hpp) to ck_tile::StreamKReductionStrategy so the two enums cannot
// silently drift. The dispatcher enum carries an extra None=0 sentinel, so the
// three real strategies are offset by one. This backend header is the single
// place that includes both definitions, so the check belongs here rather than in
// the public key header.
static_assert(static_cast<std::uint32_t>(ReductionStrategy::Atomic) ==
static_cast<std::uint32_t>(ck_tile::StreamKReductionStrategy::Atomic) + 1u,
"dispatcher ReductionStrategy drifted from ck_tile::StreamKReductionStrategy");
static_assert(static_cast<std::uint32_t>(ReductionStrategy::Linear) ==
static_cast<std::uint32_t>(ck_tile::StreamKReductionStrategy::Linear) + 1u,
"dispatcher ReductionStrategy drifted from ck_tile::StreamKReductionStrategy");
static_assert(static_cast<std::uint32_t>(ReductionStrategy::Tree) ==
static_cast<std::uint32_t>(ck_tile::StreamKReductionStrategy::Tree) + 1u,
"dispatcher ReductionStrategy drifted from ck_tile::StreamKReductionStrategy");
/**
* Kernel-instance wrapper for unified_gemm_codegen.py Stream-K kernels.
*
* Counterpart of GeneratedTileKernelInstance (regular GEMM) for the Stream-K
* variant. The difference is the host-args type: Stream-K needs
* ck_tile::StreamKHostArgs (workspace pointer + reduction strategy), which is
* ABI-incompatible with the GemmHostArgs path -- this is exactly why Stream-K
* could not previously ride the registry. With this backend it can: the
* Dispatcher selects the instance by KernelKey (which now carries streamk +
* reduction_strategy) and calls run().
*
* supports() gates on the requested reduction strategy so that the registry can
* hold atomic/linear/tree side by side and the Dispatcher's first-fit selection
* picks the one the caller asked for via Problem::reduction_strategy.
*
* NOTE (PR-C): the generated SelectedKernel::launch(StreamKHostArgs, stream)
* still owns the reduction workspace internally (DeviceMem) and does the
* per-iter reset. PR-D relocates workspace ownership + reset to Dispatcher::run()
* via get_workspace_size()/the workspace-aware run() overload.
*/
template <typename SelectedKernelType,
typename ADataType_,
typename BDataType_,
typename CDataType_,
typename AccDataType_>
class GeneratedStreamKKernelInstance : public KernelInstance
{
public:
using ADataType = ADataType_;
using BDataType = BDataType_;
using CDataType = CDataType_;
using AccDataType = AccDataType_;
using SelectedKernel = SelectedKernelType;
GeneratedStreamKKernelInstance(const KernelKey& key, const std::string& name)
: key_(key), name_(name)
{
}
const KernelKey& get_key() const override { return key_; }
std::string get_name() const override { return name_; }
/// Accept ONLY when the caller requested a Stream-K kernel with THIS
/// instance's reduction strategy. Lets atomic/linear/tree coexist in the
/// registry and be selected by Problem::reduction_strategy.
bool supports(const Problem& problem) const override
{
if(!problem.streamk)
return false;
if(problem.reduction_strategy != key_.algorithm.reduction_strategy)
return false;
// Stream-K distributes K-iterations across workgroups; padding flags
// mirror the regular backend's divisibility guard.
constexpr bool pad_m = SelectedKernel::kPadM;
constexpr bool pad_n = SelectedKernel::kPadN;
constexpr bool pad_k = SelectedKernel::kPadK;
if(!pad_m && problem.M % SelectedKernel::TileM != 0)
return false;
if(!pad_n && problem.N % SelectedKernel::TileN != 0)
return false;
if(!pad_k && problem.K % SelectedKernel::TileK != 0)
return false;
// Final feasibility: enough tiles to partition across CUs. Rejecting here
// (instead of throwing at launch) lets the dispatcher's first-fit fall back
// to a non-Stream-K kernel for too-small problems.
return SelectedKernel::IsSupported(make_args(problem));
}
/// Device workspace (bytes) needed for `problem`. 0 for Atomic; >0 for
/// Linear/Tree. The Dispatcher uses this to size the buffer it owns and then
/// passes that buffer to the workspace-aware run() below.
std::size_t get_workspace_size(const Problem& problem) const override
{
return SelectedKernel::GetWorkSpaceSize(make_args(problem));
}
/// No-workspace entry point: delegates to the workspace-aware overload with a
/// null buffer, so the generated launch() falls back to its internal
/// (self-allocating) path. Used when the caller does not own a workspace.
float run(const void* a_ptr,
const void* b_ptr,
void* c_ptr,
const void** d_ptrs,
const Problem& problem,
void* stream) const override
{
return run(a_ptr, b_ptr, c_ptr, d_ptrs, /*workspace=*/nullptr, problem, stream);
}
/// Workspace-aware execution (PR-D). `workspace` is the Dispatcher-owned
/// reduction buffer (may be null for Atomic, which needs none). When non-null
/// the generated launch() binds it instead of allocating its own DeviceMem.
float run(const void* a_ptr,
const void* b_ptr,
void* c_ptr,
const void** d_ptrs,
void* workspace,
const Problem& problem,
void* stream) const override
{
(void)d_ptrs; // Not used for Stream-K GEMM
auto args = make_args(problem, a_ptr, b_ptr, c_ptr);
const bool bench = this->benchmarking_;
ck_tile::stream_config stream_cfg;
stream_cfg.stream_id_ = reinterpret_cast<hipStream_t>(stream);
stream_cfg.time_kernel_ = bench;
stream_cfg.log_level_ = 0;
stream_cfg.cold_niters_ = bench ? 5 : 0;
stream_cfg.nrepeat_ = bench ? 10 : 1;
stream_cfg.is_gpu_timer_ = bench;
// Flush the L2 between timed iterations so the measurement is cold, like
// tile_engine and the standalone 03 driver. Leaving the cache warm here was
// the methodology artifact that over-reported TFlops and produced the
// spurious dispatcher-vs-TE "performance gap"; do not present a warm number
// as parity evidence.
stream_cfg.flush_cache_ = bench;
// NOTE: input-buffer rotation is intentionally NOT enabled (rotating_count
// = 1). Atomic reduction accumulates straight into C, and this same run()
// serves the functional path that callers verify against the reference, so
// rotating/accumulating would corrupt the output left on the device. This
// means the timing here is cold-but-non-rotated and is therefore NOT the
// fully apple-to-apple surface: for TE-calibrated numbers use the 03 driver
// (or a --validate 0 pass) which rotates 1000 input copies like tile_engine.
stream_cfg.rotating_count_ = 1;
if(workspace != nullptr)
return SelectedKernel::launch(args, stream_cfg, workspace);
return SelectedKernel::launch(args, stream_cfg);
}
bool validate(const void* a_ptr,
const void* b_ptr,
const void* c_ptr,
const void** d_ptrs,
const Problem& problem,
float tolerance) const override
{
(void)d_ptrs;
(void)tolerance;
// This backend owns no host reference, so a numeric correctness check is
// out of scope here (the TE/driver harness does that). But returning a
// blind "true" would mis-report an unrunnable config as valid, so validate
// what we CAN without a reference: non-null operands, a well-formed
// problem, and that THIS Stream-K instance actually supports it.
if(a_ptr == nullptr || b_ptr == nullptr || c_ptr == nullptr)
return false;
if(!problem.is_valid())
return false;
return supports(problem);
}
private:
/// Build StreamKHostArgs for `problem`. Leading dims are derived from the
/// kernel key's layouts so every layout works (rcr/rrr/ccr/crr, ...), not
/// just rcr: A is MxK (row->K, col->M), B is KxN (row->N, col->K), C is MxN
/// (row->N, col->M). k_batch is owned by the Stream-K tile partitioner, not
/// passed here. Pointers default to null for sizing-only use
/// (GetWorkSpaceSize). StreamKHostArgs uses ck_tile::index_t (int32); cast
/// from Problem's int64.
ck_tile::StreamKHostArgs make_args(const Problem& problem,
const void* a_ptr = nullptr,
const void* b_ptr = nullptr,
void* c_ptr = nullptr) const
{
using idx = ck_tile::index_t;
// StreamKHostArgs uses int32 index_t while Problem carries int64 dims.
// Guard the narrowing so an oversized M/N/K (or a derived leading dim)
// fails loudly instead of silently wrapping to a negative/garbage extent.
// The dimension parser was widened to std::stoll specifically to avoid
// overflow, so dropping back to int32 here must be checked, not assumed.
auto to_idx = [](std::int64_t v, const char* what) -> idx {
if(v < 0 || v > static_cast<std::int64_t>(std::numeric_limits<idx>::max()))
throw std::runtime_error(std::string("StreamK make_args: ") + what + " (" +
std::to_string(v) +
") exceeds int32 ck_tile::index_t range");
return static_cast<idx>(v);
};
const auto& sig = key_.signature;
const bool a_row = sig.layout_a == LayoutTag::RowMajor;
const bool b_row = sig.layout_b == LayoutTag::RowMajor;
const bool c_row = sig.layout_c == LayoutTag::RowMajor;
const idx M = to_idx(problem.M, "M");
const idx N = to_idx(problem.N, "N");
const idx K = to_idx(problem.K, "K");
const idx stride_a = to_idx(a_row ? problem.K : problem.M, "stride_a");
const idx stride_b = to_idx(b_row ? problem.N : problem.K, "stride_b");
const idx stride_c = to_idx(c_row ? problem.N : problem.M, "stride_c");
return ck_tile::StreamKHostArgs{a_ptr, b_ptr, c_ptr, M, N, K, stride_a, stride_b, stride_c};
}
KernelKey key_;
std::string name_;
};
/// Helper to create a Stream-K kernel-instance wrapper.
template <typename SelectedKernel,
typename ADataType,
typename BDataType,
typename CDataType,
typename AccDataType>
std::shared_ptr<KernelInstance> create_generated_streamk_kernel(const KernelKey& key,
const std::string& name)
{
return std::make_shared<GeneratedStreamKKernelInstance<SelectedKernel,
ADataType,
BDataType,
CDataType,
AccDataType>>(key, name);
}
} // namespace backends
} // namespace dispatcher
} // namespace ck_tile

36
dispatcher/include/ck_tile/dispatcher/dispatcher.hpp
View File

@@ -27,6 +27,7 @@
#include "ck_tile/dispatcher/kernel_instance.hpp"
#include "ck_tile/dispatcher/problem.hpp"
#include "ck_tile/dispatcher/registry.hpp"
#include <cstddef>
#include <functional>
#include <memory>
#include <string>
@@ -41,6 +42,16 @@ using HeuristicFunction = std::function<std::vector<std::string>(const Problem&)
/// Dispatcher: Top-level orchestration for kernel selection and execution
/// Provides unified interface for kernel dispatch across different backends
///
/// Concurrency contract: a Dispatcher instance is NOT safe for concurrent use
/// from multiple threads / HIP streams. It owns a single reduction workspace for
/// Stream-K linear/tree kernels (see workspace_ below), which would be corrupted
/// by two overlapping dispatches. Callers that need concurrency should create one
/// Dispatcher per stream/thread (the object is a lightweight handle -- just a
/// Registry* + arch string + heuristic), exactly as one would use per-stream
/// library handles. This mirrors how the workspace is zeroed on the caller's
/// stream in run() (hipMemsetAsync), so a per-stream Dispatcher stays correctly
/// ordered without any cross-stream synchronization.
class Dispatcher
{
public:
@@ -56,6 +67,18 @@ class Dispatcher
/// @param gfx_arch Target GPU architecture (e.g. "gfx950")
explicit Dispatcher(Registry* registry = nullptr, const std::string& gfx_arch = "");
/// Frees the dispatcher-owned Stream-K reduction workspace, if any.
~Dispatcher();
/// The Dispatcher owns a raw HIP reduction workspace that it frees in the
/// destructor, so it must not be copied (a copy would double-free the buffer)
/// nor moved (no use-case, and consistent with the single-stream contract
/// above). Non-copyable, non-movable.
Dispatcher(const Dispatcher&) = delete;
Dispatcher& operator=(const Dispatcher&) = delete;
Dispatcher(Dispatcher&&) = delete;
Dispatcher& operator=(Dispatcher&&) = delete;
void set_arch(const std::string& arch) { gfx_arch_ = arch; }
[[nodiscard]] const std::string& arch() const { return gfx_arch_; }
@@ -149,6 +172,19 @@ class Dispatcher
std::string gfx_arch_;
bool benchmarking_ = true;
// Dispatcher-owned, grow-on-demand reduction workspace for Stream-K kernels
// (linear/tree). Sized via KernelInstance::get_workspace_size() and reused
// across calls so we don't hipMalloc/hipFree on the hot path. Held as a raw
// pointer to keep HIP/ck_tile out of this public header.
mutable void* workspace_ = nullptr;
mutable std::size_t workspace_bytes_ = 0;
/// Ensure the owned workspace holds at least `bytes`, growing it if needed,
/// and zero the first `bytes` on `stream` (hipMemsetAsync). Not thread-safe --
/// see the Dispatcher concurrency contract above (one Dispatcher per stream).
/// `stream` is a hipStream_t held as void* to keep HIP out of this header.
void ensure_workspace(std::size_t bytes, void* stream) const;
/// Select kernel using first-fit strategy
[[nodiscard]] KernelInstancePtr select_first_fit(const Problem& problem) const;

25
dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp
View File

@@ -5,6 +5,7 @@
#include "ck_tile/dispatcher/kernel_key.hpp"
#include "ck_tile/dispatcher/problem.hpp"
#include <cstddef>
#include <memory>
#include <string>
@@ -45,6 +46,30 @@ class KernelInstance
const Problem& problem,
void* stream = nullptr) const = 0;
/// Device workspace (in bytes) this kernel needs for `problem` (0 = none).
/// Non-zero only for Stream-K linear/tree reductions; the caller (Dispatcher)
/// sizes and owns the buffer and passes it to the workspace-aware run().
[[nodiscard]] virtual std::size_t get_workspace_size(const Problem& problem) const
{
(void)problem;
return 0;
}
/// Workspace-aware execution. Default forwards to the no-workspace run(), so
/// existing (non-Stream-K) kernels need no change; the Stream-K backend
/// overrides this to set the reduction workspace pointer before launch.
[[nodiscard]] virtual float run(const void* a_ptr,
const void* b_ptr,
void* c_ptr,
const void** d_ptrs,
void* workspace,
const Problem& problem,
void* stream = nullptr) const
{
(void)workspace;
return run(a_ptr, b_ptr, c_ptr, d_ptrs, problem, stream);
}
/// Validate kernel output against reference implementation
/// @param a_ptr Pointer to matrix A (device memory)
/// @param b_ptr Pointer to matrix B (device memory)

46
dispatcher/include/ck_tile/dispatcher/kernel_key.hpp
View File

@@ -72,6 +72,30 @@ enum class Scheduler : std::uint8_t
Interwave
};
/// Stream-K partial-sum reduction strategy. `None` = not a Stream-K kernel.
/// Mirrors ck_tile::StreamKReductionStrategy (Atomic/Linear/Tree).
enum class ReductionStrategy : std::uint8_t
{
None = 0,
Atomic,
Linear,
Tree
};
/// Canonical lower-case name for a reduction strategy. Matches the codegen suffix
/// scheme (atomic -> "atomic", etc.) so callers/drivers share one spelling.
inline const char* to_string(ReductionStrategy r)
{
switch(r)
{
case ReductionStrategy::Atomic: return "atomic";
case ReductionStrategy::Linear: return "linear";
case ReductionStrategy::Tree: return "tree";
case ReductionStrategy::None: return "none";
}
return "none";
}
/// KernelKey: Compile-time kernel configuration metadata
/// Organized into Signature (what operation) and Algorithm (how it's implemented)
struct KernelKey
@@ -147,6 +171,11 @@ struct KernelKey
bool pad_m = true; // Support arbitrary M dimensions via padding
bool pad_n = true; // Support arbitrary N dimensions via padding
bool pad_k = true; // Support arbitrary K dimensions via padding
// Stream-K (workgroup K-stream) parameters
bool streamk = false; // is a Stream-K kernel
ReductionStrategy reduction_strategy = ReductionStrategy::None; // atomic / linear / tree
bool workspace = false; // needs a device accumulation buffer (linear/tree)
} algorithm;
std::string gfx_arch; // e.g. "gfx942", "gfx90a", "gfx908"
@@ -195,7 +224,10 @@ struct KernelKey
algorithm.num_wave_groups,
algorithm.pad_m,
algorithm.pad_n,
algorithm.pad_k);
algorithm.pad_k,
algorithm.streamk,
algorithm.reduction_strategy,
algorithm.workspace);
}
/// Equality comparison
@@ -445,6 +477,18 @@ inline std::string KernelKey::encode_identifier() const
if(algorithm.preshuffle)
oss << "_preshuffle";
// Stream-K suffix -- must match unified_gemm_codegen.py KernelNaming.generate():
// atomic -> "..._streamk" linear -> "..._streamk_linear" tree -> "..._streamk_tree"
// Guarded by algorithm.streamk so non-Stream-K identifiers stay byte-identical.
if(algorithm.streamk)
{
oss << "_streamk";
if(algorithm.reduction_strategy == ReductionStrategy::Linear)
oss << "_linear";
else if(algorithm.reduction_strategy == ReductionStrategy::Tree)
oss << "_tree";
}
return oss.str();
}

22
dispatcher/include/ck_tile/dispatcher/problem.hpp
View File

@@ -7,6 +7,8 @@
#include <stdexcept>
#include <string>
#include "ck_tile/dispatcher/kernel_key.hpp" // ReductionStrategy
namespace ck_tile {
namespace dispatcher {
@@ -58,6 +60,10 @@ struct Problem
// Validation control
bool enable_validation; // Enable output validation against reference
// Stream-K request: which reduction strategy the caller wants (None = non-Stream-K)
bool streamk = false;
ReductionStrategy reduction_strategy = ReductionStrategy::None;
/// Default constructor with sensible defaults
Problem()
: M(0),
@@ -66,7 +72,9 @@ struct Problem
k_batch(1),
smem_budget(0),
prefer_persistent(false),
enable_validation(false)
enable_validation(false),
streamk(false),
reduction_strategy(ReductionStrategy::None)
{
}
@@ -78,7 +86,9 @@ struct Problem
k_batch(1),
smem_budget(0),
prefer_persistent(false),
enable_validation(false)
enable_validation(false),
streamk(false),
reduction_strategy(ReductionStrategy::None)
{
}
@@ -293,6 +303,14 @@ class ProblemBuilder
return *this;
}
/// Request a Stream-K kernel with a given reduction strategy
ProblemBuilder& stream_k(ReductionStrategy strategy = ReductionStrategy::Atomic)
{
problem_.streamk = true;
problem_.reduction_strategy = strategy;
return *this;
}
/// Build the Problem
[[nodiscard]] Problem build() const
{

81
dispatcher/src/dispatcher.cpp
View File

@@ -3,6 +3,7 @@
#include "ck_tile/dispatcher/dispatcher.hpp"
#include "ck_tile/dispatcher/dispatcher_error.hpp"
#include <hip/hip_runtime.h>
#include <sstream>
#include <iostream>
@@ -17,6 +18,54 @@ Dispatcher::Dispatcher(Registry* registry, const std::string& gfx_arch)
{
}
Dispatcher::~Dispatcher()
{
if(workspace_)
{
(void)hipFree(workspace_);
workspace_ = nullptr;
workspace_bytes_ = 0;
}
}
void Dispatcher::ensure_workspace(std::size_t bytes, void* stream) const
{
// Not thread-safe: mutates the Dispatcher-owned buffer. Safe because a
// Dispatcher is used from a single stream/thread (see the concurrency
// contract in dispatcher.hpp) -- there is no shared-buffer contention to
// guard against, so no lock is needed.
if(bytes > workspace_bytes_)
{
if(workspace_)
{
(void)hipFree(workspace_);
workspace_ = nullptr;
workspace_bytes_ = 0;
}
if(hipMalloc(&workspace_, bytes) != hipSuccess)
{
workspace_ = nullptr;
workspace_bytes_ = 0;
throw DispatcherError("Dispatcher: failed to allocate Stream-K reduction workspace");
}
workspace_bytes_ = bytes;
}
// Zero the region the kernel will use. Linear/Tree reductions accumulate into
// this buffer and read it before writing, so a stale/garbage buffer corrupts
// results. Doing it here makes correctness independent of whether the backend's
// per-iteration preprocess reset runs (e.g. on the non-benchmarking nrepeat=1
// path), mirroring the internal DeviceMem::SetZero() the standalone launch does.
// Zeroed on the caller's stream so the reset is ordered against the kernel
// launch that follows (same stream) without an implicit device-wide sync.
if(bytes > 0 &&
hipMemsetAsync(workspace_, 0, bytes, static_cast<hipStream_t>(stream)) != hipSuccess)
{
throw DispatcherError("Dispatcher: failed to zero Stream-K reduction workspace");
}
}
void Dispatcher::set_heuristic(HeuristicFunction heuristic)
{
heuristic_ = heuristic;
@@ -66,7 +115,21 @@ float Dispatcher::run_fused(const void* a_ptr,
}
kernel->set_benchmarking(benchmarking_);
return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, problem, stream);
// Size and own the reduction workspace (0 for non-Stream-K and for Atomic).
// For Linear/Tree the Dispatcher owns and reuses the buffer; no lock is taken
// because a Dispatcher is single-stream (see the concurrency contract in
// dispatcher.hpp). The buffer is zeroed on the caller's stream and the kernel
// launches on the same stream, so the reset is correctly ordered.
const std::size_t ws_bytes = kernel->get_workspace_size(problem);
if(ws_bytes > 0)
{
ensure_workspace(ws_bytes, stream); // grows if needed AND zeroes ws_bytes on `stream`
return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, workspace_, problem, stream);
}
// No workspace needed (non-Stream-K / Atomic): nothing to size or zero.
return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, nullptr, problem, stream);
}
float Dispatcher::run_explicit(const std::string& kernel_id,
@@ -92,7 +155,21 @@ float Dispatcher::run_explicit(const std::string& kernel_id,
}
kernel->set_benchmarking(benchmarking_);
return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, problem, stream);
// Size and own the reduction workspace (0 for non-Stream-K and for Atomic).
// For Linear/Tree the Dispatcher owns and reuses the buffer; no lock is taken
// because a Dispatcher is single-stream (see the concurrency contract in
// dispatcher.hpp). The buffer is zeroed on the caller's stream and the kernel
// launches on the same stream, so the reset is correctly ordered.
const std::size_t ws_bytes = kernel->get_workspace_size(problem);
if(ws_bytes > 0)
{
ensure_workspace(ws_bytes, stream); // grows if needed AND zeroes ws_bytes on `stream`
return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, workspace_, problem, stream);
}
// No workspace needed (non-Stream-K / Atomic): nothing to size or zero.
return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, nullptr, problem, stream);
}
bool Dispatcher::validate(const void* a_ptr,

30
dispatcher/tests/CMakeLists.txt
View File

@@ -126,6 +126,36 @@ set_tests_properties(dispatcher_test_fmha_parity PROPERTIES
ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts"
)
# Stream-K deep-core registry test (requires GPU + hipcc; SKIPs otherwise).
# Pass the gfx target CMake already configured with so the test does not have to
# detect it at runtime (no rocminfo dependency); it falls back to ROCm env vars
# / amdgpu-arch if none is set here.
set(_streamk_test_arch "")
if(GPU_TARGETS)
list(GET GPU_TARGETS 0 _streamk_test_arch)
elseif(CMAKE_HIP_ARCHITECTURES)
list(GET CMAKE_HIP_ARCHITECTURES 0 _streamk_test_arch)
elseif(AMDGPU_TARGETS)
list(GET AMDGPU_TARGETS 0 _streamk_test_arch)
endif()
set(_streamk_arch_arg "")
if(_streamk_test_arch)
set(_streamk_arch_arg --arch ${_streamk_test_arch})
endif()
add_test(
NAME dispatcher_test_streamk_registry
COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_streamk_registry.py ${_streamk_arch_arg}
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/..
)
set_tests_properties(dispatcher_test_streamk_registry PROPERTIES
LABELS "dispatcher;python;streamk;gpu"
TIMEOUT 900
SKIP_RETURN_CODE 77
ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts"
)
# Stress Test Script
add_test(
NAME dispatcher_stress_test

266
dispatcher/tests/test_streamk_registry.py Normal file
View File

@@ -0,0 +1,266 @@
#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""
Stream-K deep-core registry test (requires a GPU + hipcc).
Guards the deep-core path that lets Stream-K ride the registry like regular GEMM:
codegen -> generated SK wrapper -> Registry -> Dispatcher::run() (workspace alloc
+ strategy-aware reset) -> generated_tile_backend_streamk -> verify vs reference.
Each reduction strategy (atomic/linear/tree) is a *distinct compiled kernel*
(SkReductionStrategy is a compile-time constexpr), so we generate all three from a
single tile config and build the 04 registry driver once per strategy, force-
including that strategy's header. For each we assert:
* the encode_identifier() suffix matches the strategy (..._streamk[_linear|_tree])
* the Dispatcher selects that kernel by Problem::reduction_strategy
* the result verifies against the reference GEMM
The test SKIPs (exit 77) when no GPU or no hipcc is available, so it is safe in
CPU-only CI; it only runs the heavy build+launch where a GPU is present.
Usage:
python3 test_streamk_registry.py
python3 test_streamk_registry.py --arch gfx942 --m 3840 --n 4096 --k 2048
"""
import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
DISPATCHER_DIR = Path(__file__).resolve().parent.parent
CK_DIR = DISPATCHER_DIR.parent
CODEGEN = DISPATCHER_DIR / "codegen" / "unified_gemm_codegen.py"
DRIVER = DISPATCHER_DIR / "examples" / "gemm" / "cpp" / "04_streamk_registry_driver.cpp"
REGISTRY_SRC = DISPATCHER_DIR / "src" / "registry.cpp"
DISPATCHER_SRC = DISPATCHER_DIR / "src" / "dispatcher.cpp"
SKIP = 77 # ctest SKIP_RETURN_CODE
# One tile config, all three reduction strategies.
TILE = "128x128x64_2x2x1_32x32x16"
TILE_CONFIG_JSON = json.dumps(
{
"tile_config": {
"tile_m": [128], "tile_n": [128], "tile_k": [64],
"warp_m": [2], "warp_n": [2], "warp_k": [1],
"warp_tile_m": [32], "warp_tile_n": [32], "warp_tile_k": [16],
"block_size": [256],
},
"trait_config": {
"pipeline": ["compv3"], "epilogue": ["cshuffle"], "scheduler": ["intrawave"],
"pad_m": [False], "pad_n": [False], "pad_k": [False], "persistent": [False],
},
"streamk_config": {"reduction_strategy": ["atomic", "linear", "tree"]},
}
)
# strategy -> (header variant suffix, expected encode_identifier suffix)
STRATEGIES = {
"atomic": ("streamk", "_streamk"),
"linear": ("streamk_linear", "_streamk_linear"),
"tree": ("streamk_tree", "_streamk_tree"),
}
# Datatypes the Stream-K dispatcher codegen supports end-to-end. fp8/bf8 inputs
# accumulate in fp32 and write an fp16 C tensor (get_output_dtype), matching
# Tile Engine; the registry identifier keys on the input dtype (dtype_a), so the
# expected encode_identifier prefix is "{dtype}_{layout}" for each.
DATATYPES = ["fp16", "bf16", "fp8", "bf8"]
# Layouts Tile Engine builds Stream-K for (all keep C row-major, which the atomic
# C-reset relies on). Full coverage = DATATYPES x LAYOUTS x STRATEGIES.
LAYOUTS = ["rcr", "rrr", "ccr", "crr"]
def detect_arch(fallback=None):
# Resolve the gfx target without shelling out to rocminfo. Preference order:
# the arch the build already configured with (passed via --arch from
# CMakeLists.txt) is handled by the caller; here we fall back to the standard
# ROCm environment variables and then the amdgpu-arch / offload-arch LLVM
# tools, which query the driver directly and ship with the ROCm/LLVM toolchain.
for env in ("PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET", "AMDGPU_TARGETS", "GPU_TARGETS"):
val = os.environ.get(env)
if val:
return re.split(r"[;,]", val)[0].strip()
for tool in ("amdgpu-arch", "offload-arch"):
exe = shutil.which(tool)
if exe:
try:
out = run([exe], timeout=30).stdout
m = re.search(r"\bgfx[0-9a-f]+\b", out)
if m:
return m.group(0)
except Exception:
pass
return fallback
def run(cmd, **kw):
return subprocess.run(cmd, capture_output=True, text=True, **kw)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--arch", default=None)
ap.add_argument("--m", type=int, default=3840)
ap.add_argument("--n", type=int, default=4096)
ap.add_argument("--k", type=int, default=2048)
ap.add_argument(
"--datatypes", default=",".join(DATATYPES),
help="Comma-separated datatypes to test (default: all TE-equivalent).",
)
ap.add_argument(
"--layouts", default=",".join(LAYOUTS),
help="Comma-separated layouts to test (default: all TE-equivalent).",
)
args = ap.parse_args()
datatypes = [d.strip() for d in args.datatypes.split(",") if d.strip()]
layouts = [l.strip() for l in args.layouts.split(",") if l.strip()]
hipcc = shutil.which("hipcc")
if not hipcc:
print("SKIP: hipcc not found")
return SKIP
arch = args.arch or detect_arch()
if not arch:
print("SKIP: no GPU / could not detect gfx arch")
return SKIP
print(f"Stream-K registry test on {arch} @ {args.m}x{args.n}x{args.k}")
inc = ["-I", str(CK_DIR / "include"), "-I", str(DISPATCHER_DIR / "include")]
with tempfile.TemporaryDirectory(prefix="sk_reg_test_") as td:
# Build the dtype-independent core objects once (no force-include).
reg_o, disp_o = Path(td) / "registry.o", Path(td) / "dispatcher.o"
for src, obj in ((REGISTRY_SRC, reg_o), (DISPATCHER_SRC, disp_o)):
c = run(
[hipcc, "-std=c++17", f"--offload-arch={arch}", "-O3", *inc,
"-c", str(src), "-o", str(obj)],
timeout=900,
)
if c.returncode != 0:
print(f"FAIL: compiling {src.name}\n" + c.stderr[-2000:])
return 1
failures = []
for dtype in datatypes:
for layout in layouts:
failures += run_for_combo(
dtype, layout, td, arch, args, hipcc, inc, reg_o, disp_o
)
if failures:
print("\nSTREAM-K REGISTRY TEST FAILED:")
for f in failures:
print(" - " + f)
return 1
print(
"All Stream-K combos registered, dispatched, and verified "
f"(datatypes: {', '.join(datatypes)} | layouts: {', '.join(layouts)})."
)
return 0
def run_for_combo(dtype, layout, td, arch, args, hipcc, inc, reg_o, disp_o):
"""Generate + build + run all reduction strategies for one (dtype, layout).
Returns a list of failure strings (empty on success)."""
failures = []
# Verify each built kernel against the CLI shape AND a small-M/N, large-K
# shape. The latter maximizes the Stream-K split factor, which is exactly
# where the split-K-aware verification tolerance matters: a plain single-pass
# tolerance spuriously FAILs correct atomic results on this shape. The driver
# binary is shape-independent, so this only adds runs, not rebuilds.
shapes = [(args.m, args.n, args.k), (128, 128, 16384)]
gen = Path(td) / f"gen_{dtype}_{layout}"
# 1) generate all three strategy headers from one tile config
g = run(
[
sys.executable, str(CODEGEN),
"--datatype", dtype, "--layout", layout,
"--gpu-target", arch, "--variants", "stream_k",
"--tile-config-json", TILE_CONFIG_JSON,
"--output-dir", str(gen),
],
timeout=600,
)
if g.returncode != 0:
return [f"{dtype}/{layout}: codegen failed\n" + g.stderr[-2000:]]
for strat, (variant, want_suffix) in STRATEGIES.items():
tag = f"{dtype}/{layout}/{strat}"
header = gen / (
f"gemm_{dtype}_{layout}_compv3_cshuffle_intrawave_"
f"False_False_False_False_{TILE}_{variant}.hpp"
)
if not header.exists():
failures.append(f"{tag}: generated header missing ({header.name})")
continue
stem = f"{dtype}_{layout}_{variant}"
drv_o, exe = Path(td) / f"d_{stem}.o", Path(td) / f"skreg_{stem}"
c = run(
[hipcc, "-std=c++17", f"--offload-arch={arch}", "-O3",
"-DCK_TILE_SINGLE_KERNEL_INCLUDE", f'-DGFX_ARCH="{arch}"',
*inc, "-I", str(gen), "-include", str(header),
"-c", str(DRIVER), "-o", str(drv_o)],
timeout=900,
)
if c.returncode != 0:
failures.append(f"{tag}: driver compile failed\n{c.stderr[-1500:]}")
continue
l = run(
[hipcc, f"--offload-arch={arch}", str(drv_o), str(disp_o),
str(reg_o), "-o", str(exe)],
timeout=300,
)
if l.returncode != 0:
failures.append(f"{tag}: link failed\n{l.stderr[-1500:]}")
continue
for (sm, sn, sk) in shapes:
r = run(
[str(exe), "--m", str(sm), "--n", str(sn),
"--k", str(sk), "--strategy", strat, "--validate", "1"],
timeout=300,
)
out = r.stdout
ok_verify = "Verification: PASS" in out
# Guard the identifier parse: a crashed/silent driver prints no
# "identifier=" token, so split(...)[1] would raise IndexError and
# abort the run instead of recording a clean failure.
ok_suffix = False
if f"identifier={dtype}_{layout}" in out and "identifier=" in out:
token = out.split("identifier=", 1)[1].split()[0]
ok_suffix = want_suffix in token
if r.returncode != 0 or not ok_verify or not ok_suffix:
failures.append(
f"{tag} @ {sm}x{sn}x{sk}: rc={r.returncode} verify={ok_verify} "
f"suffix_ok={ok_suffix}\n{out[-800:]}{r.stderr[-400:]}"
)
else:
tflops = next(
(ln for ln in out.splitlines() if "TFlops" in ln), ""
).strip()
print(
f" PASS {dtype:5s} {layout:4s} {strat:6s} {sm}x{sn}x{sk} "
f"-> {want_suffix} | {tflops}"
)
return failures
if __name__ == "__main__":
sys.exit(main())