Fix DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3 factory and compute types for input and output tensor in bwd weigth convs.

This commit is contained in:
Ville Pietilä
2026-01-05 03:09:22 -05:00
parent 4eea42c5b7
commit 1dcea1825f
14 changed files with 35 additions and 37 deletions

View File

@@ -255,7 +255,7 @@ concept SpecifiesBlockGemm = requires {
};
template <typename T>
concept SpecifiedGridwiseGemmPipeline = requires
concept SpecifiesGridwiseGemmPipeline = requires
{
{ T::pipeline_version } -> std::convertible_to<PipelineVersion>;
};

View File

@@ -1193,7 +1193,7 @@ consteval auto detailed_diagnostic_SpecifiesLdsTransfer() -> std::string {
}
template <typename T>
consteval auto detailed_diagnostic_SpecifiedGridwiseGemmPipeline() -> std::string {
consteval auto detailed_diagnostic_SpecifiesGridwiseGemmPipeline() -> std::string {
if constexpr (requires { T::pipeline_version; }) {
using PipelineType = decltype(T::pipeline_version);
constexpr bool convertible = std::convertible_to<PipelineType, PipelineVersion>;

View File

@@ -153,7 +153,7 @@ struct FwdWmmaAlgorithm {
CHECK_CONCEPT(T, SpecifiesGemmSpecialization)
CHECK_CONCEPT(T, SpecifiesNumPrefetchStages)
CHECK_CONCEPT(T, SpecifiesLoopScheduler)
CHECK_CONCEPT(T, SpecifiedGridwiseGemmPipeline)
CHECK_CONCEPT(T, SpecifiesGridwiseGemmPipeline)
static constexpr bool c1 = c_ConvAlgorithmDescriptor;
static constexpr bool c2 = c_SpecifiesThreadBlock;
@@ -166,7 +166,7 @@ struct FwdWmmaAlgorithm {
static constexpr bool c9 = c_SpecifiesGemmSpecialization;
static constexpr bool c10 = c_SpecifiesNumPrefetchStages;
static constexpr bool c11 = c_SpecifiesLoopScheduler;
static constexpr bool c12 = c_SpecifiedGridwiseGemmPipeline;
static constexpr bool c12 = c_SpecifiesGridwiseGemmPipeline;
static consteval bool is_valid() {
return c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12;
@@ -186,7 +186,7 @@ struct FwdWmmaAlgorithm {
DIAGNOSTIC_LINE(SpecifiesGemmSpecialization) +
DIAGNOSTIC_LINE(SpecifiesNumPrefetchStages) +
DIAGNOSTIC_LINE(SpecifiesLoopScheduler) +
DIAGNOSTIC_LINE(SpecifiedGridwiseGemmPipeline);
DIAGNOSTIC_LINE(SpecifiesGridwiseGemmPipeline);
}
};
@@ -479,12 +479,12 @@ template <typename T>
struct BwdWmmaAlgorithm : public BwdWmmaAlgorithmBase<T> {
CHECK_CONCEPT(T, SpecifiesNumPrefetchStages)
CHECK_CONCEPT(T, SpecifiesLoopScheduler)
CHECK_CONCEPT(T, SpecifiedGridwiseGemmPipeline)
CHECK_CONCEPT(T, SpecifiesGridwiseGemmPipeline)
CHECK_CONCEPT(T, SpecifiesGenericInstance)
static constexpr bool c9 = c_SpecifiesNumPrefetchStages;
static constexpr bool c10 = c_SpecifiesLoopScheduler;
static constexpr bool c11 = c_SpecifiedGridwiseGemmPipeline;
static constexpr bool c11 = c_SpecifiesGridwiseGemmPipeline;
static constexpr bool c12 = c_SpecifiesGenericInstance;
static consteval bool is_valid() {
@@ -497,7 +497,7 @@ struct BwdWmmaAlgorithm : public BwdWmmaAlgorithmBase<T> {
BwdWmmaAlgorithmBase<T>::message() +
DIAGNOSTIC_LINE(SpecifiesNumPrefetchStages) +
DIAGNOSTIC_LINE(SpecifiesLoopScheduler) +
DIAGNOSTIC_LINE(SpecifiedGridwiseGemmPipeline) +
DIAGNOSTIC_LINE(SpecifiesGridwiseGemmPipeline) +
DIAGNOSTIC_LINE(SpecifiesGenericInstance);
}
};

View File

@@ -35,10 +35,6 @@ struct ConvBwdWeightMultiDWmmaV3Factory
static constexpr auto BLOCK = internal::SetThreadBlockInfo<ALGORITHM>();
static constexpr auto GRIDWISE_GEMM = ALGORITHM.gridwise_gemm;
static constexpr auto GRIDWISE_GEMM_PIPELINE_VERSION =
internal::SetGridwiseGemmPipelineVersion<ALGORITHM>();
static constexpr auto LOOP_SCHEDULER = internal::SetLoopScheduler<ALGORITHM>();
static constexpr auto A_BLOCK_TRANSFER =
internal::SetBwdConvBlockTransfer<ALGORITHM.transfer.a>();
static constexpr auto B_BLOCK_TRANSFER =
@@ -51,10 +47,10 @@ struct ConvBwdWeightMultiDWmmaV3Factory
static_assert(InputVectorTransferLimits<A_BLOCK_TRANSFER>, "Invalid A block transfer config");
static_assert(InputVectorTransferLimits<B_BLOCK_TRANSFER>, "Invalid B block transfer config");
static_assert(OutputVectorTransferLimits<C_BLOCK_TRANSFER>, "Invalid C block transfer config");
static_assert(AccessOrderLimits4D<A_BLOCK_TRANSFER.thread_cluster_order>, "Invalid A thread cluster access order");
static_assert(AccessOrderLimits4D<B_BLOCK_TRANSFER.thread_cluster_order>, "Invalid B thread cluster access order");
static_assert(AccessOrderLimits4D<A_BLOCK_TRANSFER.src_access_order>, "Invalid A source access order");
static_assert(AccessOrderLimits4D<B_BLOCK_TRANSFER.src_access_order>, "Invalid B source access order");
static_assert(AccessOrderLimits3D<A_BLOCK_TRANSFER.thread_cluster_order>, "Invalid A thread cluster access order");
static_assert(AccessOrderLimits3D<B_BLOCK_TRANSFER.thread_cluster_order>, "Invalid B thread cluster access order");
static_assert(AccessOrderLimits3D<A_BLOCK_TRANSFER.src_access_order>, "Invalid A source access order");
static_assert(AccessOrderLimits3D<B_BLOCK_TRANSFER.src_access_order>, "Invalid B source access order");
// The forward convolution kernel class instance.
using Instance = ck::tensor_operation::device::DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3<
@@ -100,7 +96,9 @@ struct ConvBwdWeightMultiDWmmaV3Factory
to_sequence_v<C_BLOCK_TRANSFER.thread_cluster_dims>,
C_BLOCK_TRANSFER.scalar_per_vector,
BLOCK_GEMM.scheduler,
BLOCK_GEMM.pipeline_version>;
BLOCK_GEMM.pipeline_version,
typename Types::OutComputeType,
typename Types::InComputeType>;
};
} // namespace ck_tile::builder::factory

View File

@@ -95,8 +95,8 @@ struct ConvBwdWeightMultiDXdlFactory
C_BLOCK_TRANSFER.n_xdl_per_wave_per_shuffle,
to_sequence_v<C_BLOCK_TRANSFER.thread_cluster_dims>,
C_BLOCK_TRANSFER.scalar_per_vector,
typename Types::InComputeType,
typename Types::WeiComputeType>;
typename Types::OutComputeType,
typename Types::InComputeType>;
};
} // namespace ck_tile::builder::factory

View File

@@ -96,8 +96,8 @@ struct ConvBwdWeightTwoStageWmmaV3Factory
BLOCK_GEMM.scheduler,
BLOCK_GEMM.pipeline_version,
ALGORITHM.num_conv_groups_to_merge,
typename Types::OutComputeType,
typename Types::InComputeType,
typename Types::WeiComputeType,
ALGORITHM.max_transpose_transfer_src_scalar_per_vector,
ALGORITHM.max_transpose_transfer_dst_scalar_per_vector>;
};

View File

@@ -97,8 +97,8 @@ struct ConvBwdWeightTwoStageXdlFactory
BLOCK_GEMM.scheduler,
BLOCK_GEMM.pipeline_version,
ALGORITHM.num_conv_groups_to_merge,
typename Types::OutComputeType,
typename Types::InComputeType,
typename Types::WeiComputeType,
ALGORITHM.max_transpose_transfer_src_scalar_per_vector,
ALGORITHM.max_transpose_transfer_dst_scalar_per_vector>;
};

