mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-12 01:10:17 +00:00

Files

Vidyasagar Ananthan 920acd2c12 [rocm-libraries] ROCm/rocm-libraries#5168 (commit 8b5afcb)

[CK] [CK_Tile] Add GroupConv to Kernel Dispatcher

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

2026-04-09 17:39:35 +00:00

ctypes

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

2026-04-09 17:39:35 +00:00

README.md

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

2026-04-09 17:39:35 +00:00

README.md

CK Tile Dispatcher - Language Bindings

This directory contains language bindings for the CK Tile Dispatcher.

Structure

bindings/
|---- ctypes/              # Python ctypes bindings (C API)
|   |---- gemm_ctypes_lib.cpp      # GEMM dispatcher C API
|   |---- conv_ctypes_lib.cpp      # Grouped conv dispatcher C API (fwd + bwd_data)
|   |---- conv_bwdw_ctypes_lib.cpp # Grouped conv backward weight C API (separate library)
|   |---- gpu_helper.cpp           # CLI helper for Python
|   +---- CMakeLists.txt
+---- README.md

ctypes Bindings

The ctypes bindings provide a C API that Python can load via ctypes.CDLL().

Building

cd build
cmake .. -DCMAKE_PREFIX_PATH=/opt/rocm
make dispatcher_gemm_lib dispatcher_conv_lib gpu_helper

Usage from Python

import ctypes

# Load the library
lib = ctypes.CDLL("path/to/libdispatcher_gemm_lib.so")

# Initialize
lib.dispatcher_init()

# Check if problem is supported
is_supported = lib.dispatcher_is_supported(M, N, K)

# Run GEMM
time_ms = ctypes.c_float()
result = lib.dispatcher_run_gemm(
    A_ptr, B_ptr, C_ptr,
    M, N, K,
    ctypes.byref(time_ms)
)

# Cleanup
lib.dispatcher_cleanup()

GEMM API

Function	Description
`dispatcher_init()`	Initialize the dispatcher
`dispatcher_is_supported(M, N, K)`	Check if problem size is supported
`dispatcher_select_kernel(M, N, K, name_buf, buf_size)`	Get kernel name for problem
`dispatcher_run_gemm(A, B, C, M, N, K, time_ms)`	Execute GEMM
`dispatcher_get_kernel_count()`	Get number of registered kernels
`dispatcher_export_registry_json()`	Export registry as JSON
`dispatcher_cleanup()`	Release resources

Grouped Convolution API

Function	Description
`conv_dispatcher_init()`	Initialize the dispatcher
`conv_dispatcher_is_supported(prob)`	Check if problem is supported
`conv_dispatcher_select_kernel(prob, name_buf, buf_size)`	Get kernel name
`conv_dispatcher_run(input, weight, output, prob, stream)`	Execute convolution
`conv_dispatcher_get_kernel_count()`	Get number of registered kernels
`conv_dispatcher_cleanup()`	Release resources

GPU Helper

The gpu_helper executable provides a CLI interface for Python:

./gpu_helper 1024 1024 1024 --validate

Output is JSON for easy parsing:

{
  "problem": {"M": 1024, "N": 1024, "K": 1024},
  "kernel": "gemm_fp16_rcr_...",
  "execution": {
    "time_ms": 0.5,
    "tflops": 4.2
  },
  "validation": {
    "accuracy": 100.0
  },
  "status": "success"
}

Examples

See the examples that use these bindings:

GEMM: dispatcher/examples/gemm/python/

Grouped Convolution

Grouped convolution C++ headers and Python utilities are in:

C++ Headers: dispatcher/include/ck_tile/dispatcher/grouped_conv_*.hpp
Python Utils: dispatcher/python/grouped_conv_utils.py
Build Script: dispatcher/scripts/compile_grouped_conv_examples.py