diff --git a/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp b/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp index 1347b132e8..d970de795f 100644 --- a/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp +++ b/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp @@ -74,7 +74,7 @@ struct FwdXdlV3Algorithm { }; template -struct FwdXdlAlgorithm { +struct FwdXdlAlgorithmBase { CHECK_CONCEPT(T, ConvAlgorithmDescriptor) CHECK_CONCEPT(T, SpecifiesThreadBlock) CHECK_CONCEPT(T, SpecifiesBlockTransfer) @@ -106,8 +106,7 @@ struct FwdXdlAlgorithm { } static consteval auto message() -> std::string { - return std::string("\n=== Forward XDL Algorithm Diagnostic (closest match) ===\n" - "Concepts for FwdXdl Algorithm:\n") + + return DIAGNOSTIC_LINE(ConvAlgorithmDescriptor) + DIAGNOSTIC_LINE(SpecifiesThreadBlock) + DIAGNOSTIC_LINE(SpecifiesBlockTransfer) + @@ -123,6 +122,24 @@ struct FwdXdlAlgorithm { } }; +template +struct FwdXdlAlgorithm : public FwdXdlAlgorithmBase{ + CHECK_CONCEPT(T, SpecifiesGenericInstance) + + static constexpr bool c13 = c_SpecifiesGenericInstance; + + static consteval bool is_valid() { + return c13 && FwdXdlAlgorithmBase::is_valid(); + } + + static consteval auto message() -> std::string { + return std::string("\n=== Forward XDL Algorithm Diagnostic (closest match) ===\n" + "Concepts for FwdXdl Algorithm:\n") + + FwdXdlAlgorithmBase::message() + + DIAGNOSTIC_LINE(SpecifiesGenericInstance); + } +}; + template struct FwdWmmaAlgorithm { CHECK_CONCEPT(T, ConvAlgorithmDescriptor) @@ -244,7 +261,7 @@ struct TileAlgorithm { }; template -struct LargeTensorAlgorithm : public FwdXdlAlgorithm +struct LargeTensorAlgorithm : public FwdXdlAlgorithmBase { CHECK_CONCEPT(T, SpecifiesLargeTensorSupport) @@ -256,7 +273,9 @@ struct LargeTensorAlgorithm : public FwdXdlAlgorithm } static consteval auto message() -> std::string { - return FwdXdlAlgorithm::message() + + return std::string("\n=== Forward XDL Large Tensor Algorithm Diagnostic (closest match) ===\n" + "Concepts for FwdLargeTensorXdl Algorithm:\n") + + FwdXdlAlgorithmBase::message() + DIAGNOSTIC_LINE(SpecifiesLargeTensorSupport); } }; @@ -441,7 +460,7 @@ struct BwdTwoStageXdlAlgorithm { }; template -struct BwdWmmaV3Algorithm { +struct BwdWmmaV3AlgorithmBase { CHECK_CONCEPT(T, ConvAlgorithmDescriptor) CHECK_CONCEPT(T, SpecifiesThreadBlock) CHECK_CONCEPT(T, SpecifiesBlockTransfer) @@ -452,7 +471,6 @@ struct BwdWmmaV3Algorithm { CHECK_CONCEPT(T, SpecifiesBwdWeightConvSpecialization) CHECK_CONCEPT(T, SpecifiesBlockGemm) CHECK_CONCEPT(T, TransposeTransferWellDefinedIfProvided) - CHECK_CONCEPT(T, SpecifiesGenericInstance) static constexpr bool c1 = c_ConvAlgorithmDescriptor; static constexpr bool c2 = c_SpecifiesThreadBlock; @@ -464,15 +482,13 @@ struct BwdWmmaV3Algorithm { static constexpr bool c8 = c_SpecifiesBwdWeightConvSpecialization; static constexpr bool c9 = c_SpecifiesBlockGemm; static constexpr bool c10 = c_TransposeTransferWellDefinedIfProvided; - static constexpr bool c11 = c_SpecifiesGenericInstance; static consteval bool is_valid() { - return c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 & c11; + return c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10; } static consteval auto message() -> std::string { - return std::string("\n=== Backward WMMA V3 Algorithm Diagnostic (closest match) ===\n" - "Concepts for BwdXdlV3 Algorithm:\n") + + return DIAGNOSTIC_LINE(ConvAlgorithmDescriptor) + DIAGNOSTIC_LINE(SpecifiesThreadBlock) + DIAGNOSTIC_LINE(SpecifiesBlockTransfer) + @@ -482,11 +498,51 @@ struct BwdWmmaV3Algorithm { DIAGNOSTIC_LINE(SpecifiesGridwiseWmmaGemm) + DIAGNOSTIC_LINE(SpecifiesBwdWeightConvSpecialization) + DIAGNOSTIC_LINE(SpecifiesBlockGemm) + - DIAGNOSTIC_LINE(TransposeTransferWellDefinedIfProvided) + + DIAGNOSTIC_LINE(TransposeTransferWellDefinedIfProvided); + } +}; + +template +struct BwdWmmaV3Algorithm : public BwdWmmaV3AlgorithmBase +{ + CHECK_CONCEPT(T, SpecifiesGenericInstance) + + static constexpr bool c11 = c_SpecifiesGenericInstance; + + static consteval bool is_valid() { + return c11 && BwdWmmaV3AlgorithmBase::is_valid(); + } + + static consteval auto message() -> std::string { + return std::string("\n=== Backward WMMA V3 Algorithm Diagnostic (closest match) ===\n" + "Concepts for BwdWmmaV3 Algorithm:\n") + + BwdWmmaV3AlgorithmBase::message() + DIAGNOSTIC_LINE(SpecifiesGenericInstance); } }; +template +struct BwdTwoStageWmmaV3Algorithm : public BwdWmmaV3AlgorithmBase +{ + CHECK_CONCEPT(T, SpecifiesTwoStageSupport) + CHECK_CONCEPT(T, SpecifiesGemmBatchOptions) + + static constexpr bool c11 = c_SpecifiesTwoStageSupport; + static constexpr bool c12 = c_SpecifiesGemmBatchOptions; + + static consteval bool is_valid() { + return c11 && c12 && BwdWmmaV3AlgorithmBase::is_valid(); + } + + static consteval auto message() -> std::string { + return std::string("\n=== Backward Two Stage WMMA V3 Algorithm Diagnostic (closest match) ===\n" + "Concepts for BwdTwoStageWmmaV3 Algorithm:\n") + + BwdWmmaV3AlgorithmBase::message() + + DIAGNOSTIC_LINE(SpecifiesGemmBatchOptions) + + DIAGNOSTIC_LINE(SpecifiesTwoStageSupport); + } +}; + template struct BwdDlAlgorithm { CHECK_CONCEPT(T, ConvAlgorithmDescriptor) @@ -531,7 +587,7 @@ consteval int count_matches_fwd_xdl_v3() { template consteval int count_matches_fwd_xdl() { using Alg = FwdXdlAlgorithm; - 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; + 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 + Alg::c13; } template @@ -576,6 +632,12 @@ consteval int count_matches_bwd_wmma_v3() { return Alg::c1 + Alg::c2 + Alg::c3 + Alg::c4 + Alg::c5 + Alg::c6 + Alg::c7 + Alg::c8 + Alg::c9 + Alg::c10 + Alg::c11; } +template +consteval int count_matches_bwd_two_stage_wmma_v3() { + using Alg = BwdTwoStageWmmaV3Algorithm; + 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 consteval int count_matches_bwd_dl() { using Alg = BwdDlAlgorithm; @@ -677,6 +739,7 @@ consteval void diagnose_bwd_weight_algorithm_signature() constexpr int dl_matches = count_matches_bwd_dl(); constexpr int multi_d_xdl_matches = count_matches_bwd_multi_d_xdl(); constexpr int wmma_v3_matches = count_matches_bwd_wmma_v3(); + constexpr int two_stage_wmma_v3_matches = count_matches_bwd_two_stage_wmma_v3(); // Check whether we have XDL or WMMA algorithm if constexpr (SpecifiesGridwiseBwdXdlGemm) @@ -709,11 +772,15 @@ consteval void diagnose_bwd_weight_algorithm_signature() } else if constexpr (SpecifiesGridwiseWmmaGemm) { - constexpr int max_matches = wmma_v3_matches; + constexpr int max_matches = wmma_v3_matches > two_stage_wmma_v3_matches ? wmma_v3_matches : two_stage_wmma_v3_matches; if constexpr (max_matches == wmma_v3_matches) { using Alg = BwdWmmaV3Algorithm; static_assert(Alg::is_valid(), Alg::message()); } + else if constexpr (max_matches == two_stage_wmma_v3_matches) { + using Alg = BwdTwoStageWmmaV3Algorithm; + static_assert(Alg::is_valid(), Alg::message()); + } } else { diff --git a/experimental/builder/include/ck_tile/builder/factory/conv_bwd_weight_two_stage_wmma_v3_factory.hpp b/experimental/builder/include/ck_tile/builder/factory/conv_bwd_weight_two_stage_wmma_v3_factory.hpp new file mode 100644 index 0000000000..635b822293 --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/factory/conv_bwd_weight_two_stage_wmma_v3_factory.hpp @@ -0,0 +1,105 @@ +// 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_two_stage_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 DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffle_V3 instance +// of a grouped bwd weight convolution kernel. +template + requires ConvDirectionIsBackwardWeight +struct ConvBwdWeightTwoStageWmmaV3Factory +{ + static constexpr size_t SPATIAL_DIM = SIGNATURE.spatial_dim; + using Layouts = internal::ConvTensorLayouts; + using Types = internal::BwdWeightConvTensorDataTypes; + using Ops = internal::ElementwiseOps; + using AlgorithmType = decltype(ALGORITHM); + + static constexpr auto BWD_CONV_SPECIALIZATION = internal::SetBwdWeightConvSpecialization(); + + static constexpr auto BLOCK = internal::SetThreadBlockInfo(); + static constexpr auto GRIDWISE_GEMM = ALGORITHM.gridwise_gemm; + static constexpr auto A_BLOCK_TRANSFER = + internal::SetBwdConvBlockTransfer(); + static constexpr auto B_BLOCK_TRANSFER = + internal::SetBwdConvBlockTransfer(); + static constexpr auto C_BLOCK_TRANSFER = internal::SetCBlockTransfer(); + static constexpr auto BLOCK_GEMM = internal::SetBlockGemm(); + + // Check limits for the algorithm parameters. + // TODO: Add more limits checks as needed. + static_assert(InputVectorTransferLimits, "Invalid A block transfer config"); + static_assert(InputVectorTransferLimits, "Invalid B block transfer config"); + static_assert(OutputVectorTransferLimits, "Invalid C block transfer config"); + static_assert(AccessOrderLimits3D, "Invalid A thread cluster access order"); + static_assert(AccessOrderLimits3D, "Invalid B thread cluster access order"); + static_assert(AccessOrderLimits3D, "Invalid A source access order"); + static_assert(AccessOrderLimits3D, "Invalid B source access order"); + + // The forward convolution kernel class instance. + using Instance = ck::tensor_operation::device::DeviceGroupedConvBwdWeightTwoStage_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, + to_sequence_v, + to_sequence_v, + 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, + to_sequence_v, + to_sequence_v, + 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.scalar_per_vector, + BLOCK_GEMM.scheduler, + BLOCK_GEMM.pipeline_version, + ALGORITHM.num_conv_groups_to_merge, + 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 diff --git a/experimental/builder/include/ck_tile/builder/factory/conv_dispatcher.hpp b/experimental/builder/include/ck_tile/builder/factory/conv_dispatcher.hpp index 6aecd779db..557f35f861 100644 --- a/experimental/builder/include/ck_tile/builder/factory/conv_dispatcher.hpp +++ b/experimental/builder/include/ck_tile/builder/factory/conv_dispatcher.hpp @@ -74,6 +74,7 @@ #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" +#include "ck_tile/builder/factory/conv_bwd_weight_two_stage_wmma_v3_factory.hpp" namespace ck_tile::builder::factory { @@ -176,6 +177,10 @@ constexpr auto make_conv_instance() { return typename ConvBwdWeightWmmaV3Factory::Instance{}; } + else if constexpr (BwdTwoStageWmmaV3Algorithm::is_valid()) + { + return typename ConvBwdWeightTwoStageWmmaV3Factory::Instance{}; + } else { diagnose_bwd_weight_algorithm_signature(); diff --git a/experimental/builder/test/CMakeLists.txt b/experimental/builder/test/CMakeLists.txt index 3d6e448d84..aedcbb8110 100644 --- a/experimental/builder/test/CMakeLists.txt +++ b/experimental/builder/test/CMakeLists.txt @@ -176,6 +176,7 @@ set(BWD_WEIGHT_TESTS if (CK_USE_WMMA) list(APPEND BWD_WEIGHT_TESTS conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle_v3.cpp + conv/ck/test_ckb_conv_bwd_weight_two_stage_wmma_cshuffle_v3.cpp ) endif() diff --git a/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_two_stage_wmma_cshuffle_v3.cpp b/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_two_stage_wmma_cshuffle_v3.cpp new file mode 100644 index 0000000000..e581da4951 --- /dev/null +++ b/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_two_stage_wmma_cshuffle_v3.cpp @@ -0,0 +1,47 @@ +// 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 = 2, + .direction = ckb::ConvDirection::BACKWARD_WEIGHT, + .data_type = ckb::DataType::FP16, + .accumulation_data_type = ckb::DataType::FP32, + .input = {.config = {.layout = NGCHW}}, + .weight = {.config = {.layout = GKYXC}}, + .output = {.config = {.layout = NGKHW}}}; + +constexpr auto ALGORITHM = cku::ConvAlgorithm_DeviceGroupedConvBwdWeight_TwoStage_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) + .with_bwd_specialization(ckb::ConvSpecialization::DEFAULT) + .with_block_gemm(cku::BlockGemmDesc_v1_intrawave) + .with_num_conv_groups_to_merge(2) + .with_transpose_params(2,2); + +using Builder = ckb::ConvBuilder; +using Instance = Builder::Instance; + +TEST(BwdWeight_2DFp16_TwoStage_Wmma_CShuffle_V3, Create) +{ + const auto expected_transfer_parameters = to_string(ALGORITHM); + std::cout << "Expected Transfer Parameters: " << expected_transfer_parameters << std::endl; + cku::run_test({"DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3", + expected_transfer_parameters, + "Default", + "NGCHW,GKYXC,NGKHW", + "PassThrough,PassThrough,PassThrough", + "Intrawave", + "v1", + "fp16,fp16,2,2>"}); // Check compute types and transpose params. +} diff --git a/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle_v3.cpp b/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle_v3.cpp index 4a1a60e852..f58ec11129 100644 --- a/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle_v3.cpp +++ b/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_wmma_cshuffle_v3.cpp @@ -34,11 +34,13 @@ using Instance = Builder::Instance; TEST(BwdWeight_1DBf16_Wmma_CShuffle_V3, Create) { const auto expected_transfer_parameters = to_string(ALGORITHM); + std::cout << "Expected Transfer Parameters: " << expected_transfer_parameters << std::endl; cku::run_test({"DeviceGroupedConvBwdWeight_Wmma_CShuffleV3", expected_transfer_parameters, "Filter1x1Stride1Pad0", "NGCW,GKXC,NGKW", "PassThrough,PassThrough,PassThrough", "Intrawave", - "v1"}); + "v1", + "bf16,bf16,4,4>"}); } diff --git a/experimental/builder/test/impl/conv_algorithm_types.hpp b/experimental/builder/test/impl/conv_algorithm_types.hpp index b562df2608..404e14fced 100644 --- a/experimental/builder/test/impl/conv_algorithm_types.hpp +++ b/experimental/builder/test/impl/conv_algorithm_types.hpp @@ -604,4 +604,8 @@ using ConvAlgorithm_DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle = using ConvAlgorithm_DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3 = ConvAlgorithmTemplate, ConvSpecializationBwdWeight_, BlockGemm_, TransposeParams_>; + +using ConvAlgorithm_DeviceGroupedConvBwdWeight_TwoStage_Wmma_CShuffle_V3 = + ConvAlgorithmTemplate, ConvSpecializationBwdWeight_, BlockGemm_, TransposeParams_, GemmBatchOptions_, TwoStageSpecialization_>; + } // namespace ck_tile::builder::test diff --git a/experimental/builder/test/utils/conv_algorithm_type_utils.hpp b/experimental/builder/test/utils/conv_algorithm_type_utils.hpp index bfd16814da..161d6be268 100644 --- a/experimental/builder/test/utils/conv_algorithm_type_utils.hpp +++ b/experimental/builder/test/utils/conv_algorithm_type_utils.hpp @@ -433,6 +433,16 @@ inline std::string to_string +inline std::string to_string( + ConvAlgorithm_DeviceGroupedConvBwdWeight_TwoStage_Wmma_CShuffle_V3 t) +{ + std::ostringstream oss; + oss << to_string(static_cast(t)) << "," << to_string(static_cast(t)) + << "," << to_string(static_cast>(t)); + return oss.str(); +} + template <> inline std::string to_string( ConvAlgorithm_DeviceGroupedConvBwdWeight_TwoStage_Xdl_CShuffle t)