View File

@@ -95,8 +95,8 @@ struct ConvBwdWeightWmmaV3Factory
C_BLOCK_TRANSFER.scalar_per_vector,
BLOCK_GEMM.scheduler,
BLOCK_GEMM.pipeline_version,
typename Types::OutComputeType,
typename Types::InComputeType,
typename Types::WeiComputeType,
ALGORITHM.max_transpose_transfer_src_scalar_per_vector,
ALGORITHM.max_transpose_transfer_dst_scalar_per_vector>;
};

View File

@@ -93,8 +93,8 @@ struct ConvBwdWeightXdlFactory
C_BLOCK_TRANSFER.n_xdl_per_wave_per_shuffle,
to_sequence_v<C_BLOCK_TRANSFER.thread_cluster_dims>,
C_BLOCK_TRANSFER.scalar_per_vector,
typename Types::OutComputeType,
typename Types::InComputeType,
typename Types::WeiComputeType,
ALGORITHM.max_transpose_transfer_src_scalar_per_vector,
ALGORITHM.max_transpose_transfer_dst_scalar_per_vector>;
};

View File

@@ -96,8 +96,8 @@ struct ConvBwdWeightXdlV3Factory
C_BLOCK_TRANSFER.scalar_per_vector,
BLOCK_GEMM.scheduler,
BLOCK_GEMM.pipeline_version,
typename Types::InComputeType,
typename Types::WeiComputeType>;
typename Types::OutComputeType,
typename Types::InComputeType>;
};
} // namespace ck_tile::builder::factory

