ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make BlockTransferDescriptor concept parametrized. Introduce a common TileTransferParameters concept for conv algorithms.

Browse Source
This commit is contained in:
Ville Pietilä
2026-01-13 04:17:06 -05:00
parent 1d519792ca
commit a8f1d44078
3 changed files with 37 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
experimental/builder/include/ck_tile/builder/conv_algorithm_concepts.hpp
View File

@@ -56,7 +56,7 @@ concept GridwiseWmmaGemmDescriptor = requires(T t) {
// Concept for vectorized data transfer for convolution input tensors.
template <typename T>
concept BlockTransferDescriptor = requires(T t) {
concept BlockTransferDescriptor3D = requires(T t) {
{ t.k0 } -> SizeType;
{ t.m_n } -> SizeType;
{ t.k1 } -> SizeType;
@@ -70,6 +70,10 @@ concept BlockTransferDescriptor4D = requires(T t) {
{ t.k_batch_size } -> SizeType;
};
template <typename T, size_t ThreadClusterRank>
concept BlockTransferDescriptor = (ThreadClusterRank == 3 && BlockTransferDescriptor3D<T>) ||
(ThreadClusterRank == 4 && BlockTransferDescriptor4D<T>);
// Concept for thread cluster dimensions for GEMM output tensor.
template <typename T>
concept ThreadClusterDescriptor = requires(T t) {
@@ -202,19 +206,10 @@ concept SpecifiesGridwiseWmmaGemm = requires(T t) {
};
// Concept to check if a struct specifies convolution input and output block transfer info.
template <typename T>
template <typename T, size_t ThreadClusterRank = 3>
concept SpecifiesBlockTransfer = requires(T t) {
{ T::transfer.a.block_transfer } -> BlockTransferDescriptor;
{ T::transfer.b.block_transfer } -> BlockTransferDescriptor;
{ T::transfer.c.thread_cluster_dims } -> ThreadClusterDescriptor;
};
// Concept to check if a struct specifies convolution input and output block transfer info
// for 4D thread slices.
template <typename T>
concept SpecifiesBlockTransfer4D = requires(T t) {
{ T::transfer.a.block_transfer } -> BlockTransferDescriptor4D;
{ T::transfer.b.block_transfer } -> BlockTransferDescriptor4D;
{ T::transfer.a.block_transfer } -> BlockTransferDescriptor<ThreadClusterRank>;
{ T::transfer.b.block_transfer } -> BlockTransferDescriptor<ThreadClusterRank>;
{ T::transfer.c.thread_cluster_dims } -> ThreadClusterDescriptor;
};

49
experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp
View File

@@ -8,39 +8,46 @@
namespace ck_tile::builder::factory {
// Base algorithm concepts
template <typename T, size_t ThreadClusterRank = 3>
concept TileTransferParameters =
SpecifiesBlockTransfer<T, ThreadClusterRank> && SpecifiesLdsTransfer<T> &&
SpecifiesThreadClusterAccessOrder<T> && SpecifiesSourceAccessOrder<T>;
template <typename T>
concept SpecifiesTileTransferParameters3D = TileTransferParameters<T, 3>;
template <typename T>
concept SpecifiesTileTransferParameters4D = TileTransferParameters<T, 4>;
template <typename T>
concept FwdXdlAlgorithmBase =
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesBlockTransfer<T> &&
SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
SpecifiesSourceAccessOrder<T> && SpecifiesGridwiseFwdXdlGemm<T> &&
SpecifiesFwdConvSpecialization<T> && SpecifiesGemmSpecialization<T> &&
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesTileTransferParameters3D<T> &&
SpecifiesGridwiseFwdXdlGemm<T> && SpecifiesFwdConvSpecialization<T> && SpecifiesGemmSpecialization<T> &&
SpecifiesNumPrefetchStages<T> && SpecifiesNumGroupsToMerge<T> && SpecifiesLoopScheduler<T>;
template <typename T>
concept BwdXdlAlgorithmBase =
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesBlockTransfer4D<T> &&
SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
SpecifiesSourceAccessOrder<T> && SpecifiesGridwiseBwdXdlGemm<T> &&
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> &&
SpecifiesTileTransferParameters4D<T> && SpecifiesGridwiseBwdXdlGemm<T> &&
SpecifiesBwdWeightConvSpecialization<T>;
template <typename T>
concept BwdXdlV3AlgorithmBase =
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesBlockTransfer<T> &&
SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
SpecifiesSourceAccessOrder<T> && SpecifiesGridwiseBwdXdlGemm<T> &&
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> &&
SpecifiesTileTransferParameters3D<T> && SpecifiesGridwiseBwdXdlGemm<T> &&
SpecifiesBwdWeightConvSpecialization<T> && SpecifiesBlockGemm<T>;
template <typename T>
concept BwdWmmaAlgorithmBase =
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesBlockTransfer<T> &&
SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
SpecifiesSourceAccessOrder<T> && SpecifiesGridwiseWmmaGemm<T> &&
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> &&
SpecifiesTileTransferParameters3D<T> && SpecifiesGridwiseWmmaGemm<T> &&
SpecifiesBwdWeightConvSpecialization<T>;
template <typename T>
concept BwdWmmaV3AlgorithmBase =
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesBlockTransfer<T> &&
SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
SpecifiesSourceAccessOrder<T> && SpecifiesGridwiseWmmaGemm<T> &&
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> &&
SpecifiesTileTransferParameters3D<T> && SpecifiesGridwiseWmmaGemm<T> &&
SpecifiesBwdWeightConvSpecialization<T> && SpecifiesBlockGemm<T>;
// Reference algorithm concept
@@ -62,17 +69,15 @@ concept LargeTensorAlgorithm = FwdXdlAlgorithmBase<T> && SpecifiesLargeTensorSup
template <typename T>
concept FwdXdlV3Algorithm =
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesBlockTransfer<T> &&
SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
SpecifiesSourceAccessOrder<T> && SpecifiesGridwiseFwdXdlGemm<T> &&
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> &&
SpecifiesTileTransferParameters3D<T> && SpecifiesGridwiseFwdXdlGemm<T> &&
SpecifiesFwdConvSpecialization<T> && SpecifiesGemmSpecialization<T> && SpecifiesBlockGemm<T>;
// FWD WMMA algorithm concepts
template <typename T>
concept FwdWmmaAlgorithm =
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesBlockTransfer<T> &&
SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
SpecifiesSourceAccessOrder<T> && SpecifiesGridwiseWmmaGemm<T> &&
ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> &&
SpecifiesTileTransferParameters3D<T> && SpecifiesGridwiseWmmaGemm<T> &&
SpecifiesFwdConvSpecialization<T> && SpecifiesGemmSpecialization<T> &&
SpecifiesNumPrefetchStages<T> && SpecifiesLoopScheduler<T> && SpecifiesGridwiseGemmPipeline<T>;

4
experimental/builder/test/impl/conv_algorithm_types.hpp
View File

@@ -90,8 +90,8 @@ struct BlockTransfer<3>
size_t m_n;
size_t k1;
};
static_assert(ckb::BlockTransferDescriptor<BlockTransfer<>>);
static_assert(ckb::BlockTransferDescriptor<BlockTransfer<4>>);
static_assert(ckb::BlockTransferDescriptor<BlockTransfer<3>, 3>);
static_assert(ckb::BlockTransferDescriptor<BlockTransfer<4>, 4>);
// Describe C block transfer thread cluster lengths.
struct ThreadCluster