[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).