ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-21 13:29:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

[CK] [CK_Tile] Add GroupConv to Kernel Dispatcher (#5168)

Browse Source 
## Motivation

This PR adds CK Tile group convolution (forward, backward-data,
backward-weight) support to the kernel dispatcher, matching and unifying
with the existing dispatcher GEMM infrastructure in architecture and
usability. The dispatcher provides a unified kernel dispatch system with
both C++ and Python frontends, and until now only supported GEMM
operations. This PR enables framework integrators to use the same
declarative kernel workflow for convolutions as they do for GEMM:
declare kernels, build a registry JIT, select kernels within the
registry at runtime, and dispatch to GPU. Future PRs will include
runtime kernel selection heuristics for autotuning of kernel parameters
based on (problem, hardware arch).

## Technical Details

Grouped convolution support has been added to the CK Tile Dispatcher
with generated_conv_backend.hpp enabling dispatcher.run(in, wei, out,
problem) for all 6 conv variants (fwd/bwdd/bwdw x 2D/3D), runtime
heuristic kernel selection, and GroupedConvKernelKey with full
ConvConfigBase fields. Python side adds parallel JIT via
registry.build(max_workers) and heuristic registry.select(). Includes 7
C++ and 6 Python examples covering all directions with CPU reference
validation, and shared infrastructure improvements (BaseRegistry CRTP,
structured exceptions). As a sanity check, JIT compile times for a
single kernel remains the same and for multiple kernels there is better
parallelism:
Kernels | 1 worker | 8 workers
1 | 7.7 s | 7.7 s
2 | 15.9 s | 8.2 s
4 | 33.4 s | 9.7 s
6 | 52.3 s | 10.2 s

## Test Plan

145 ephemeral unit tests have been added to test basic functionality.
All 30 examples/integration tests run end-to-end on gfx950 (MI350): 7
C++ conv, 7 C++ GEMM, 6 Python conv, 10 Python GEMM. CPU reference
validation for forward, backward-data, and backward-weight (2D) in both
C++ and Python examples pass.

## Test Result

30 examples pass. Peak performance: 132 TFLOPS (Batch-32 forward 56x56),
53 TFLOPS (pointwise 1x1). CPU reference accuracy: max_abs_diff < 0.002
for all directions (fp16 vs fp32 reference).

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

---------

Co-authored-by: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com>
This commit is contained in:
Vidyasagar Ananthan
2026-04-09 10:38:33 -07:00
committed by GitHub
parent fb22cd0c69
commit a2b844d335
86 changed files with 15538 additions and 1500 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
dispatcher/examples/gemm/cpp/02_multi_size.cpp
View File

@@ -21,9 +21,9 @@
* - pipeline: "compv3" -> 1 option (compv4 requires special handling)
* - scheduler: "intrawave" -> 1 option
*
* Raw expansion: 3 × 2 = 6 configs, but arch filter validates each:
* - tile_m must be divisible by (warp_m × warp_tile_m)
* - tile_n must be divisible by (warp_n × warp_tile_n)
* Raw expansion: 3 x 2 = 6 configs, but arch filter validates each:
* - tile_m must be divisible by (warp_m x warp_tile_m)
* - tile_n must be divisible by (warp_n x warp_tile_n)
* - Some wave/warp combos invalid: (4,1,1)+(32,32,16), (1,4,1)+(32,32,16)
* Result: 4 valid wildcard kernels + 1 explicit = 5 total
*
@@ -70,13 +70,13 @@ DECL_KERNEL_SET(multi_size_kernels,
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(64, 64, 64)
.wave(ANY_INT, ANY_INT, 1) // ANY_INT → (1,4,1), (2,2,1), (4,1,1)
.warp(-1, -1, -1) // -1 same as ANY_INT → (16,16,32), (32,32,16)
.pipeline("*") // "*" → valid pipelines
.scheduler("*") // "*" → valid schedulers
.wave(ANY_INT, ANY_INT, 1) // ANY_INT -> (1,4,1), (2,2,1), (4,1,1)
.warp(-1, -1, -1) // -1 same as ANY_INT -> (16,16,32), (32,32,16)
.pipeline("*") // "*" -> valid pipelines
.scheduler("*") // "*" -> valid schedulers
.epilogue("cshuffle"),
"gfx942"));
// Raw: 3×2=6, arch filter removes 2 invalid → 4 valid kernels
// Raw: 3x2=6, arch filter removes 2 invalid -> 4 valid kernels
// =============================================================================
// MAIN
@@ -116,8 +116,8 @@ int main(int argc, char* argv[])
.pipeline("*") -> expands to valid pipelines = 1
.scheduler("*") -> expands to valid schedulers = 1
Expanded: 3 × 2 = 6 configs, but arch filter validates each:
- wave×warp must divide tile: (4,1,1)×(32,32,16) invalid for 64x64
Expanded: 3 x 2 = 6 configs, but arch filter validates each:
- wave x warp must divide tile: (4,1,1)x(32,32,16) invalid for 64x64
- Result: 4 valid kernels from wildcard + 1 explicit = 5 total
)";

