Add factory and tests for DeviceGroupedConvBwdWeight_Wmma_CShuffleV3.

This commit is contained in:
Ville Pietilä
2026-01-02 06:08:29 -05:00
parent d045923a0d
commit bc3cba873b
11 changed files with 268 additions and 45 deletions

View File

@@ -570,6 +570,12 @@ consteval int count_matches_bwd_two_stage_xdl() {
return Alg::c1 + Alg::c2 + Alg::c3 + Alg::c4 + Alg::c5 + Alg::c6 + Alg::c7 + Alg::c8 + Alg::c9 + Alg::c10 + Alg::c11 + Alg::c12;
}
template <typename T>
consteval int count_matches_bwd_wmma_v3() {
using Alg = BwdWmmaV3Algorithm<T>;
return Alg::c1 + Alg::c2 + Alg::c3 + Alg::c4 + Alg::c5 + Alg::c6 + Alg::c7 + Alg::c8 + Alg::c9 + Alg::c10 + Alg::c11;
}
template <typename T>
consteval int count_matches_bwd_dl() {
using Alg = BwdDlAlgorithm<T>;
@@ -599,26 +605,26 @@ consteval void diagnose_fwd_algorithm_signature()
constexpr int large_tensor_matches = count_matches_large_tensor<AlgoType>();
constexpr int tile_matches = count_matches_tile<AlgoType>();
// Find maximum matches across all variants
constexpr int max_1 = xdl_v3_matches > xdl_matches ? xdl_v3_matches : xdl_matches;
constexpr int max_2 = wmma_matches > dl_matches ? wmma_matches : dl_matches;
constexpr int max_3 = max_1 > max_2 ? max_1 : max_2;
constexpr int max_4 = max_3 > large_tensor_matches ? max_3 : large_tensor_matches;
constexpr int max_matches = max_4 > tile_matches ? max_4 : tile_matches;
// Check whether we have XDL or WMMA algorithm
if constexpr (SpecifiesGridwiseFwdXdlGemm<AlgoType>)
{
constexpr int max_1 = xdl_v3_matches > xdl_matches ? xdl_v3_matches : xdl_matches;
constexpr int max_2 = max_1 > dl_matches ? max_1 : dl_matches;
constexpr int max_matches = large_tensor_matches > max_2 ? large_tensor_matches : max_2;
if constexpr(max_matches == xdl_v3_matches) {
using Alg = FwdXdlV3Algorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
} else if constexpr(max_matches == xdl_matches) {
using Alg = FwdXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
} else if constexpr(max_matches == dl_matches) {
using Alg = FwdDlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
} else if constexpr (max_matches == large_tensor_matches) {
using Alg = LargeTensorAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
}
}
else if constexpr (SpecifiesGridwiseWmmaGemm<AlgoType>)
{
@@ -627,6 +633,13 @@ consteval void diagnose_fwd_algorithm_signature()
}
else
{
// Find maximum matches across all variants
constexpr int max_1 = xdl_v3_matches > xdl_matches ? xdl_v3_matches : xdl_matches;
constexpr int max_2 = wmma_matches > dl_matches ? wmma_matches : dl_matches;
constexpr int max_3 = max_1 > max_2 ? max_1 : max_2;
constexpr int max_4 = max_3 > large_tensor_matches ? max_3 : large_tensor_matches;
constexpr int max_matches = max_4 > tile_matches ? max_4 : tile_matches;
// If we cannot match with neither WMMA nor XDL, try all algorithms for diagnostics
// and see whichi is the closest match.
if constexpr(max_matches == xdl_v3_matches) {
@@ -663,35 +676,78 @@ consteval void diagnose_bwd_weight_algorithm_signature()
constexpr int two_stage_xdl_matches = count_matches_bwd_two_stage_xdl<AlgoType>();
constexpr int dl_matches = count_matches_bwd_dl<AlgoType>();
constexpr int multi_d_xdl_matches = count_matches_bwd_multi_d_xdl<AlgoType>();
constexpr int wmma_v3_matches = count_matches_bwd_wmma_v3<AlgoType>();
constexpr int max1 = xdl_v3_matches > xdl_matches ? xdl_v3_matches : xdl_matches;
constexpr int max2 = max1 > two_stage_xdl_matches ? max1 : two_stage_xdl_matches;
constexpr int max3 = max2 > dl_matches ? max2 : dl_matches;
constexpr int max_matches = max3 > multi_d_xdl_matches ? max3 : multi_d_xdl_matches;
// Check whether we have XDL or WMMA algorithm
if constexpr (SpecifiesGridwiseBwdXdlGemm<AlgoType>)
{
constexpr int max1 = xdl_v3_matches > xdl_matches ? xdl_v3_matches : xdl_matches;
constexpr int max2 = max1 > two_stage_xdl_matches ? max1 : two_stage_xdl_matches;
constexpr int max3 = max2 > dl_matches ? max2 : dl_matches;
constexpr int max_matches = max3 > multi_d_xdl_matches ? max3 : multi_d_xdl_matches;
if constexpr (max_matches == xdl_matches) {
using Alg = BwdXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == xdl_v3_matches) {
using Alg = BwdXdlV3Algorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
if constexpr (max_matches == xdl_matches) {
using Alg = BwdXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == xdl_v3_matches) {
using Alg = BwdXdlV3Algorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == two_stage_xdl_matches) {
using Alg = BwdTwoStageXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == dl_matches) {
using Alg = BwdDlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == multi_d_xdl_matches) {
using Alg = BwdMultiDXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
}
else if constexpr (max_matches == two_stage_xdl_matches) {
using Alg = BwdTwoStageXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
else if constexpr (SpecifiesGridwiseWmmaGemm<AlgoType>)
{
constexpr int max_matches = wmma_v3_matches;
if constexpr (max_matches == wmma_v3_matches) {
using Alg = BwdWmmaV3Algorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
}
else if constexpr (max_matches == dl_matches) {
using Alg = BwdDlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == multi_d_xdl_matches) {
using Alg = BwdMultiDXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else {
// This should never happen
static_assert(false, "Internal Error: No matching algorithm variant found for diagnostics.");
else
{
// If we cannot match with neither WMMA nor XDL, try all algorithms for diagnostics
// and see which is the closest match.
constexpr int max1 = xdl_v3_matches > xdl_matches ? xdl_v3_matches : xdl_matches;
constexpr int max2 = max1 > two_stage_xdl_matches ? max1 : two_stage_xdl_matches;
constexpr int max3 = max2 > dl_matches ? max2 : dl_matches;
constexpr int max_matches = max3 > multi_d_xdl_matches ? max3 : multi_d_xdl_matches;
if constexpr (max_matches == xdl_matches) {
using Alg = BwdXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == xdl_v3_matches) {
using Alg = BwdXdlV3Algorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == two_stage_xdl_matches) {
using Alg = BwdTwoStageXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == dl_matches) {
using Alg = BwdDlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else if constexpr (max_matches == multi_d_xdl_matches) {
using Alg = BwdMultiDXdlAlgorithm<AlgoType>;
static_assert(Alg::is_valid(), Alg::message());
}
else {
// This should never happen
static_assert(false, "Internal Error: No matching algorithm variant found for diagnostics.");
}
}
}

View File

@@ -10,13 +10,12 @@
#include "ck_tile/builder/factory/helpers/ck/conv_tensor_layout.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_tensor_type.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_elementwise_op.hpp"
//#include "ck_tile/builder/factory/helpers/ck/conv_tuning_params.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_thread_block.hpp"
namespace ck_tile::builder::factory {
// Factory for DeviceGroupedConvBwdWeight_Dl instance
// of a grouped forward convolution kernel.
// of a grouped bwd weight convolution kernel.
template <ConvSignatureDescriptor auto SIGNATURE,
ConvAlgorithmDescriptor auto ALGORITHM,
StringLiteral VERSION>

View File

@@ -17,8 +17,8 @@
namespace ck_tile::builder::factory {
// Factory for DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle instance
// of a grouped forward convolution kernel.
// Factory for DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle instance
// of a grouped bwd weight convolution kernel.
template <ConvSignatureDescriptor auto SIGNATURE,
ConvAlgorithmDescriptor auto ALGORITHM,
StringLiteral VERSION>

View File

@@ -17,8 +17,8 @@
namespace ck_tile::builder::factory {
// Factory for DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle instance
// of a grouped forward convolution kernel.
// Factory for DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle instance
// of a grouped bwd weight convolution kernel.
template <ConvSignatureDescriptor auto SIGNATURE,
ConvAlgorithmDescriptor auto ALGORITHM,
StringLiteral VERSION>

View File

@@ -0,0 +1,104 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp"
#include "ck_tile/builder/conv_signature_concepts.hpp"
#include "ck_tile/builder/conv_algorithm_concepts.hpp"
#include "ck_tile/builder/conv_algorithm_limits.hpp"
#include "ck_tile/builder/builder_utils.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_tensor_layout.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_tensor_type.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_elementwise_op.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_tuning_params.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_block_transfer.hpp"
#include "ck_tile/builder/factory/helpers/ck/conv_thread_block.hpp"
namespace ck_tile::builder::factory {
// Factory for DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3 instance
// of a grouped bwd weight convolution kernel.
template <ConvSignatureDescriptor auto SIGNATURE,
ConvAlgorithmDescriptor auto ALGORITHM,
StringLiteral VERSION>
requires ConvDirectionIsBackwardWeight<SIGNATURE>
struct ConvBwdWeightWmmaV3Factory
{
static constexpr size_t SPATIAL_DIM = SIGNATURE.spatial_dim;
using Layouts = internal::ConvTensorLayouts<SIGNATURE, SPATIAL_DIM>;
using Types = internal::BwdWeightConvTensorDataTypes<SIGNATURE>;
using Ops = internal::ElementwiseOps<SIGNATURE>;
using AlgorithmType = decltype(ALGORITHM);
static constexpr auto BWD_CONV_SPECIALIZATION = internal::SetBwdWeightConvSpecialization<ALGORITHM>();
static constexpr auto BLOCK = internal::SetThreadBlockInfo<ALGORITHM>();
static constexpr auto GRIDWISE_GEMM = ALGORITHM.gridwise_gemm;
static constexpr auto A_BLOCK_TRANSFER =
internal::SetBwdConvBlockTransfer<ALGORITHM.transfer.a>();
static constexpr auto B_BLOCK_TRANSFER =
internal::SetBwdConvBlockTransfer<ALGORITHM.transfer.b>();
static constexpr auto C_BLOCK_TRANSFER = internal::SetCBlockTransfer<SIGNATURE, ALGORITHM>();
static constexpr auto BLOCK_GEMM = internal::SetBlockGemm<ALGORITHM>();
// Check limits for the algorithm parameters.
// TODO: Add more limits checks as needed.
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(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::DeviceGroupedConvBwdWeight_Wmma_CShuffleV3<
SPATIAL_DIM,
typename Layouts::InLayout,
typename Layouts::WeiLayout,
typename Layouts::OutLayout,
typename Types::InDataType,
typename Types::WeiDataType,
typename Types::OutDataType,
typename Types::AccDataType,
typename Ops::InElementwiseOp,
typename Ops::WeiElementwiseOp,
typename Ops::OutElementwiseOp,
BWD_CONV_SPECIALIZATION,
BLOCK.block_size,
BLOCK.per_block.m,
BLOCK.per_block.n,
BLOCK.per_block.k,
GRIDWISE_GEMM.k1,
GRIDWISE_GEMM.m_per_wmma,
GRIDWISE_GEMM.n_per_wmma,
GRIDWISE_GEMM.m_wmma_per_wave,
GRIDWISE_GEMM.n_wmma_per_wave,
to_sequence_v<A_BLOCK_TRANSFER.thread_cluster_dims>,
to_sequence_v<A_BLOCK_TRANSFER.thread_cluster_order>,
to_sequence_v<A_BLOCK_TRANSFER.src_access_order>,
A_BLOCK_TRANSFER.src_vector_dim,
A_BLOCK_TRANSFER.src_scalar_per_vector,
A_BLOCK_TRANSFER.lds_dst_scalar_per_vector,
A_BLOCK_TRANSFER.lds_padding,
to_sequence_v<B_BLOCK_TRANSFER.thread_cluster_dims>,
to_sequence_v<B_BLOCK_TRANSFER.thread_cluster_order>,
to_sequence_v<B_BLOCK_TRANSFER.src_access_order>,
B_BLOCK_TRANSFER.src_vector_dim,
B_BLOCK_TRANSFER.src_scalar_per_vector,
B_BLOCK_TRANSFER.lds_dst_scalar_per_vector,
B_BLOCK_TRANSFER.lds_padding,
C_BLOCK_TRANSFER.m_xdl_per_wave_per_shuffle,
C_BLOCK_TRANSFER.n_xdl_per_wave_per_shuffle,
to_sequence_v<C_BLOCK_TRANSFER.thread_cluster_dims>,
C_BLOCK_TRANSFER.scalar_per_vector,
BLOCK_GEMM.scheduler,
BLOCK_GEMM.pipeline_version,
typename Types::InComputeType,
typename Types::WeiComputeType,
ALGORITHM.max_transpose_transfer_src_scalar_per_vector,
ALGORITHM.max_transpose_transfer_dst_scalar_per_vector>;
};
} // namespace ck_tile::builder::factory

View File

@@ -17,8 +17,8 @@
namespace ck_tile::builder::factory {
// Factory for DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle instance
// of a grouped forward convolution kernel.
// Factory for DeviceGroupedConvBwdWeight_Xdl_CShuffle instance
// of a grouped bwd weight convolution kernel.
template <ConvSignatureDescriptor auto SIGNATURE,
ConvAlgorithmDescriptor auto ALGORITHM,
StringLiteral VERSION>

View File

@@ -17,8 +17,8 @@
namespace ck_tile::builder::factory {
// Factory for DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle instance
// of a grouped forward convolution kernel.
// Factory for DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 instance
// of a grouped bwd weight convolution kernel.
template <ConvSignatureDescriptor auto SIGNATURE,
ConvAlgorithmDescriptor auto ALGORITHM,
StringLiteral VERSION>

View File

@@ -73,6 +73,7 @@
#include "ck_tile/builder/factory/conv_bwd_weight_two_stage_xdl_factory.hpp"
#include "ck_tile/builder/factory/conv_bwd_weight_dl_factory.hpp"
#include "ck_tile/builder/factory/conv_bwd_weight_multi_d_xdl_factory.hpp"
#include "ck_tile/builder/factory/conv_bwd_weight_wmma_v3_factory.hpp"
namespace ck_tile::builder::factory {
@@ -173,8 +174,7 @@ constexpr auto make_conv_instance()
}
else if constexpr (BwdWmmaV3Algorithm<AlgoType>::is_valid())
{
static_assert(false,
"Backward weight convolution: WMMA V3 algorithm not yet implemented.");
return typename ConvBwdWeightWmmaV3Factory<SIGNATURE, ALGORITHM, VERSION>::Instance{};
}
else
{

View File

@@ -170,6 +170,7 @@ add_ck_builder_test(test_ckb_build_bwd_weight_instances
conv/ck/test_ckb_conv_bwd_weight_multi_d_xdl_cshuffle.cpp
conv/ck/test_ckb_conv_bwd_weight_xdl_cshuffle_v3.cpp
conv/ck/test_ckb_conv_bwd_weight_dl.cpp
conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle_v3.cpp
conv/ck_tile/test_ckb_conv_bwd_weight_2d_fp16_v3.cpp
)
target_link_libraries(test_ckb_build_bwd_weight_instances PRIVATE utility)

View File

@@ -0,0 +1,44 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include "utils/ckb_conv_test_configs.hpp"
#include "utils/ckb_conv_test_utils.hpp"
#include "utils/conv_algorithm_type_utils.hpp"
#include "ck_tile/host/device_prop.hpp"
namespace ckb = ck_tile::builder;
namespace ckt = ck_tile::builder::test;
namespace cku = ck_tile::builder::test_utils;
using enum ck_tile::builder::TensorLayout;
constexpr auto SIGNATURE =
ckt::ConvSignature{.spatial_dim = 1,
.direction = ckb::ConvDirection::BACKWARD_WEIGHT,
.data_type = ckb::DataType::BF16,
.accumulation_data_type = ckb::DataType::FP32,
.input = {.config = {.layout = NGCW}},
.weight = {.config = {.layout = GKXC}},
.output = {.config = {.layout = NGKW}}};
constexpr auto ALGORITHM = cku::ConvAlgorithm_DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3{}
.with_thread_block(cku::ThreadBlock_64_32x32x32)
.with_gemm_config(cku::GemmParams_Wmma_2x1_per_wave)
.with_transfer(cku::BwdTransfer_4x8x1_4x16x1_v3)
.with_bwd_specialization(ckb::ConvSpecialization::FILTER_1X1_STRIDE1_PAD0)
.with_block_gemm(cku::BlockGemmDesc_v2_intrawave)
.with_transpose_params(4,4);
using Builder = ckb::ConvBuilder<SIGNATURE, ALGORITHM>;
using Instance = Builder::Instance;
TEST(BwdWeight_1DBf16_Wmma_CShuffle_V3, Create)
{
const auto expected_transfer_parameters = to_string(ALGORITHM);
cku::run_test<Builder>({"DeviceGroupedConvBwdWeight_Wmma_CShuffleV3",
expected_transfer_parameters,
"Filter1x1Stride1Pad0",
"NGCW,GKXC,NGKW",
"PassThrough,PassThrough,PassThrough",
"Intrawave",
"v2"});
}

View File

@@ -237,6 +237,15 @@ inline std::string to_string<DlEpilogue>(DlEpilogue t)
return oss.str();
}
template <>
inline std::string to_string<TransposeParams_>(TransposeParams_ t)
{
std::ostringstream oss;
oss << t.max_transpose_transfer_src_scalar_per_vector << ","
<< t.max_transpose_transfer_dst_scalar_per_vector;
return oss.str();
}
template <>
inline std::string to_string<DlTransfer<4>>(DlTransfer<4> t)
{
@@ -414,6 +423,16 @@ inline std::string to_string<ConvAlgorithm_DeviceGroupedConvBwdWeight_Xdl_CShuff
return oss.str();
}
template <>
inline std::string to_string<ConvAlgorithm_DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3>(
ConvAlgorithm_DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3 t)
{
std::ostringstream oss;
oss << to_string(static_cast<ThreadBlock_>(t)) << "," << to_string(static_cast<WmmaGemm_>(t))
<< "," << to_string(static_cast<Transfer_<>>(t));
return oss.str();
}
template <>
inline std::string to_string<ConvAlgorithm_DeviceGroupedConvBwdWeight_TwoStage_Xdl_CShuffle>(
ConvAlgorithm_DeviceGroupedConvBwdWeight_TwoStage_Xdl_CShuffle t)