ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 17:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity
3,461 Commits 779 Branches 40 Tags
6648115aed5183e9906ec8fc0b368635c5fb83a4
Commit Graph

3 Commits

Author SHA1 Message Date
Muhammed Emin Ozturk
6648115aed [rocm-libraries] ROCm/rocm-libraries#9000 (commit 9faa8de) 
feat(ck-tile): add grouped GEMM variant to TE to dispatcher
 bridge (#9000)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

> Re-opened from #8130 with a policy-compliant branch name
(`users/muozturk/ck-tile/dispatcher-te-bridge-grouped-gemm`). Supersedes
#8130.

## What this PR does

Routes the **grouped_gemm** variant through the Tile Engine (TE) →
Dispatcher **bridge**: TE only generates configs and benchmarks; the
Dispatcher owns codegen, build, and runtime. This is the grouped
counterpart of the regular-GEMM bridge (#8123/#8479), the fp8/bf8/int8
bridge (#8887), and the Stream-K bridge (#8136).

**This PR now also contains the grouped Dispatcher codegen** that
previously lived in #8075 — that PR has been **closed in favor of this
one** to keep the grouped codegen in a single place (it was otherwise
duplicated across both).

## Why grouped needs special handling

Grouped GEMM is **multi-problem**: one launch runs a *list* of `(M, N,
K)` sub-problems with arrays of A/B/C device pointers.

1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`)
cannot express a list of problems.
2. The generated registry wrapper (`generated_tile_backend.hpp::run()`)
hard-codes the single-problem launch and won't compile against a grouped
`SelectedKernel`.

So the grouped path **bypasses the registry**: a dedicated ctypes lib
calls the generated `SelectedKernel::launch(descs, stream)` directly and
reports the name from the compile-time `KERNEL_NAME` macro.

## Changes

**Codegen (absorbed from #8075)**
- `codegen/arch_filter.py` — `GEMM_GROUPED` operator tile constraints.
- `codegen/unified_gemm_codegen.py` — `GemmVariant.GROUPED`, the grouped
launch generator (DeviceMem internal workspace via `MakeKargs`,
persistent/non-persistent grid), `grouped` in `--variants`.
- `examples/gemm/cpp/02_grouped_gemm_driver.cpp` — standalone,
layout/dtype-generic grouped driver with per-group reference
verification.
- `codegen/README.md` + `examples/gemm/cpp/README.md` — grouped
sections.

**Bridge**
- `bindings/ctypes/grouped_gemm_ctypes_lib.cpp` — multi-problem,
registry-bypass C ABI; per-group device alloc/copy; strides derived from
the compile-time `ALayout/BLayout/CLayout`; warmup/repeat timing matched
to Old-TE (`CK_TILE_BENCH_WARMUP/REPEAT`).
- `python/gemm_utils.py` — `GroupedGemmProblem`/`GroupedGemmResult`,
`GpuGroupedGemmRunner`, `run_grouped`, fp16/bf16/fp8(E4M3 FNUZ)/bf8(E5M2
FNUZ) codecs, output-dtype-aware C buffer.
- `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` +
`run_one_grouped_gemm_kernel.py` — TE driver + worker for the parity
sweep.
- `bindings/ctypes/GROUPED_GEMM_BRIDGE.md` — design README.

## Coverage (= Old-TE grouped runnable set on develop)

| Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) |
|---|---|---|---|---|
| rcr / rrr / ccr / crr | ✓ | ✓ | ✓ | ✓ |

C is always row-major. `int8` (rejected by the TE grouped builder) and
`fp32`/`fp64` (no MFMA warp tiles) are excluded on both sides.

## Parity vs Old-TE (MI300X / gfx942)

Apples-to-apples (same warmup=50/repeat=100 both sides, A/B interleaved,
single GPU, both engines rebuilt fresh, stale-`.so` guard, matched
compile flags):

- **Correctness: 64/64 PASS.**
- **Performance: 64/64 within ±15%.**
- The 5 small-shape (1024³ fp8/bf8) rows that initially read >15% were
proven by `rocprof` to be a **measurement-harness artifact** (Old-TE's
JSON `latency(ms)` rounded to 2 decimals → 30–50% TFLOPS swing on ~0.02
ms kernels), **not** a kernel/codegen difference — bridge and Old-TE
launch byte-identical kernels (same grid/VGPR/SGPR, duration ≤3.22%);
full-precision re-measure collapses all 5 to <3%.

## Notes

- Targets `develop`. Depends on #8997 (fp16/bf16 bridge) and #8998
(fp8/bf8/int8 bridge) merging to `develop` first; until then this PR's
diff also shows their content, after which it reduces to the
grouped-only files.
- Supersedes #8075 (closed).
 2026-07-16 02:55:42 +00:00
Muhammed Emin Ozturk
1c455b1bf5 [rocm-libraries] ROCm/rocm-libraries#8998 (commit 5501ef1) 
feat(ck-tile): TE to dispatcher GEMM bridge for fp8/bf8/int8
 (all layouts) (#8998)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## 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.
 2026-07-09 23:58:00 +00:00
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