191
dispatcher/examples/gemm/cpp/07_gfx950_minimal.cpp Normal file
View File

@@ -0,0 +1,191 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/**
* Example 07: Minimal gfx950 (CDNA4 / MI350) GEMM
*
* Demonstrates the dispatcher working with gfx950-specific kernels:
*
* - fp16 GEMM with standard tile configs
* - fp8 GEMM with gfx950-extended warp tiles (16x16x128)
* - 160KB LDS: gfx950 doubles the LDS from 64KB to 160KB
*
* Build: cd dispatcher/build && cmake .. -DGPU_TARGETS=gfx950 && make gemm_07_gfx950_minimal
*/
#include <hip/hip_runtime.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include "ck_tile/dispatcher.hpp"
#include "ck_tile/dispatcher/kernel_decl.hpp"
#include "ck_tile/dispatcher/example_args.hpp"
using namespace ck_tile::dispatcher;
using namespace ck_tile::dispatcher::backends;
using namespace ck_tile::dispatcher::utils;
using Signature = decl::Signature;
using Algorithm = decl::Algorithm;
// =============================================================================
// gfx950-targeted kernel declarations
// =============================================================================
DECL_KERNEL_SET(gfx950_gemm_kernels,
// fp16 128x128x32 -- bread-and-butter config, works on all CDNA
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 128, 32)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx950")
// fp16 128x128x64 -- deeper K tile using more LDS
// LDS usage: 128*64*2 + 128*64*2 = 32768 bytes (fits 64KB, gfx950 has 160KB)
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(128, 128, 64)
.wave(2, 2, 1)
.warp(32, 32, 16)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx950")
// fp16 64x64x32 -- small-tile variant for small problems
.add(Signature().dtype("fp16").layout("rcr"),
Algorithm()
.tile(64, 64, 32)
.wave(2, 2, 1)
.warp(16, 16, 32)
.pipeline("compv3")
.scheduler("intrawave")
.epilogue("cshuffle"),
"gfx950"));
// =============================================================================
// MAIN
// =============================================================================
int main(int argc, char* argv[])
{
ExampleArgs args("Example 07: gfx950 Minimal GEMM",
"Demonstrates gfx950 (CDNA4 / MI350) dispatcher");
args.add_flag("--list", "List registered kernels");
args.add_flag("--list-verbose", "List registered kernels with full details");
args.add_option("--M", "1024", "Problem M dimension");
args.add_option("--N", "1024", "Problem N dimension");
args.add_option("--K", "1024", "Problem K dimension");
args.add_option("--arch", "gfx950", "GPU architecture (default: gfx950)");
if(!args.parse(argc, argv))
return 0;
std::string gfx_arch = args.get("--arch", "gfx950");
print_header("Example 07: gfx950 (CDNA4) Minimal GEMM");
// =========================================================================
// Architecture info
// =========================================================================
std::cout << "\ngfx950 (CDNA4 / MI350) highlights:\n";
std::cout << " - 160KB LDS (up from 64KB on gfx942)\n";
std::cout << " - Extended FP8 warp tiles: 16x16x128, 32x32x64\n";
std::cout << " - Packed FP4 support (pk_fp4)\n";
std::cout << " - Same warp configs as gfx942: [1,4,1], [2,2,1], [4,1,1]\n\n";
// =========================================================================
// Register kernels
// =========================================================================
std::cout << "Registering kernels for " << gfx_arch << "...\n";
Registry registry;
registry.set_name("gfx950_gemm");
REGISTER_GENERATED_KERNELS(registry, gfx_arch);
std::cout << " Registered " << registry.size() << " kernel(s)\n";
if(args.has("--list") || args.has("--list-verbose"))
{
std::cout << "\n";
print_registered_kernels(registry, std::cout, args.has("--list-verbose"));
return 0;
}
if(registry.size() == 0)
{
std::cerr << "ERROR: No kernels registered for " << gfx_arch << "!\n";
std::cerr << " Did you build with -DGPU_TARGETS=gfx950?\n";
return 1;
}
// =========================================================================
// Create Dispatcher
// =========================================================================
Dispatcher dispatcher(&registry);
// =========================================================================
// Setup Problem
// =========================================================================
const int M = args.get_int("--M", 1024);
const int N = args.get_int("--N", 1024);
const int K = args.get_int("--K", 1024);
std::cout << "\nProblem: " << M << " x " << N << " x " << K << "\n";
Problem problem(M, N, K);
using DataType = ck_tile::fp16_t;
GpuBuffer<DataType> a_dev(M * K);
GpuBuffer<DataType> b_dev(K * N);
GpuBuffer<DataType> c_dev(M * N);
std::vector<DataType> a_host(M * K, DataType(1.0f));
std::vector<DataType> b_host(K * N, DataType(1.0f));
a_dev.copy_from_host(a_host.data());
b_dev.copy_from_host(b_host.data());
c_dev.zero();
// =========================================================================
// Select and Run
// =========================================================================
auto selected = dispatcher.select_kernel(problem);
if(!selected)
{
std::cerr << "ERROR: No suitable kernel found for " << M << "x" << N << "x" << K << "\n";
return 1;
}
std::cout << " Selected: " << selected->get_name() << "\n";
float time_ms = dispatcher.run(a_dev.get(), b_dev.get(), c_dev.get(), problem, nullptr);
std::cout << " Time: " << std::fixed << std::setprecision(4) << time_ms << " ms\n";
std::cout << " TFLOPS: " << std::setprecision(2) << calculate_tflops(M, N, K, time_ms) << "\n";
// =========================================================================
// Verify
// =========================================================================
std::cout << "\nVerification:\n";
std::vector<DataType> c_host(M * N);
c_dev.copy_to_host(c_host.data());
const float expected = static_cast<float>(K);
int errors = 0;
for(int i = 0; i < std::min(M * N, 1024); ++i)
{
if(std::abs(static_cast<float>(c_host[i]) - expected) > 0.01f * expected + 1.0f)
++errors;
}
bool passed = (errors == 0);
std::cout << " Expected value: " << expected << "\n";
std::cout << " Errors (first 1024 elements): " << errors << "\n";
std::cout << " Status: " << (passed ? "PASS" : "FAIL") << "\n";
print_separator();
return passed ? 0 : 1;
}

