Files
composable_kernel/dispatcher/examples/gemm/cpp
Muhammed Emin Ozturk 5d3380aa30 [rocm-libraries] ROCm/rocm-libraries#8985 (commit 3d4cbef)
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).
2026-07-15 16:12:14 +00:00
..

GEMM C++ Examples

CK Tile Dispatcher C++ examples for GEMM (General Matrix Multiplication) operations.

Main Documentation: Dispatcher README | Examples Overview

Quick Start

Build and Run

cd /path/to/composable_kernel/dispatcher
mkdir -p build && cd build

cmake .. \
  -DCMAKE_PREFIX_PATH=/opt/rocm \
  -DCMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc \
  -DBUILD_DISPATCHER_EXAMPLES=ON

# Build (kernels generated automatically by CMake)
make -j$(nproc)

# Run examples
cd examples
./gemm_01_basic
./gemm_03_benchmark_validation
./gemm_04_heuristics

Examples

Example Description
01_basic_gemm.cpp Basic GEMM with declarative API, autofill, autocorrect
02_multi_size.cpp Wildcard expansion for multiple configurations
03_benchmark_validation.cpp Performance benchmarking with CPU reference validation
04_heuristics.cpp Heuristic-based kernel selection
05_json_export.cpp Registry JSON export for external tools
06_multi_registry.cpp Multiple registries with named kernel sets

Example Details

01_basic_gemm.cpp - Basic GEMM

Demonstrates the declarative kernel API with three patterns:

  1. Autofill Pattern - Minimal specification, defaults filled automatically
  2. Autocorrect Pattern - Invalid parameters corrected at build time
  3. Full Specification Pattern - Complete kernel configuration
DECL_KERNEL_SET(basic_kernels,
    // Pattern 1: Autofill - minimal specification
    .add(
        Signature().dtype("fp16").layout("rcr"),
        Algorithm(),  // Defaults filled by autofill
        "gfx942"
    )
    // Pattern 2: Full specification
    .add(
        Signature().dtype("fp16").layout("rcr"),
        Algorithm().tile(256, 256, 32).wave(2, 2, 1).warp(32, 32, 16)
                   .pipeline("compv4").scheduler("intrawave"),
        "gfx942"
    )
);

Features:

  • Uses generic REGISTER_GENERATED_KERNELS macro
  • print_registered_kernels() utility for debugging
  • Demonstrates autofill messages during build

02_multi_size.cpp - Wildcard Expansion

Demonstrates automatic generation of multiple kernel configurations:

DECL_KERNEL_SET(multi_kernels,
    .add(
        Signature().dtype("fp16").layout("rcr"),
        Algorithm().tile(*, *, 32)     // Wildcard tile M and N
                   .wave(2, 2, 1)
                   .warp(32, 32, 16)
                   .pipeline("compv4")
                   .scheduler("intrawave"),
        "gfx942"
    )
);

Wildcard Values:

  • *, -1, or ANY_INT expand to all valid configurations
  • Architecture filter prunes invalid combinations automatically
  • Example generates 5 valid kernels after arch filtering (from 7 expansions)

03_benchmark_validation.cpp - Benchmark + Validation

Consolidated example combining performance benchmarking with correctness validation:

# Benchmark only
./gemm_03_benchmark_validation --warmup 10 --iterations 100

# With CPU validation
./gemm_03_benchmark_validation --verify 1 --rtol 1e-3 --atol 1e-3

# With GPU reference validation (faster for large matrices)
./gemm_03_benchmark_validation --verify 2

Features:

  • Warmup iterations (discarded from timing)
  • Benchmark iterations with statistics (min/max/mean/median)
  • CPU reference validation using ck_tile::reference_gemm
  • GPU reference validation using ck_tile::reference_gemm_gpu
  • Configurable tolerances

04_heuristics.cpp - Heuristic Selection

Demonstrates custom kernel selection based on problem characteristics:

// Problem size analysis
auto heuristic = [](const Problem& p) -> std::optional<KernelKey> {
    if (p.M() * p.N() < 256 * 256) {
        return small_kernel_key;   // Memory-bound heuristic
    } else {
        return large_kernel_key;   // Compute-bound heuristic
    }
};

dispatcher.set_heuristic(heuristic);

Features:

  • Problem size analysis (small vs large matrices)
  • Compute-bound vs memory-bound selection
  • Custom heuristic function registration

05_json_export.cpp - JSON Export

Exports registry information to JSON for external tool integration:

auto json = registry.to_json();
std::ofstream file("kernels.json");
file << json;

Use Cases:

  • Kernel metadata serialization
  • External analysis tools
  • Configuration management

06_multi_registry.cpp - Multiple Registries

Demonstrates using multiple registries with named kernel sets:

// Define separate kernel sets
DECL_KERNEL_SET(compute_optimized, ...);
DECL_KERNEL_SET(latency_optimized, ...);

// Register to specific registries
Registry compute_registry, latency_registry;
REGISTER_KERNEL_SET(compute_optimized, compute_registry);
REGISTER_KERNEL_SET(latency_optimized, latency_registry);

// Use appropriate registry based on workload
Dispatcher compute_dispatcher(compute_registry);
Dispatcher latency_dispatcher(latency_registry);

Features:

  • Named kernel set registration with REGISTER_KERNEL_SET macro
  • Separate registries for different optimization goals
  • Dynamic kernel set selection by name

Benchmark Parameters (stream_config)

CK Tile uses stream_config for benchmark control:

ck_tile::stream_config cfg{
    nullptr,    // stream_id       - HIP stream (nullptr = default)
    true,       // time_kernel     - Enable timing
    1,          // log_level       - Verbosity (0=quiet, 1=normal)
    5,          // cold_niters     - Warmup iterations
    20,         // nrepeat         - Benchmark iterations
    true,       // is_gpu_timer    - Use GPU events vs CPU chrono
    false,      // flush_cache     - Flush L2 cache between iterations
    1           // rotating_count  - Rotating buffers for cache simulation
};
Parameter CLI Option Default Description
cold_niters_ --warmup 5 Warmup iterations
nrepeat_ --iterations 100 Benchmark iterations
flush_cache_ - false Flush L2 cache
rotating_count_ - 1 Rotating buffers
is_gpu_timer_ - true GPU timer vs CPU

Declarative Kernel Pattern

All examples use the declarative DECL_KERNEL_SET macro:

DECL_KERNEL_SET(my_kernels,
    .add(
        Signature()               // WHAT: operation signature
            .dtype("fp16")        // Data type
            .layout("rcr"),       // Matrix layouts (A=row, B=col, C=row)
        Algorithm()               // HOW: implementation details  
            .tile(256, 256, 32)   // Tile sizes (M, N, K)
            .wave(2, 2, 1)        // Wave configuration
            .warp(32, 32, 16)     // Warp tile sizes
            .pipeline("compv4")   // Pipeline type
            .scheduler("intrawave"), // Scheduler type
        "gfx942"                  // WHERE: target architecture
    )
);

Key Macros:

  • DECL_KERNEL_SET(name, ...) - Declare a kernel set
  • REGISTER_GENERATED_KERNELS - Register all kernels from this example
  • REGISTER_KERNEL_SET(name, registry) - Register specific kernel set to a registry