Files
composable_kernel/dispatcher/codegen
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
..

CK Tile Unified Code Generators

Single source of truth for GEMM and Grouped Convolution kernel generation.

See also: Main Dispatcher README for installation and core concepts.

Shared Infrastructure

Both GEMM and Grouped Conv generators share common code via codegen_common.py:

  • TileConfig - Dataclass for tile dimensions
  • TraitConfigBase - Base for kernel trait configurations with arch-aware validation
  • CommonTypeMappings - Dtype-to-C++ type mappings
  • parallel_generate() - Parallel kernel generation with per-kernel progress logging
  • Arch-aware expansion helpers (valid_wave_configs, valid_warp_configs, etc.)

Quick Start

GEMM

cd dispatcher/codegen

# Generate standard FP16 kernels
python3 unified_gemm_codegen.py \
    --output-dir ../build/generated_kernels \
    --datatype fp16 \
    --layout rcr \
    --variants standard

# Generate all variants
python3 unified_gemm_codegen.py \
    --output-dir ../build/generated_kernels \
    --variants standard preshuffle multi_d

Grouped Convolution

cd dispatcher/codegen

# Generate forward FP16 grouped conv kernels
python3 unified_grouped_conv_codegen.py \
    --output-dir ../build/generated_kernels \
    --datatype fp16 \
    --variant forward \
    --ndim-spatial 2

# Generate backward data kernels
python3 unified_grouped_conv_codegen.py \
    --output-dir ../build/generated_kernels \
    --variant backward_data \
    --ndim-spatial 2

Using from Python

from ctypes_utils import CodegenRunner, KernelConfig

# Generate from specific config
config = KernelConfig(tile_m=256, tile_n=256, tile_k=64)
codegen = CodegenRunner()
result = codegen.generate_from_config(config)

# Generate variant
result = codegen.generate("preshuffle")

# Generate all
results = codegen.generate_all()

Command Line Options

Option Values Description
--output-dir path Output directory
--datatype fp16, bf16, fp32, int8 Data type
--layout rcr, rrr, crr, ccr Matrix layouts
--gpu-target gfx942, gfx90a, gfx950 Target GPU
--variants standard, preshuffle, multi_d Kernel variants
--preselected fp16_rcr_essential, etc. Predefined kernel set

Layout Notation

  • R = Row-major, C = Column-major
  • Order: A, B, C (e.g., rcr = A row, B col, C row)

Variants

Standard

Basic GEMM: C = A x B

PreShuffle

Optimized weight access with LDS pre-shuffling. Best for large matrices.

Multi-D

Element-wise fusion: C = op(A x B + D0 + D1 + ...)

Supported ops: PassThrough, MultiDAdd, Relu, Gelu, Sigmoid, Tanh

Output Structure

generated_kernels/
|---- gemm_fp16_rcr_compv4_..._128x128x32_....hpp          # GEMM kernels
|---- gemm_fp16_rcr_compv4_..._preshuffle.hpp
|---- gemm_fp16_rcr_compv4_..._multid_Relu_d1.hpp
|---- grouped_conv_fwd_fp16_nhwgc_..._128x128x32_....hpp   # Grouped conv kernels
+---- ...

Configuration Files

arch_specs.json

GPU architecture specifications (single source of truth):

{
  "architectures": {
    "gfx942": {
      "family": "cdna3",
      "warp_size": 64,
      "warp_configs": [[2, 2, 1], [4, 4, 1]],
      ...
    }
  }
}

preselected_kernels.py

Curated kernel sets for common use cases.

Adding New GPU Support

See ADDING_NEW_GPU.md for complete guide.

Quick steps:

  1. Edit arch_specs.json
  2. Run python generate_arch_specs.py
  3. Rebuild

Troubleshooting

Issue Solution
"Arguments not supported" Check tile config validity
Missing element-wise op Check elementwise_ops.hpp
Compilation errors Verify C++17, include paths

More info: See ../README.md for full documentation.