mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-18 01:28:27 +00:00
58c1bcbe2efb88e75dcf9099cace9827c3f887cc
4 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
58c1bcbe2e |
[rocm-libraries] ROCm/rocm-libraries#9028 (commit 2d6a3d6)
feat(ck-tile): stream-K GEMM TE to dispatcher bridge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > Re-opened from #8136 with a policy-compliant branch name. Supersedes #8136. ## Summary Routes the **stream_k** GEMM variant through the same Tile Engine (TE) → Dispatcher bridge already landed for regular GEMM (#8123) and grouped GEMM (#8130). The Dispatcher stays the single source of truth for codegen/build/runtime; TE only produces configs and benchmarks. ## Design note Stream-K is a single-problem GEMM with the **same C ABI** as regular GEMM, so the Python runner (`GpuGemmRunner`/`GemmProblem`) and the GPU worker are reused unchanged. The one twist: the registry path can't compile against a Stream-K `SelectedKernel`, so the Stream-K ctypes lib **bypasses the registry** and calls `SelectedKernel::launch(args, stream)` directly — the same approach grouped GEMM uses. The generated launch owns the reduction workspace internally (`DeviceMem`) and uses the Atomic strategy. ## Changes **New** - `bindings/ctypes/streamk_gemm_ctypes_lib.cpp` — single-problem C ABI (`dispatcher_run_gemm`), builds `StreamKHostArgs`, direct launch; returns `0` / `-1` (HIP/throw) / `-2` (args unsupported). - `tile_engine/ops/gemm/streamk_gemm_full_benchmark.py` + `run_one_streamk_gemm_kernel.py` — 3-phase driver (expand → build → subprocess-isolated benchmark) and disposable GPU worker. - `tile_engine/ops/gemm/gemm_streamk/configs/default_config.json` — small sweep config. **Modified** - `dispatcher/python/gemm_utils.py`, `ctypes_utils.py` — thread `variant="stream_k"` through codegen/build and `.so` selection. ## Validation fp16/rcr on gfx942/MI300X: numeric parity vs an fp32 numpy reference (widened fp16-atomic tolerance) and a full driver run of **16/16 OK** with end-to-end name parity; unsupported tiny shapes are reported gracefully (`status -2`), not crashes. fp8/bf8/bf16 now supported via `ml_dtypes` FNUZ codecs. Full tables and the bridge-vs-old-TE comparison are in the comments. ## Test plan - [x] codegen emits `*_streamk.hpp` with stem == `GemmKernelConfig(variant="stream_k").name` - [x] build/link against `streamk_gemm_ctypes_lib.cpp` - [x] numeric parity passes (fp16 atomic tolerance) - [x] full driver run 16/16 OK, name parity end-to-end - [x] unsupported shape → `status -2` handled gracefully ## Next Land #8123, then this; afterwards remove the legacy `tile_engine/ops/gemm_streamk/` machinery (Phase 4). |
||
|
|
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). |
||
|
|
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. |
||
|
|
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. |