diff --git a/dispatcher/README.md b/dispatcher/README.md index 307e612305..19b97a91bb 100644 --- a/dispatcher/README.md +++ b/dispatcher/README.md @@ -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). --- diff --git a/dispatcher/STREAMK.md b/dispatcher/STREAMK.md new file mode 100644 index 0000000000..4832bca296 --- /dev/null +++ b/dispatcher/STREAMK.md @@ -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= --overlap \ + --container-image=/cluster/images/ck/ck_rocm7.1.1_therock_.sqsh \ + --container-mounts=$HOME:$HOME \ + bash -lc '' +``` + +--- + +> 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___compv3_cshuffle_intrawave______.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` | diff --git a/dispatcher/codegen/arch_filter.py b/dispatcher/codegen/arch_filter.py index 63dbee2dd7..6a2dff0b39 100644 --- a/dispatcher/codegen/arch_filter.py +++ b/dispatcher/codegen/arch_filter.py @@ -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 diff --git a/dispatcher/codegen/unified_gemm_codegen.py b/dispatcher/codegen/unified_gemm_codegen.py index 47cb3680f1..3b26e5a112 100755 --- a/dispatcher/codegen/unified_gemm_codegen.py +++ b/dispatcher/codegen/unified_gemm_codegen.py @@ -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 +#include +#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::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; + using SkUniversalGemmProblem = UniversalGemmPipelineProblem< + ADataType, BDataType, AccDataType, TileShape, SkGemmUniversalTraits, SkScheduler>; + using SkGemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}; + {self._epilogue_code(config)} + using SkStreamKTilePartitioner = + ck_tile::StreamKTilePartitioner; + using StreamKGemmKernel = + ck_tile::StreamKKernel; + + // 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 preprocess = reset_data_buffers; + + float ave_time = launch_kernel_time_mask(stream, preprocess, + make_kernel(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 preprocess = reset_data_buffers; + + float ave_time = launch_kernel_time_mask(stream, preprocess, + make_kernel(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" + f'(key, "{kernel_name}");' + ) + else: + ret_stmt = ( + "return std::make_shared>" + 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>(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", ) diff --git a/dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp b/dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp new file mode 100644 index 0000000000..f8c032ca4d --- /dev/null +++ b/dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp @@ -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 /include -I \ + * -include /_streamk.hpp \ + * 03_streamk_gemm_driver.cpp -o streamk_gemm_driver + */ + +#include + +#include +#include +#include +#include +#include + +#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 a_host( + ck_tile::host_tensor_descriptor(M, K, sA, is_row_major(ALayout{}))); + ck_tile::HostTensor b_host( + ck_tile::host_tensor_descriptor(K, N, sB, is_row_major(BLayout{}))); + ck_tile::HostTensor c_host( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_host); + ck_tile::FillUniformDistribution{-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(flop) / 1.E9 / ave_time; + const float gbps = static_cast(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 ref( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + ref.SetZero(); + ck_tile::reference_gemm(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; + auto kargs = SelectedKernel::StreamKGemmKernel::MakeKernelArgs(args); + const ck_tile::index_t num_wgs_per_tile = + std::max(1, kargs.tile_partitioner.estimate_num_wgs_per_tile()); + const auto tol = + streamk_tolerance(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; +} diff --git a/dispatcher/examples/gemm/cpp/04_streamk_registry_driver.cpp b/dispatcher/examples/gemm/cpp/04_streamk_registry_driver.cpp new file mode 100644 index 0000000000..5cd56d46bc --- /dev/null +++ b/dispatcher/examples/gemm/cpp/04_streamk_registry_driver.cpp @@ -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 /include -I /dispatcher/include -I \ + * -include /_streamk.hpp \ + * 04_streamk_registry_driver.cpp -o streamk_registry_driver + */ + +#include + +#include +#include +#include +#include +#include + +#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(); + key.signature.dtype_b = dtype_enum_of(); + key.signature.dtype_c = dtype_enum_of(); + key.signature.dtype_acc = dtype_enum_of(); + key.signature.layout_a = layout_tag_of(); + key.signature.layout_b = layout_tag_of(); + key.signature.layout_c = layout_tag_of(); + 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(SelectedKernel::WarpTileM), + static_cast(SelectedKernel::WarpTileN), + static_cast(SelectedKernel::WarpTileK)}; + key.algorithm.wave_shape = {static_cast(SelectedKernel::WarpPerBlock_M), + static_cast(SelectedKernel::WarpPerBlock_N), + static_cast(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(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 a_host( + ck_tile::host_tensor_descriptor(M, K, sA, is_row_major(ALayout{}))); + ck_tile::HostTensor b_host( + ck_tile::host_tensor_descriptor(K, N, sB, is_row_major(BLayout{}))); + ck_tile::HostTensor c_host( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_host); + ck_tile::FillUniformDistribution{-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(flop) / 1.E9 / ave_time; + const float gbps = static_cast(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 ref( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + ref.SetZero(); + ck_tile::reference_gemm(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; + auto kargs = SelectedKernel::StreamKGemmKernel::MakeKernelArgs(sk_args); + const ck_tile::index_t num_wgs_per_tile = + std::max(1, kargs.tile_partitioner.estimate_num_wgs_per_tile()); + const auto tol = + streamk_tolerance(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; +} diff --git a/dispatcher/examples/gemm/cpp/streamk_driver_common.hpp b/dispatcher/examples/gemm/cpp/streamk_driver_common.hpp new file mode 100644 index 0000000000..e30784619e --- /dev/null +++ b/dispatcher/examples/gemm/cpp/streamk_driver_common.hpp @@ -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 +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" + +#include "ck_tile/dispatcher/kernel_key.hpp" + +template +constexpr auto is_row_major(Layout) +{ + return ck_tile::bool_constant< + std::is_same_v, 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 +constexpr ck_tile::dispatcher::DataType dtype_enum_of() +{ + using U = ck_tile::remove_cvref_t; + if constexpr(std::is_same_v) + return ck_tile::dispatcher::DataType::FP16; + else if constexpr(std::is_same_v) + return ck_tile::dispatcher::DataType::BF16; + else if constexpr(std::is_same_v) + return ck_tile::dispatcher::DataType::FP8; + else if constexpr(std::is_same_v) + return ck_tile::dispatcher::DataType::BF8; + else if constexpr(std::is_same_v) + return ck_tile::dispatcher::DataType::INT8; + else if constexpr(std::is_same_v) + return ck_tile::dispatcher::DataType::FP32; + else + return ck_tile::dispatcher::DataType::UNKNOWN; +} + +template +constexpr ck_tile::dispatcher::LayoutTag layout_tag_of() +{ + return std::is_same_v, 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 +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(k_per_split); + const double atol_base = ck_tile::get_absolute_threshold( + 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(num_wgs_per_tile); + const double atol_split_k = + ck_tile::get_absolute_threshold(maxv, num_wgs_per_tile); + return {std::max(rtol_base, rtol_split_k), std::max(atol_base, atol_split_k)}; +} diff --git a/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp b/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp new file mode 100644 index 0000000000..7badf44900 --- /dev/null +++ b/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp @@ -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 +#include +#include +#include +#include + +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(ReductionStrategy::Atomic) == + static_cast(ck_tile::StreamKReductionStrategy::Atomic) + 1u, + "dispatcher ReductionStrategy drifted from ck_tile::StreamKReductionStrategy"); +static_assert(static_cast(ReductionStrategy::Linear) == + static_cast(ck_tile::StreamKReductionStrategy::Linear) + 1u, + "dispatcher ReductionStrategy drifted from ck_tile::StreamKReductionStrategy"); +static_assert(static_cast(ReductionStrategy::Tree) == + static_cast(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 +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(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::numeric_limits::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(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 +std::shared_ptr create_generated_streamk_kernel(const KernelKey& key, + const std::string& name) +{ + return std::make_shared>(key, name); +} + +} // namespace backends +} // namespace dispatcher +} // namespace ck_tile diff --git a/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp b/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp index d266d693da..80a90ea064 100644 --- a/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp +++ b/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp @@ -27,6 +27,7 @@ #include "ck_tile/dispatcher/kernel_instance.hpp" #include "ck_tile/dispatcher/problem.hpp" #include "ck_tile/dispatcher/registry.hpp" +#include #include #include #include @@ -41,6 +42,16 @@ using HeuristicFunction = std::function(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; diff --git a/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp b/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp index b6ef76e4f8..9b577ece35 100644 --- a/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp +++ b/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp @@ -5,6 +5,7 @@ #include "ck_tile/dispatcher/kernel_key.hpp" #include "ck_tile/dispatcher/problem.hpp" +#include #include #include @@ -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) diff --git a/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp b/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp index af903534ec..637eea7441 100644 --- a/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp +++ b/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp @@ -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(); } diff --git a/dispatcher/include/ck_tile/dispatcher/problem.hpp b/dispatcher/include/ck_tile/dispatcher/problem.hpp index 5bffb56b49..266a568328 100644 --- a/dispatcher/include/ck_tile/dispatcher/problem.hpp +++ b/dispatcher/include/ck_tile/dispatcher/problem.hpp @@ -7,6 +7,8 @@ #include #include +#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 { diff --git a/dispatcher/src/dispatcher.cpp b/dispatcher/src/dispatcher.cpp index 133485b248..3aad515c3d 100644 --- a/dispatcher/src/dispatcher.cpp +++ b/dispatcher/src/dispatcher.cpp @@ -3,6 +3,7 @@ #include "ck_tile/dispatcher/dispatcher.hpp" #include "ck_tile/dispatcher/dispatcher_error.hpp" +#include #include #include @@ -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(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, diff --git a/dispatcher/tests/CMakeLists.txt b/dispatcher/tests/CMakeLists.txt index a18663f76d..3976988ab0 100644 --- a/dispatcher/tests/CMakeLists.txt +++ b/dispatcher/tests/CMakeLists.txt @@ -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 diff --git a/dispatcher/tests/test_streamk_registry.py b/dispatcher/tests/test_streamk_registry.py new file mode 100644 index 0000000000..5e6c92b9d4 --- /dev/null +++ b/dispatcher/tests/test_streamk_registry.py @@ -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())