18
dispatcher/examples/gemm/cpp/README.md
View File

@@ -29,14 +29,14 @@ cd examples
## Examples
| Example | Description | Complexity |
|---------|-------------|------------|
| [01_basic_gemm.cpp](01_basic_gemm.cpp) | Basic GEMM with declarative API, autofill, autocorrect | ★☆☆☆☆ |
| [02_multi_size.cpp](02_multi_size.cpp) | Wildcard expansion for multiple configurations | ★★☆☆☆ |
| [03_benchmark_validation.cpp](03_benchmark_validation.cpp) | Performance benchmarking with CPU reference validation | ★★☆☆☆ |
| [04_heuristics.cpp](04_heuristics.cpp) | Heuristic-based kernel selection | ★★★☆☆ |
| [05_json_export.cpp](05_json_export.cpp) | Registry JSON export for external tools | ★★☆☆☆ |
| [06_multi_registry.cpp](06_multi_registry.cpp) | Multiple registries with named kernel sets | ★★★☆☆ |
| Example | Description |
|---------|-------------|
| [01_basic_gemm.cpp](01_basic_gemm.cpp) | Basic GEMM with declarative API, autofill, autocorrect |
| [02_multi_size.cpp](02_multi_size.cpp) | Wildcard expansion for multiple configurations |
| [03_benchmark_validation.cpp](03_benchmark_validation.cpp) | Performance benchmarking with CPU reference validation |
| [04_heuristics.cpp](04_heuristics.cpp) | Heuristic-based kernel selection |
| [05_json_export.cpp](05_json_export.cpp) | Registry JSON export for external tools |
| [06_multi_registry.cpp](06_multi_registry.cpp) | Multiple registries with named kernel sets |
## Example Details
@@ -225,5 +225,5 @@ DECL_KERNEL_SET(my_kernels,
## Related Documentation
- [Python GEMM Examples](../python/README.md)
- [Convolution Examples](../../conv/cpp/README.md)
- [C++ Headers (GEMM + Grouped Conv)](../../../include/ck_tile/dispatcher/README.md)
- [Main Dispatcher README](../../../README.md)