mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-08 16:17:16 +00:00

Files

Muhammed Emin Ozturk de292a24f9 [rocm-libraries] ROCm/rocm-libraries#8997 (commit 6e9bfd9)

feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all
 layouts) (#8997)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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

2026-07-07 01:15:38 +00:00

CMakeLists.txt

[rocm-libraries] ROCm/rocm-libraries#5168 (commit 8b5afcb)

2026-04-09 10:38:33 -07:00

ctypes_utils.py

[rocm-libraries] ROCm/rocm-libraries#8997 (commit 6e9bfd9)

2026-07-07 01:15:38 +00:00

dispatcher_common.py

[rocm-libraries] ROCm/rocm-libraries#5260 (commit a1834d2)

2026-05-17 00:29:40 -07:00

fmha_utils.py

[rocm-libraries] ROCm/rocm-libraries#5260 (commit a1834d2)

2026-05-17 00:29:40 -07:00

gemm_utils.py

[rocm-libraries] ROCm/rocm-libraries#8997 (commit 6e9bfd9)

2026-07-07 01:15:38 +00:00

grouped_conv_utils.py

[rocm-libraries] ROCm/rocm-libraries#6327 (commit 1e7a12e)

2026-05-08 13:47:13 -07:00

pytest.ini

Adding dispatcher architecture (#3300 )

2026-01-22 09:34:33 -08:00

README.md

[rocm-libraries] ROCm/rocm-libraries#5168 (commit 8b5afcb)

2026-04-09 10:38:33 -07:00

requirements.txt

[rocm-libraries] ROCm/rocm-libraries#5676 (commit 1d18339)

2026-04-01 19:25:55 -07:00

README.md

CK Tile Dispatcher Python Utilities

This directory contains Python utilities used by the dispatcher examples.

Shared Utilities (used by both GEMM and Grouped Conv)

dispatcher_common.py - Shared dispatcher infrastructure
- Path helpers (get_dispatcher_root, get_build_dir, etc.)
- ValidationResultBase - Structured validation feedback
- validate_wave_config, validate_warp_tile_config, validate_trait_combo
- auto_correct_wave, auto_correct_trait - Auto-correction helpers
- Colors - Cross-platform ANSI color support
- print_phase, print_success, print_error, print_info - Phased output
- cleanup_generated_kernels - Cleanup helper

GEMM Utilities

ctypes_utils.py - Core ctypes utilities for GEMM Python examples
- KernelConfig - Kernel configuration dataclass
- setup_gemm_dispatcher() - Setup dispatcher with auto-correction
- cleanup_gemm() - Cleanup dispatcher resources
- GemmRunner - GPU execution helper
- Auto-correction and validation utilities

Grouped Convolution Utilities

grouped_conv_utils.py - Utilities for grouped convolution
- GroupedConvValidationResult - Validation result (extends ValidationResultBase)
- validate_grouped_conv_config - Validate a grouped conv config
- auto_correct_grouped_conv_config - Auto-correct invalid configs
- get_grouped_conv_default_config - Get default config for a variant
- GroupedConvDataType - Data type enum (FP16, BF16, FP32, FP8, BF8, INT8)
- format_grouped_conv_summary - Human-readable config summary

Usage

GEMM Examples

The GEMM Python examples in dispatcher/examples/gemm/python/ import:

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))

from ctypes_utils import (
    KernelConfig,
    setup_gemm_dispatcher,
    cleanup_gemm,
    GemmRunner,
)

Grouped Conv Usage

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))

from grouped_conv_utils import (
    validate_grouped_conv_config,
    auto_correct_grouped_conv_config,
    get_grouped_conv_default_config,
    GroupedConvDataType,
)

# Get a default config
config = get_grouped_conv_default_config(variant="forward", arch="gfx942")

# Validate
result = validate_grouped_conv_config(config)
print(f"Valid: {result.is_valid}")

Requirements

Python 3.8+
NumPy
HIP runtime (for GPU execution)

README.md

CK Tile Dispatcher Python Utilities

Contents

Shared Utilities (used by both GEMM and Grouped Conv)

GEMM Utilities

Grouped Convolution Utilities

Usage

GEMM Examples

Grouped Conv Usage

Requirements