View File

@@ -178,7 +178,7 @@ if (CK_USE_WMMA)
conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle_v3.cpp
conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle.cpp
conv/ck/test_ckb_conv_bwd_weight_two_stage_wmma_cshuffle_v3.cpp
conv/ck/test_ckb_conv_bwd_weight_multi_d_wmma_cshuffle.cpp
conv/ck/test_ckb_conv_bwd_weight_multi_d_wmma_cshuffle_v3.cpp
)
endif()

View File

@@ -11,15 +11,15 @@ namespace ckt = ck_tile::builder::test;
namespace cku = ck_tile::builder::test_utils;
constexpr auto SIGNATURE =
ckt::ConvSignature{.spatial_dim = 2,
ckt::ConvSignature{.spatial_dim = 3,
.direction = ckb::ConvDirection::BACKWARD_WEIGHT,
.data_type = ckb::DataType::FP16,
.accumulation_data_type = ckb::DataType::FP32,
.input = {.config = {.layout = ckb::TensorLayout::GNHWC}},
.weight = {.config = {.layout = ckb::TensorLayout::GKYXC}},
.output = {.config = {.layout = ckb::TensorLayout::GNHWK}}};
.input = {.config = {.layout = ckb::TensorLayout::GNDHWC}},
.weight = {.config = {.layout = ckb::TensorLayout::GKZYXC}},
.output = {.config = {.layout = ckb::TensorLayout::GNDHWK}}};
constexpr auto ALGORITHM = cku::ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle{}
constexpr auto ALGORITHM = cku::ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle_V3{}
.with_thread_block(cku::ThreadBlock_64_32x32x32)
.with_gemm_config(cku::GemmParams_Wmma_16x16_2x1_per_wave)
.with_transfer(cku::BwdTransfer_4x8x1_4x16x1_v3)
@@ -29,14 +29,14 @@ constexpr auto ALGORITHM = cku::ConvAlgorithm_DeviceGroupedConvBwdWeightMultiple
using Builder = ckb::ConvBuilder<SIGNATURE, ALGORITHM>;
using Instance = Builder::Instance;
TEST(BwdWeight_2DFp16_MultiD_Wmma_Shuffle_GNHWC, Create)
TEST(BwdWeight_3DFp16_MultiD_Wmma_ShuffleV3_GNHWC, Create)
{
const auto expected_transfer_parameters = to_string(ALGORITHM);
std::cout << "Expected Transfer Parameters: " << expected_transfer_parameters << std::endl;
cku::run_test<Builder>({"DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle",
cku::run_test<Builder>({"DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3",
expected_transfer_parameters,
"Default",
"GNHWC,GKYXC,GNHWK",
"GNDHWC,GKZYXC,GNDHWK",
"PassThrough,PassThrough,PassThrough",
"fp16,fp16>"}); // check compute types
}

View File

@@ -609,7 +609,7 @@ using ConvAlgorithm_DeviceGroupedConvBwdWeight_TwoStage_Wmma_CShuffle_V3 =
using ConvAlgorithm_DeviceGroupedConvBwdWeight_Wmma_CShuffle =
ConvAlgorithmTemplate<ThreadBlock_, WmmaGemm_, Transfer_<4>, ConvSpecializationBwdWeight_, GridGemm_, Prefetch_>;
using ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle =
using ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle_V3 =
ConvAlgorithmTemplate<ThreadBlock_, WmmaGemm_, Transfer_<>, ConvSpecializationBwdWeight_, BlockGemm_, MultipleDSpecialization_>;

View File

@@ -444,8 +444,8 @@ inline std::string to_string<ConvAlgorithm_DeviceGroupedConvBwdWeight_Wmma_CShuf
}
template <>
inline std::string to_string<ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle>(
ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle t)
inline std::string to_string<ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle_V3>(
ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle_V3 t)
{
std::ostringstream oss;
oss << to_string(static_cast<ThreadBlock_>(t)) << "," << to_string(static_cast<WmmaGemm_>(t))