ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-21 05:19: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

181
dispatcher/src/registry.cpp
View File

@@ -5,39 +5,32 @@
#include "ck_tile/dispatcher/json_export.hpp"
#include "ck_tile/dispatcher/arch_filter.hpp"
#include <algorithm>
#include <fstream>
#include <iostream>
namespace ck_tile {
namespace dispatcher {
Registry::Registry()
: name_("default"),
auto_export_enabled_(false),
auto_export_include_statistics_(true),
auto_export_on_every_registration_(true)
{
}
Registry::Registry() = default;
Registry::~Registry()
{
// Perform auto-export on destruction if enabled (regardless of export_on_every_registration
// setting)
if(auto_export_enabled_)
{
perform_auto_export();
}
}
Registry::Registry(Registry&& other) noexcept
: mutex_() // mutex is not movable, create new one
,
kernels_(std::move(other.kernels_)),
name_(std::move(other.name_)),
auto_export_enabled_(other.auto_export_enabled_),
auto_export_filename_(std::move(other.auto_export_filename_)),
auto_export_include_statistics_(other.auto_export_include_statistics_),
auto_export_on_every_registration_(other.auto_export_on_every_registration_)
Registry::Registry(Registry&& other) noexcept : Base(std::move(other))
{
// Disable auto-export on the moved-from object to prevent double export
// Base move constructor already locked+released other.mutex_.
// Re-acquire to safely read the remaining fields.
std::lock_guard<std::mutex> lock(other.mutex());
auto_export_enabled_ = other.auto_export_enabled_;
auto_export_filename_ = std::move(other.auto_export_filename_);
auto_export_include_statistics_ = other.auto_export_include_statistics_;
auto_export_on_every_registration_ = other.auto_export_on_every_registration_;
other.auto_export_enabled_ = false;
}
@@ -45,11 +38,7 @@ Registry& Registry::operator=(Registry&& other) noexcept
{
if(this != &other)
{
std::lock_guard<std::mutex> lock(mutex_);
std::lock_guard<std::mutex> other_lock(other.mutex_);
kernels_ = std::move(other.kernels_);
name_ = std::move(other.name_);
Base::operator=(std::move(other));
auto_export_enabled_ = other.auto_export_enabled_;
auto_export_filename_ = std::move(other.auto_export_filename_);
auto_export_include_statistics_ = other.auto_export_include_statistics_;
@@ -64,55 +53,27 @@ Registry& Registry::operator=(Registry&& other) noexcept
bool Registry::register_kernel(KernelInstancePtr instance, Priority priority)
{
if(!instance)
{
return false;
}
const std::string identifier = instance->get_key().encode_identifier();
bool registered = false;
if(Base::register_kernel(instance->get_name(), instance, priority))
{
std::lock_guard<std::mutex> lock(mutex_);
auto it = kernels_.find(identifier);
if(it != kernels_.end())
if(auto_export_enabled_ && auto_export_on_every_registration_)
{
// Kernel with this identifier already exists
// Only replace if new priority is higher
if(priority > it->second.priority)
{
it->second.instance = instance;
it->second.priority = priority;
registered = true;
}
}
else
{
// New kernel, insert it
kernels_[identifier] = RegistryEntry{instance, priority};
registered = true;
perform_auto_export();
}
return true;
}
// Perform auto-export if enabled and configured to export on every registration
if(registered && auto_export_enabled_ && auto_export_on_every_registration_)
{
perform_auto_export();
}
return registered;
return false;
}
KernelInstancePtr Registry::lookup(const std::string& identifier) const
{
std::lock_guard<std::mutex> lock(mutex_);
auto it = kernels_.find(identifier);
if(it != kernels_.end())
std::lock_guard<std::mutex> lock(mutex());
auto it = entries().find(identifier);
if(it != entries().end())
{
return it->second.instance;
}
return nullptr;
}
@@ -121,75 +82,23 @@ KernelInstancePtr Registry::lookup(const KernelKey& key) const
return lookup(key.encode_identifier());
}
std::vector<KernelInstancePtr> Registry::get_all() const
{
std::lock_guard<std::mutex> lock(mutex_);
std::vector<KernelInstancePtr> result;
result.reserve(kernels_.size());
for(const auto& pair : kernels_)
{
result.push_back(pair.second.instance);
}
return result;
}
std::vector<KernelInstancePtr> Registry::get_all() const { return Base::get_all_instances(); }
std::vector<KernelInstancePtr>
Registry::filter(std::function<bool(const KernelInstance&)> predicate) const
{
std::lock_guard<std::mutex> lock(mutex_);
std::lock_guard<std::mutex> lock(mutex());
std::vector<KernelInstancePtr> result;
for(const auto& pair : kernels_)
for(const auto& [name, entry] : entries())
{
if(predicate(*pair.second.instance))
if(predicate(*(entry.instance)))
{
result.push_back(pair.second.instance);
result.push_back(entry.instance);
}
}
return result;
}
std::size_t Registry::size() const
{
std::lock_guard<std::mutex> lock(mutex_);
return kernels_.size();
}
bool Registry::empty() const
{
std::lock_guard<std::mutex> lock(mutex_);
return kernels_.empty();
}
void Registry::clear()
{
std::lock_guard<std::mutex> lock(mutex_);
kernels_.clear();
}
const std::string& Registry::get_name() const
{
std::lock_guard<std::mutex> lock(mutex_);
return name_;
}
void Registry::set_name(const std::string& name)
{
std::lock_guard<std::mutex> lock(mutex_);
name_ = name;
}
Registry& Registry::instance()
{
static Registry global_registry;
return global_registry;
}
std::string Registry::export_json(bool include_statistics) const
{
return export_registry_json(*this, include_statistics);
@@ -204,7 +113,7 @@ void Registry::enable_auto_export(const std::string& filename,
bool include_statistics,
bool export_on_every_registration)
{
std::lock_guard<std::mutex> lock(mutex_);
std::lock_guard<std::mutex> lock(mutex());
auto_export_enabled_ = true;
auto_export_filename_ = filename;
auto_export_include_statistics_ = include_statistics;
@@ -213,13 +122,13 @@ void Registry::enable_auto_export(const std::string& filename,
void Registry::disable_auto_export()
{
std::lock_guard<std::mutex> lock(mutex_);
std::lock_guard<std::mutex> lock(mutex());
auto_export_enabled_ = false;
}
bool Registry::is_auto_export_enabled() const
{
std::lock_guard<std::mutex> lock(mutex_);
std::lock_guard<std::mutex> lock(mutex());
return auto_export_enabled_;
}
@@ -230,7 +139,7 @@ void Registry::perform_auto_export()
bool include_stats;
{
std::lock_guard<std::mutex> lock(mutex_);
std::lock_guard<std::mutex> lock(mutex());
if(!auto_export_enabled_)
{
return;
@@ -243,31 +152,15 @@ void Registry::perform_auto_export()
export_json_to_file(filename, include_stats);
}
std::size_t Registry::merge_from(const Registry& other, Priority priority)
{
auto other_kernels = other.get_all();
std::size_t merged_count = 0;
for(const auto& kernel : other_kernels)
{
if(register_kernel(kernel, priority))
{
merged_count++;
}
}
return merged_count;
}
std::size_t Registry::filter_by_arch(const std::string& gpu_arch)
{
ArchFilter filter(gpu_arch);
std::vector<std::string> to_remove;
{
std::lock_guard<std::mutex> lock(mutex_);
std::lock_guard<std::mutex> lock(mutex());
for(const auto& pair : kernels_)
for(const auto& pair : entries())
{
if(!filter.is_valid(pair.second.instance->get_key()))
{
@@ -277,12 +170,18 @@ std::size_t Registry::filter_by_arch(const std::string& gpu_arch)
for(const auto& key : to_remove)
{
kernels_.erase(key);
entries_mut().erase(key);
}
}
return to_remove.size();
}
Registry& Registry::instance()
{
static Registry global_registry;
return global_registry;
}
} // namespace dispatcher
} // namespace ck_tile
} // namespace ck_tile