From 5878c32c147eaec6c3ccbddab4e3b1547be3af5c Mon Sep 17 00:00:00 2001 From: John Shumway Date: Thu, 18 Sep 2025 05:12:20 +0000 Subject: [PATCH] Clean up factory for backwards convolutions. Adds some comments and simplifies the block transfer settings by splitting into more funcctions. --- .../include/ck_tile/builder/conv_factory.hpp | 134 ++++++++++-------- 1 file changed, 76 insertions(+), 58 deletions(-) diff --git a/experimental/builder/include/ck_tile/builder/conv_factory.hpp b/experimental/builder/include/ck_tile/builder/conv_factory.hpp index c9d5658de4..264604ef5c 100644 --- a/experimental/builder/include/ck_tile/builder/conv_factory.hpp +++ b/experimental/builder/include/ck_tile/builder/conv_factory.hpp @@ -1,9 +1,37 @@ #pragma once -// #include -// "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_comp_instance.hpp" +// A factory for instantiating CK convolution kernels. +// +// This file translates a semantic description of a convolution operation +// (`ConvSignatureDescriptor` and `ConvAlgorithmDescriptor`) into specific, +// low-level template arguments required by the underlying CK device-level +// kernel implementations. This abstraction also enables more complex build +// time logic and simplifies the kernel specificatoin. +// +// Key Components: +// +// Template Metaprogram: +// - ConvFactory: The main factory, with specializations for different +// convolution directions. +// +// Template Metaprogram Helpers: +// - ConvTensorLayouts: Maps layout enums to CK layout types. +// - ConvTensorTypes: Maps data type enums to C++ types used by CK. +// - ConvPassThroughOps: Hard-coded pass-through element-wise operations. +// +// `constexpr` Helper Functions: +// - SetThreadBlockInfo: Determines thread block dimensions from the algorithm +// descriptor or provides defaults. +// - SetConvTuningInfo: Sets low-level tuning parameters. +// - Set*BlockTransfer: Configures tensor data movement parameters for +// tensors A, B, and C. +// - SetBlockGemmPipelineVersion: Selects the GEMM pipeline version. +// +// The primary entry point is the `ConvFactory` struct, which is specialized +// for forward and backward-data convolutions. + #include -#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp" +#include #include #include #include @@ -26,7 +54,6 @@ struct ConvTensorLayouts template <> struct ConvTensorLayouts { - // Channels first convolution layout. using ALayout = ck::tensor_layout::convolution::NHWGC; using BLayout = ck::tensor_layout::convolution::GKCYX; using DsLayout = ck::Tuple<>; @@ -36,7 +63,6 @@ struct ConvTensorLayouts struct ConvTensorLayouts { - // Channels last convolution layout. using ALayout = ck::tensor_layout::convolution::NGKHW; using BLayout = ck::tensor_layout::convolution::GKYXC; using DsLayout = ck::Tuple<>; @@ -46,7 +72,6 @@ struct ConvTensorLayouts struct ConvTensorLayouts { - // Channels last convolution layout. using ALayout = ck::tensor_layout::convolution::NHWGC; using BLayout = ck::tensor_layout::convolution::GKYXC; using DsLayout = ck::Tuple<>; @@ -56,7 +81,6 @@ struct ConvTensorLayouts struct ConvTensorLayouts { - // Channels last convolution layout. using ALayout = ck::tensor_layout::convolution::NDHWGC; using BLayout = ck::tensor_layout::convolution::GKZYXC; using DsLayout = ck::Tuple<>; @@ -107,6 +131,7 @@ struct ConvTensorTypes }; // Hard-coded pass-through ops. +// TODO: Generalize this for more fused operations. struct ConvPassThroughOps { using AElementwiseOp = ck::tensor_operation::element_wise::PassThrough; @@ -126,7 +151,7 @@ struct ConvSpec ck::tensor_operation::device::GemmSpecialization gemm_spec; }; -// Deduction guide for ConvSpec simplifies brace initialization. +// Deduction guide for ConvSpec to simplify brace initialization. template ConvSpec(CONV_ENUM, GEMM_ENUM) -> ConvSpec; @@ -226,23 +251,10 @@ struct CBlockTransfer int scaler_per_vector = 8; }; -template -constexpr BlockTransfer SetABlockTransfer() +template +constexpr BlockTransfer SetFwdConvABlockTransfer() { using AlgorithmType = decltype(ALGORITHM); - if constexpr(ConvDirectionIsBackwardData) - { - // Different default values for backward data. - return BlockTransfer{ - .thread_cluster_dims = {4, 16, 1}, - .thread_cluster_order = {1, 0, 2}, - .src_access_order = {1, 0, 2}, - .src_vector_dim = 2, - .src_scaler_per_vector = 8, - .dest_scaler_per_vector_k1 = 8, - .add_extra = 1, - }; - } BlockTransfer block_transfer{ .thread_cluster_dims = {4, 64, 1}, .thread_cluster_order = {1, 0, 2}, @@ -261,23 +273,24 @@ constexpr BlockTransfer SetABlockTransfer() return block_transfer; } -template -constexpr BlockTransfer SetBBlockTransfer() +template +constexpr BlockTransfer SetBwdDataConvABlockTransfer() +{ + return BlockTransfer{ + .thread_cluster_dims = {4, 16, 1}, + .thread_cluster_order = {1, 0, 2}, + .src_access_order = {1, 0, 2}, + .src_vector_dim = 2, + .src_scaler_per_vector = 8, + .dest_scaler_per_vector_k1 = 8, + .add_extra = 1, + }; +} + +template +constexpr BlockTransfer SetFwdConvBBlockTransfer() { using AlgorithmType = decltype(ALGORITHM); - if constexpr(ConvDirectionIsBackwardData) - { - // Different default values for backward data. - return BlockTransfer{ - .thread_cluster_dims = {4, 8, 1}, - .thread_cluster_order = {0, 2, 1}, - .src_access_order = {0, 2, 1}, - .src_vector_dim = 1, - .src_scaler_per_vector = 8, - .dest_scaler_per_vector_k1 = 8, - .add_extra = 1, - }; - } BlockTransfer block_transfer{ .thread_cluster_dims = {4, 64, 1}, .thread_cluster_order = {1, 0, 2}, @@ -292,10 +305,24 @@ constexpr BlockTransfer SetBBlockTransfer() constexpr auto& TCL = ALGORITHM.block_transfer.thread_cluster_dims_b; block_transfer.thread_cluster_dims = {TCL.k0, TCL.n, TCL.k1}; } - // Default. return block_transfer; } +template +constexpr BlockTransfer SetBwdDataConvBBlockTransfer() +{ + // Different default values for backward data. + return BlockTransfer{ + .thread_cluster_dims = {4, 8, 1}, + .thread_cluster_order = {0, 2, 1}, + .src_access_order = {0, 2, 1}, + .src_vector_dim = 1, + .src_scaler_per_vector = 8, + .dest_scaler_per_vector_k1 = 8, + .add_extra = 1, + }; +} + template constexpr CBlockTransfer SetCBlockTransfer() { @@ -353,11 +380,11 @@ template struct ConvFactory; -// Factory builds an instance of a grouped forward convolution kernel. +// Factory specialization for an instance of a grouped forward convolution kernel. template - requires SupportedVersion && ConvDirectionIsForward + StringLiteral VERSION> + requires ConvDirectionIsForward struct ConvFactory { static constexpr int SPATIAL_DIM = SIGNATURE.spatial_dim; @@ -370,12 +397,12 @@ struct ConvFactory }; static constexpr ConvBlock BLOCK = SetThreadBlockInfo(); static constexpr ConvTuning TUNING = SetConvTuningInfo(); - static constexpr BlockTransfer A_BLOCK_TRANSFER = SetABlockTransfer(); - static constexpr BlockTransfer B_BLOCK_TRANSFER = SetBBlockTransfer(); + static constexpr BlockTransfer A_BLOCK_TRANSFER = SetFwdConvABlockTransfer(); + static constexpr BlockTransfer B_BLOCK_TRANSFER = SetFwdConvBBlockTransfer(); static constexpr CBlockTransfer C_BLOCK_TRANSFER = SetCBlockTransfer(); static constexpr auto PIPELINE_SCHEDULER = ck::BlockGemmPipelineScheduler::Intrawave; static constexpr auto PIPELINE_VERSION = SetBlockGemmPipelineVersion(); - // The convlution kernel class instance. + // The forward convolution kernel class instance. using Instance = ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3< // SPATIAL_DIM, @@ -426,16 +453,7 @@ struct ConvFactory PIPELINE_VERSION>; }; -// clang-format off - // ##############################################| NDim| ALayout| BLayout| DsLayout| ELayout| AData| BData| AccData| CShuffle| DsData| EData| AElementwise| BElementwise| CDEElementwise| ConvolutionBackward| DoPad| DoPad| NumGemmK| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffleMXdl| CShuffleNXdl| CDEBlockTransfer| CDEBlockTransfer| - // ##############################################| Spatial| | | | | Type| Type| Type| DataType| Type| Type| Operation| Operation| Operation| DataSpecialization| GemmM| GemmN| PrefetchStage| Size| Block| Block| Block| | | XDL| XDL| PerWave| PerWave| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| ExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| ExtraN| PerWave| PerWave| _MBlock_MPerBlock| ScalarPerVector| - // ##############################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | Lengths_AK0_M_AK1| ArrangeOrder| | | PerVector| PerVector_AK1| | Lengths_BK0_N_BK1| ArrangeOrder| | | PerVector| PerVector_BK1| | PerShuffle| PerShuffle| _NBlock_NPerBlock| _NPerBlock| - // ##############################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - // generic instance - // DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1, S<1, 0, 2>, S<1, 0, 2>, 2, 1, 8, 1, S<4, 8, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 1, 1, 1, S<1, 16, 1, 4>, 1> -// clang-format on - -// Factory builds an instance of a grouped backward-data convolution kernel. +// Factory specialization for an instance of a grouped backward-data convolution kernel. template @@ -452,12 +470,12 @@ struct ConvFactory }; static constexpr ConvBlock BLOCK = SetThreadBlockInfo(); static constexpr ConvTuning TUNING = SetConvTuningInfo(); - static constexpr BlockTransfer A_BLOCK_TRANSFER = SetABlockTransfer(); - static constexpr BlockTransfer B_BLOCK_TRANSFER = SetBBlockTransfer(); + static constexpr BlockTransfer A_BLOCK_TRANSFER = SetBwdDataConvABlockTransfer(); + static constexpr BlockTransfer B_BLOCK_TRANSFER = SetBwdDataConvBBlockTransfer(); static constexpr CBlockTransfer C_BLOCK_TRANSFER = SetCBlockTransfer(); static constexpr auto PIPELINE_SCHEDULER = ck::BlockGemmPipelineScheduler::Intrawave; static constexpr auto PIPELINE_VERSION = SetBlockGemmPipelineVersion(); - + // The backward-data convolution kernel class instance. using Instance = ck::tensor_operation::device::DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1< SPATIAL_DIM,