feat(ck_tile): TE -> Dispatcher GEMM bridge (all layouts, fp16/bf16)#8479
feat(ck_tile): TE -> Dispatcher GEMM bridge (all layouts, fp16/bf16)#8479ozturkosu wants to merge 16 commits into
Conversation
Consolidated, single-commit GEMM bridge routing the Tile Engine regular-GEMM sweep through the Dispatcher (codegen -> build -> runtime), so the Dispatcher is the single source of truth and the Tile Engine owns only the config search space and the benchmark loop. Mirrors the FMHA/Conv reference binding end to end. Scope: - Regular GEMM bridge: unified_gemm_codegen.py, gemm_ctypes_lib.cpp (flat extern "C" ABI, host-pointer model), gemm_utils.py (GemmKernelConfig with byte-exact .name, one-.so-per-kernel build), 3-phase TE driver + subprocess worker (gemm_full_benchmark.py / run_one_gemm_kernel.py). - Trait-derived registry KernelKey (replaces the hard-coded fp16/rcr key). - bf16 support and all four layouts (rcr/rrr/crr/ccr; row-major C only). - Tile Engine AMDGPU -mllvm codegen-flag parity + arch-validated tile filtering. - --verify fp32-reference correctness gate; multi-GPU fan-out. - Runnable example (examples/gemm/python/12_te_bridge.py) and parity/unit tests. - Removes the legacy standalone gemm_universal build path and the old test/ck_tile/gemm_tile_engine harness; promotes sweep configs to the op-root flat configs/ directory (fmha/grouped_conv convention). Validated on gfx942 / MI300X (fp16 + bf16, all four layouts) against an fp32 numpy reference via --verify.
There was a problem hiding this comment.
Pull request overview
This PR routes Tile Engine regular-GEMM sweeping/benchmarking through the Dispatcher bridge, making the Dispatcher the single codegen→build→runtime path while TE retains config expansion and the benchmark driver/worker. It also removes the legacy standalone gemm_universal build/test harness and adds parity/unit tests plus an end-to-end example.
Changes:
- Add the Dispatcher-owned GEMM bridge (
dispatcher/python/gemm_utils.py) and TE driver/worker scripts for subprocess-isolated, multi-GPU benchmarking. - Extend dispatcher codegen/ctypes runtime to register kernels with trait-derived
KernelKeymetadata and add benchmark knob parity. - Remove the legacy
gemm_universalstandalone build path and the oldgemm_tile_enginetest harness; add new bridge-focused tests and examples.
Reviewed changes
Copilot reviewed 33 out of 35 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py | New per-kernel subprocess worker for loading .so and running/optionally verifying GEMM. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py | New 3-phase driver: parallel build + problem load + multi-GPU batched benchmarking. |
| projects/composablekernel/tile_engine/ops/gemm/README.md | Documents the dispatcher bridge workflow, scripts, layouts, and removal of standalone path. |
| projects/composablekernel/tile_engine/ops/gemm/configs/default_ci_config.json | Tweaks CI sweep tile params. |
| projects/composablekernel/tile_engine/ops/gemm/configs/default_config.json | Adds full sweep config for bridged GEMM. |
| projects/composablekernel/tile_engine/ops/gemm/configs/example_problems.json | Adds example GEMM problem shapes for the driver. |
| projects/composablekernel/tile_engine/ops/gemm/configs/user_provided_config.json | Adds a user scratch sweep config template. |
| projects/composablekernel/tile_engine/ops/gemm/CMakeLists.txt | Drops gemm_universal subdir and updates sampling-op lists accordingly. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_profiler.hpp | Removes legacy standalone GEMM universal profiler. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py | Removes legacy standalone GEMM universal instance builder. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py | Removes legacy standalone GEMM universal benchmark driver. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.hpp | Removes legacy standalone GEMM universal benchmark helpers. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark_single.cpp | Removes legacy single-kernel benchmark executable. |
| projects/composablekernel/tile_engine/ops/gemm/gemm_universal/CMakeLists.txt | Removes legacy standalone build integration for universal GEMM. |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/test_gemm_simple.cpp | Removes old tile-engine GEMM gtest harness. |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/README.md | Removes documentation for the deleted GEMM tile-engine tests. |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/extract_test_params.py | Removes helper used by deleted GEMM tile-engine tests. |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/small_datatype_config.json | Removes old test configs (small datatype). |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json | Removes old test configs (simple). |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/quick_coverage_config.json | Removes old test configs (quick coverage). |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/padding_coverage_config.json | Removes old test configs (padding coverage). |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/large_datatype_config.json | Removes old test configs (large datatype). |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/comprehensive_coverage_config.json | Removes old test configs (comprehensive coverage). |
| projects/composablekernel/test/ck_tile/gemm_tile_engine/CMakeLists.txt | Removes CMake for the deleted GEMM tile-engine tests. |
| projects/composablekernel/test/ck_tile/CMakeLists.txt | Removes commented-out gemm_tile_engine test subdir stanza. |
| projects/composablekernel/dispatcher/python/gemm_utils.py | New bridge contract + build API + ctypes wrapper + runner + TE sweep expansion. |
| projects/composablekernel/dispatcher/python/ctypes_utils.py | Threads requested GEMM variant into codegen (--variants). |
| projects/composablekernel/dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp | Builds registry KernelKey from codegen-emitted macros; adds indexed kernel-name ABI. |
| projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py | Emits GEMM_KEY_* macros; fixes arch validation to respect actual pipeline/scheduler. |
| projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp | Maps scheduler string "default" to Scheduler::Auto. |
| projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp | Adds env-overridable benchmark knobs and updates defaults to match TE. |
| projects/composablekernel/dispatcher/tests/test_gemm_utils.py | New CPU-only tests for bf16 encoding and name parsing/contract. |
| projects/composablekernel/dispatcher/tests/test_gemm_parity.py | New GPU-gated parity suite vs NumPy reference across dtype/layout surface. |
| projects/composablekernel/dispatcher/examples/gemm/python/12_te_bridge.py | New end-to-end bridge example (matrix/layouts, padding behavior, mini-sweep). |
| projects/composablekernel/dispatcher/parity_diag/regression/ab_same_harness.py | Adds diagnostic harness to compare bridge vs old-TE kernel under the same worker. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
The bridge dispatcher's tile-divisibility gate rejected any problem where
M % TileM != 0 for every layout, returning status -2 ("No suitable kernel")
at runtime even though the .so built fine. This wrongly excluded bf16 rcr/rrr
kernels with a non-power-of-two TileM (e.g. 192) on standard shapes like
1024^3 -- cases Old-TE compiles, runs, and verifies as correct.
Root cause: supports() was layout-blind, while the underlying
ck_tile::GemmKernel::IsSupportedArgument only constrains a dimension when an
operand whose inner axis is that dimension participates without padding:
RowMajor A -> K, ColMajor A -> M
RowMajor B -> N, ColMajor B -> K
RowMajor C -> N, ColMajor C -> M
So for rcr (RowMajor A & C) M is never gated, which is why Old-TE runs M=192
tiles on M-indivisible problems.
Make supports() compute require_m/n/k from the kernel key's A/B/C layouts so
it mirrors IsSupportedArgument exactly (also honoring k_batch in the K grain).
Anything it now lets through is still validated by the kernel's own
IsSupportedArgument inside launch(), so the bridge stays a strict functional
equivalent of Old-TE. Applied to both generated_tile_backend.hpp (the GEMM
.so path) and the sibling tile_backend.hpp.
Validated on gfx942 (MI300X): 85 previously status-2 rcr/rrr bf16 192-tile
.so now run at 1024^3 (Old-TE runs the same, verification correct); the 8
remaining rejects are tile N=192 cases that Old-TE also reports "Arguments
not supported" at N=1024 -- parity preserved in both directions.
Fix: layout-aware
|
| layout | RowMajor gates | ColMajor gates |
|---|---|---|
| A | K | M |
| B | N | K |
| C | N | M |
For rcr/rrr (RowMajor A & C), M is never gated — which is why Old-TE runs TileM=192 on M-indivisible problems.
Fix
supports() now computes require_m/n/k from the kernel key's A/B/C layouts so it mirrors IsSupportedArgument exactly (K grain also honors k_batch). Anything it lets through is still validated by the kernel's own IsSupportedArgument inside launch(), so the bridge stays a strict functional equivalent of Old-TE. Applied to both generated_tile_backend.hpp (the GEMM .so path) and the sibling tile_backend.hpp.
Validation (gfx942 / MI300X)
- Old-TE baseline, TileM=192 at M=1024³, CPU verify:
rcr&rrr× {192x128x64, 192x256x64, 192x64x128, 192x64x64} → allverification result: correct. - Bridge after fix: 85/93 rebuilt
rcr/rrrTileM=192.sonow run at 1024³ (were status -2). The 8 still rejected are TileN=192 (192x192x64) cases, which Old-TE also rejects with"Arguments not supported"at N=1024 — and both run at N=1152. Parity holds in both directions. crr/ccrat TileM=192 are out of scope: they fail to compile on both engines (ColMajor A →static_assert(X0*Y1==warp_size)), so there is no Old-TE binary or bridge.sofor them — not a parity gap.
…oding rcr dispatcher_initialize() in gemm_ctypes_lib.cpp hardcoded the KernelKey layout to rcr (RowMajor/ColMajor/RowMajor) for every kernel. Now that supports() is layout-aware, that wrong key layout makes the dispatcher reject valid problems: a crr kernel does not gate K (neither A=ColMajor nor B=RowMajor has K as its inner axis), but with a hardcoded rcr key supports() applies rcr's K-gate and returns status -2 for TileK=192 problems (e.g. crr 64x64x192 at 1024^3) that Old-TE compiles, runs, and verifies (~87 TFLOPS). Derive signature.layout_a/b/c from the force-included kernel's own ALayout/BLayout/CLayout types via std::is_same_v with tensor_layout::gemm::RowMajor. The key now matches the kernel, so the layout-aware gate is correct for all four layouts. Execution was already layout-correct (the kernel uses its own compile-time layouts); only the host-side selection metadata was wrong. Validated on gfx942 (MI300X): crr 64x64x192 now runs on the bridge (93 TFLOPS), restoring parity with Old-TE.
The >=20% bridge-vs-old-TE perf gaps in the parity sweep are a harness artifact: the sweep timed the bridge in-process but timed old-TE via its separate standalone benchmark binary, which runs the byte-identical kernel at a lower sustained SCLK. Measured through one harness the gap is <1%. ab_same_harness.py removed that artifact but hardcoded the old-TE header dir to fp16/rcr. Derive it per stem as <base>/<dtype>/<layout> so one run covers rcr/rrr/ccr/crr and fp16+bf16, add a --stems-file/--csv resume-aware sweep mode, and use the median (not max) per point.
For a full ~2000-stem sweep on a single GPU: batch all shapes into one worker call per side (5x fewer process startups), cache the compiled old-TE .so, and add a parallel --build-only pre-pass so hipcc compilation uses all CPU cores while GPU measurement stays serial.
Follow-up: the ">=20% bridge-vs-old-TE perf gaps" are a harness artifact, not a kernel differenceWhile reviewing the parity sweep Root cause — apples-to-oranges measurement. The sweep timed the bridge in-process (via Proof — measure both sides through ONE harness (build the old-TE kernel into a
Holds across rcr/rrr/ccr/crr and fp16+bf16; worst residual across spot checks is +2.75% (single 2048³ wave-quant noise sample). No kernel/codegen/config change is needed — the bridge is performance-equivalent to old-TE. Fix (this PR). Generalized the same-harness A/B tool ( |
|
Follow-up perf opportunity (not in this PR): gate -enable-post-misched=0 on persistent While validating bridge↔old-TE parity I confirmed there is no bridge codegen regression — the What: In gemm_utils.py, gate the -mllvm -enable-post-misched=0 codegen flag on config.persistent Why: A flag-bisection (same source, 8 variants) shows this is the only perf-relevant flag, and ┌─────────────────────────────────┬─────────┬────────────────────┐ The flag helps persistent kernels (occupancy-bound, +6%) but hurts non-persistent large tiles Why held back:
Proposing it as a separate follow-up PR once the broader sweep confirms no persistent That wraps it up. Summary of where things stand:
|
… guard) The bridge-vs-old-TE A/B reported phantom regressions from two MEASUREMENT bugs, not real codegen gaps: - ab_same_harness.py built the old-TE side WITHOUT the TE codegen flags the bridge (and real old-TE's own CMake) use, so -enable-post-misched defaulted back on and old-TE ran ~10-40% faster -> the bridge looked regressed when it is at parity. Now both sides build with identical flags. - ab_efficient_sweep.py measured whatever libgemm_<stem>.so existed with no freshness check, so 3-day-old binaries built from an obsolete codegen showed up as -78%/+703% gaps. Added a guard: skip any .so older than its generated header (treated as missing) instead of reporting a phantom gap. With both fixes the 41 former >15% outlier stems measure within +/-10% (median +0.01%); no bridge codegen regression exists. Note: a separate, deliberately UNCOMMITTED perf change in gemm_utils.py (gate -enable-post-misched=0 on persistent) gives non-persistent large tiles ~9-40%; held back pending a broader persistent-kernel no-regression sweep.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Codegen validity gate: reject invalid non-power-of-2-repeat tiles (commit a7e01c4)Follow-up to the 192-tile correctness finding. The CShuffle epilogue stores the accumulator through LDS in power-of-two Verified on MI350: This commit adds the gate to both tile-expansion paths so invalid tiles are dropped rather than emitted/run:
(The parity harness was also updated locally to run with fp32 verification on, so a numerically-wrong kernel is recorded as |
Re-stacks the grouped_gemm bridge on top of the all-layouts GEMM bridge (muozturk/gemm-bridge-all-layouts-bf16, #8479) instead of the now-closed #8123 base. Squashes the prior grouped-branch history (was 5e665a3.. d075b90 on the old base) into the net grouped delta replayed onto #8479. Contents: - codegen: grouped variant + launch generator (arch_filter.py, unified_gemm_codegen.py), plus the standalone 02_grouped_gemm_driver.cpp and README sections (absorbed from the closed #8075). - bridge: multi-problem registry-bypass ctypes ABI (grouped_gemm_ctypes_lib.cpp), GpuGroupedGemmRunner + dtype/layout codecs (gemm_utils.py), TE driver/worker harness, GROUPED_GEMM_BRIDGE.md. Kept both #8479's GEMM_KEY_* registry descriptors and grouped's ALayout/BLayout/CLayout single-include exports. Coverage: {rcr,rrr,ccr,crr} x {fp16,bf16,fp8,bf8}; validated at Old-TE parity on MI300X/gfx942.
Grouped GEMM variant of the bridge, stacked on the fp8/bf8/int8 bridge (#8887, itself on #8479). The grouped kernel is multi-problem, so it uses a dedicated registry-bypass ctypes ABI; TE only generates configs and benchmarks. - codegen: grouped variant + launch generator (arch_filter.py, unified_gemm_codegen.py), standalone 02_grouped_gemm_driver.cpp, README. - bridge: grouped_gemm_ctypes_lib.cpp (multi-problem ABI), GpuGroupedGemmRunner + dtype/layout codecs (gemm_utils.py), TE driver/worker harness, GROUPED_GEMM_BRIDGE.md. Coverage: {rcr,rrr,ccr,crr} x {fp16,bf16,fp8,bf8}; validated at Old-TE parity on MI300X/gfx942 (64/64 correct, 64/64 within +/-15%).
❌ PR Check — Action Required
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🚫 Please fix the failed policies before requesting reviews. The following policy checks failed:
The |
ozturkosu
changed the title
…8997) > Re-opened from #8479 with a compliant branch name (users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16). Supersedes #8479. ## Summary This PR routes the **Tile Engine (TE) regular-GEMM sweep through the Dispatcher**, making the Dispatcher the single source of truth for **codegen → build → runtime** while the Tile Engine keeps only the **config search space** and the **benchmark loop**. It is the consolidated, **single-commit** GEMM bridge covering **all four layouts (`rcr`/`rrr`/`crr`/`ccr`)** and **both `fp16` and `bf16`**. It is a clean re-roll of the earlier bridge work (previously split across #8123 + the stacked key/bf16/layouts/parity/example PRs and consolidated in #8261). Those branches accumulated unrelated cross-project commits through repeated `develop` merges; **this branch is a single clean commit off the latest `develop`** containing only the GEMM-bridge files. It supersedes and replaces #8123 / #8261. ## Motivation The Tile Engine historically owned its own codegen/build/runtime for GEMM (`tile_engine/ops/gemm/gemm_universal/`). The consolidation goal is for the **Dispatcher** to own all of that — exactly as it already does for **FMHA** and **Grouped Conv** — so there is one kernel-generation/build/runtime path and the TE shrinks to a config+benchmark frontend. This PR brings regular GEMM in line with that reference binding. ## The binding (mirrors the FMHA/Conv reference, six stages) 1. **Config JSON (TE side)** — the sweep search space lives in `tile_engine/ops/gemm/configs/` (flat op-root layout, matching the `fmha/` and `grouped_conv/` bridges). 2. **Codegen (Dispatcher)** — `dispatcher/codegen/unified_gemm_codegen.py` emits one fully-typed `.hpp` per kernel; `GemmKernelConfig.name` reproduces `KERNEL_NAME` **byte-for-byte** (the thread tying config → kernel → runtime). 3. **Compile to `.so`** — a single static `gemm_ctypes_lib.cpp` is force-included (`-include <kernel.hpp>`); one `.so` per kernel. 4. **Flat `extern "C"` ABI** — `dispatcher_run_gemm(A, B, C, M, N, K, time_ms)` + the kernel-name enumeration entry points. **Host-pointer** memory model (the C lib `hipMalloc`s internally) — the FMHA-forward branch of the reference. 5. **Python ctypes wrapper** — `dispatcher/python/gemm_utils.py` (`GemmDispatcherLib` + `GpuGemmRunner`). 6. **TE driver (3 phases)** — `gemm_full_benchmark.py` (parallel codegen+build → `expand_sweep` → subprocess-isolated benchmark) + the disposable per-kernel worker `run_one_gemm_kernel.py`. ## What's included **Bridge core** - `dispatcher/codegen/unified_gemm_codegen.py` — GEMM codegen, byte-exact naming. - `dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp` — flat C ABI, host-pointer model. - `dispatcher/python/gemm_utils.py` — `GemmKernelConfig`, multi-kernel build (`setup_multiple_gemm_dispatchers`), `expand_sweep`, one-`.so`-per-kernel. - `tile_engine/ops/gemm/gemm_full_benchmark.py` + `run_one_gemm_kernel.py` — 3-phase, multi-GPU, subprocess-isolated driver/worker. **Feature surface (the point of this PR)** - **All four layouts** `rcr`/`rrr`/`crr`/`ccr` (row-major C only — ck_tile rejects column-major C at build) with layout-aware host transpose. - **`fp16` + `bf16`** (bf16 via uint16 byte-encoding; dtype derived from kernel name). - **Trait-derived registry `KernelKey`** — replaces the earlier hard-coded fp16/rcr key so the registry path generalizes across dtype/layout/tile. **Correctness & performance hygiene** - **`--verify`** opt-in fp32 numpy-reference gate (global `max|out-ref|/max|ref|`), `verified`/`max_rel` columns in the CSV; a mismatch counts as a failure. - **Tile Engine AMDGPU `-mllvm` codegen-flag parity** (without these the kernel builds with different occupancy and the timing diverges) and **arch-validated tile filtering** against the real pipeline/scheduler. - **Multi-GPU** fan-out across all visible GPUs (`--devices`, device-pinned `HIP_VISIBLE_DEVICES` workers). **Example & tests** - `dispatcher/examples/gemm/python/12_te_bridge.py` — runnable end-to-end example. - `dispatcher/tests/test_gemm_parity.py`, `test_gemm_utils.py`, and a parity regression harness. **Cleanup** - Removes the legacy standalone `gemm_universal` build path (`gemm_universal_instance_builder.py`, `*_benchmark*.{py,cpp,hpp}`, `gemm_universal/CMakeLists.txt`) and the old `test/ck_tile/gemm_tile_engine/` harness; promotes the sweep configs to the flat op-root `configs/`. ## Design decisions (consistent with the reference) - **Host-pointer memory ownership** (C lib owns device memory) — matches FMHA-forward; the Python runner passes host numpy arrays straight through. - **One `.so` per kernel** — packaging choice; the multi-kernel name ABI is retained (`get_kernel_name_at(0)` reports the single kernel), so the Python enumeration path is unchanged from FMHA/Conv. - **Flat `configs/`** at the op root — matches the `fmha/`/`grouped_conv/` convention; the not-yet-bridged variants keep their per-variant `configs/` dirs, selected by `--variant`. ## Validation (gfx942 / MI300X) - Bridge build + benchmark + `--verify` across **`fp16` and `bf16`** and **all four layouts**, checked against an fp32 numpy reference (`A @ B`). - **Name parity** holds end-to-end: each `.so`'s reported runtime name equals `GemmKernelConfig(...).name`. - bf16 passes under a widened fp16/bf16 tolerance; fp16 within the standard `max_rel` gate. ## Test plan - [ ] `gemm_full_benchmark.py --verify` over `configs/default_ci_config.json` for `fp16` and `bf16`, each of `rcr`/`rrr`/`crr`/`ccr`. - [ ] `unified_gemm_codegen.py` emits a header whose stem == `GemmKernelConfig.name`. - [ ] `setup_multiple_gemm_dispatchers` builds + links each config against `gemm_ctypes_lib.cpp`. - [ ] `pytest dispatcher/tests/test_gemm_parity.py dispatcher/tests/test_gemm_utils.py`. - [ ] `examples/gemm/python/12_te_bridge.py` runs end to end. ## Notes - Single clean commit off the latest `develop`; the diff is **35 files, all under `projects/composablekernel/`** (dispatcher + tile_engine/ops/gemm + test/ck_tile). - **Supersedes #8123 and #8261**, which will be closed. - Stream-K (#8136) and grouped GEMM are separate bridge efforts, not in this PR. --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Muhammed Ozturk <muozturk@ctr2-alola-ctrl-01.amd.com>
…outs) (#8998) ## Summary Extends the Tile Engine ↔ Dispatcher GEMM **bridge** to the remaining data types TE's plain GEMM has MFMA warp tiles for, beyond the fp16/bf16 surface of #8479: - **fp8** (E4M3) and **bf8** (E5M2) → fp16 output, fp32 accumulate - **int8** → int32 output and accumulate (gfx942) All four A/B layout combinations per dtype (row-major C only, matching #8479). `fp32`/`fp64` are intentionally **excluded** — they appear in TE's dtype-string map but have no MFMA warp tiles in `GEMM_WARP_TILE_SUPPORTED_COMBINATIONS`, so no kernel can be generated/run. **Depends on the fp16/bf16 bridge in #8997** (`users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16`), which carries the bridge infrastructure and is not yet merged. This PR targets `develop`, so until #8997 merges its diff also includes the base bridge changes; please merge #8997 first. ## Changes - **Codegen** (`codegen_common.py`, `unified_gemm_codegen.py`): add `int32` to the dtype maps; `get_output_dtype` int8→int32; new `get_acc_dtype` (int8→int32, else fp32); derive `AccDataType`/`CDataType`, the `GEMM_KEY_DTYPE_{C,ACC}` macros, and the registry `dtype_c`/`dtype_acc` from the dtype instead of hard-coding `float`/`fp32`. - **Host harness** (`gemm_utils.py`): fp8/bf8 **FNUZ** (gfx942) uint8 codecs — exact decode (matches device `fp8_t`/`bf8_t`), nearest-representable saturating encode (same pattern as the existing bf16 helper); `GpuGemmRunner.run` encodes A/B and sizes the C buffer per dtype; `expand_sweep` sets `dtype_c`/`dtype_acc`. - **Tests**: `test_gemm_utils.py` adds CPU-only fp8/bf8 codec + output-dtype tests (all green); `test_gemm_parity.py` adds fp8/bf8/int8 cases with dtype-aware inputs/references/tolerances (int8 is bit-exact), GPU-gated like the existing cases. ## Verification done - `test_gemm_utils.py` + `test_codegen_common.py`: **54 passed** (CPU). - Codegen smoke: fp8/int8/fp16 each generate 1 kernel + 1 wrapper, 0 failed; emitted `ADataType/CDataType/AccDataType` and `GEMM_KEY_*` macros are correct (int8→int32_t acc/C; fp8→fp16_t C). - `test_gemm_parity.py` collects 60 cases and skips cleanly without a GPU. - The 16 unrelated failures in `test_examples_integration` / `test_grouped_conv_codegen` / `test_library_caching` are **pre-existing** (verified identical on the base branch; they require a built dispatcher `.a` / GPU). ## Test plan - [x] Merge #8997 (fp16/bf16 bridge), then this reduces to just the fp8/bf8/int8 delta on `develop`. - [x] On an MI300X (gfx942) node: run `python3 tests/test_gemm_parity.py` and confirm fp8/bf8/int8 parity; tune the fp8/bf8 tolerances if needed (current values are first-cut headroom). - [x] FNUZ vs OCP: the fp8/bf8 host codec targets the gfx942 FNUZ format; validate / extend for gfx950 (OCP) before enabling there. --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Muhammed Ozturk <muozturk@ctr2-alola-ctrl-01.amd.com>
Summary
This PR routes the Tile Engine (TE) regular-GEMM sweep through the Dispatcher,
making the Dispatcher the single source of truth for codegen → build → runtime
while the Tile Engine keeps only the config search space and the benchmark
loop. It is the consolidated, single-commit GEMM bridge covering all four
layouts (
rcr/rrr/crr/ccr) and bothfp16andbf16.It is a clean re-roll of the earlier bridge work (previously split across
#8123 + the stacked key/bf16/layouts/parity/example PRs and consolidated in
#8261). Those branches accumulated unrelated cross-project commits through repeated
developmerges; this branch is a single clean commit off the latestdevelopcontaining only the GEMM-bridge files. It supersedes and replaces #8123 / #8261.
Motivation
The Tile Engine historically owned its own codegen/build/runtime for GEMM
(
tile_engine/ops/gemm/gemm_universal/). The consolidation goal is for theDispatcher to own all of that — exactly as it already does for FMHA and
Grouped Conv — so there is one kernel-generation/build/runtime path and the
TE shrinks to a config+benchmark frontend. This PR brings regular GEMM in line
with that reference binding.
The binding (mirrors the FMHA/Conv reference, six stages)
tile_engine/ops/gemm/configs/(flat op-root layout, matching thefmha/andgrouped_conv/bridges).dispatcher/codegen/unified_gemm_codegen.pyemitsone fully-typed
.hppper kernel;GemmKernelConfig.namereproducesKERNEL_NAMEbyte-for-byte (the thread tying config → kernel → runtime)..so— a single staticgemm_ctypes_lib.cppis force-included(
-include <kernel.hpp>); one.soper kernel.extern "C"ABI —dispatcher_run_gemm(A, B, C, M, N, K, time_ms)+the kernel-name enumeration entry points. Host-pointer memory model (the C
lib
hipMallocs internally) — the FMHA-forward branch of the reference.dispatcher/python/gemm_utils.py(
GemmDispatcherLib+GpuGemmRunner).gemm_full_benchmark.py(parallel codegen+build →expand_sweep→ subprocess-isolated benchmark) + the disposable per-kernelworker
run_one_gemm_kernel.py.What's included
Bridge core
dispatcher/codegen/unified_gemm_codegen.py— GEMM codegen, byte-exact naming.dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp— flat C ABI, host-pointer model.dispatcher/python/gemm_utils.py—GemmKernelConfig, multi-kernel build(
setup_multiple_gemm_dispatchers),expand_sweep, one-.so-per-kernel.tile_engine/ops/gemm/gemm_full_benchmark.py+run_one_gemm_kernel.py—3-phase, multi-GPU, subprocess-isolated driver/worker.
Feature surface (the point of this PR)
rcr/rrr/crr/ccr(row-major C only — ck_tile rejectscolumn-major C at build) with layout-aware host transpose.
fp16+bf16(bf16 via uint16 byte-encoding; dtype derived from kernel name).KernelKey— replaces the earlier hard-codedfp16/rcr key so the registry path generalizes across dtype/layout/tile.
Correctness & performance hygiene
--verifyopt-in fp32 numpy-reference gate (globalmax|out-ref|/max|ref|),verified/max_relcolumns in the CSV; a mismatch counts as a failure.-mllvmcodegen-flag parity (without these the kernelbuilds with different occupancy and the timing diverges) and
arch-validated tile filtering against the real pipeline/scheduler.
--devices, device-pinnedHIP_VISIBLE_DEVICESworkers).Example & tests
dispatcher/examples/gemm/python/12_te_bridge.py— runnable end-to-end example.dispatcher/tests/test_gemm_parity.py,test_gemm_utils.py, and a parityregression harness.
Cleanup
gemm_universalbuild path(
gemm_universal_instance_builder.py,*_benchmark*.{py,cpp,hpp},gemm_universal/CMakeLists.txt) and the oldtest/ck_tile/gemm_tile_engine/harness; promotes the sweep configs to the flat op-root
configs/.Design decisions (consistent with the reference)
FMHA-forward; the Python runner passes host numpy arrays straight through.
.soper kernel — packaging choice; the multi-kernel name ABI isretained (
get_kernel_name_at(0)reports the single kernel), so the Pythonenumeration path is unchanged from FMHA/Conv.
configs/at the op root — matches thefmha//grouped_conv/convention; the not-yet-bridged variants keep their per-variant
configs/dirs, selected by
--variant.Validation (gfx942 / MI300X)
--verifyacrossfp16andbf16and allfour layouts, checked against an fp32 numpy reference (
A @ B)..so's reported runtime name equalsGemmKernelConfig(...).name.max_relgate.Test plan
gemm_full_benchmark.py --verifyoverconfigs/default_ci_config.jsonforfp16andbf16, each ofrcr/rrr/crr/ccr.unified_gemm_codegen.pyemits a header whose stem ==GemmKernelConfig.name.setup_multiple_gemm_dispatchersbuilds + links each config againstgemm_ctypes_lib.cpp.pytest dispatcher/tests/test_gemm_parity.py dispatcher/tests/test_gemm_utils.py.examples/gemm/python/12_te_bridge.pyruns end to end.Notes
develop; the diff is 35 files, all underprojects/composablekernel/(dispatcher + tile_engine/ops/gemm + test/ck_tile).