diff --git a/experimental/convolution_builder/convolution_builder.hpp b/experimental/convolution_builder/convolution_builder.hpp index 0e6496a020..dc671bfd01 100644 --- a/experimental/convolution_builder/convolution_builder.hpp +++ b/experimental/convolution_builder/convolution_builder.hpp @@ -1,15 +1,92 @@ #pragma once #include "convolution_problem_descriptor.hpp" +#include "convolution_implementation_descriptor.hpp" #include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp" #include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp" -template +template struct ConvolutionBuilder; -template -struct ConvolutionBuilder { +template +struct ConvolutionBuilder { +public: + static constexpr auto GetInstance() { + using DataType = typename Implementation::DataType; + using AccDataType = std::conditional_t, int32_t, float>; + using InLayout = std::tuple_element<0, decltype(GetLayout())>::type; + using WeiLayout = std::tuple_element<1,decltype( GetLayout())>::type; + using OutLayout = std::tuple_element<2,decltype( GetLayout())>::type; + using GroupedConvFwdMultipleABD_Xdl_CShuffleInstance = ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle< Implementation::NDimSpatial_, InLayout, WeiLayout, decltype(GetMultiDLayout()), OutLayout, DataType, DataType, DataType, AccDataType, typename Implementation::ElementwiseOpDataTypes, DataType, ck::tensor_operation::element_wise::PassThrough, ck::tensor_operation::element_wise::PassThrough, decltype(GetOutElementwiseOp()), GetConvSpecialization(), ck::tensor_operation::device::GemmSpecialization::MNKPadding, 1, Implementation::BlockSize_, Implementation::TileSizes_::At(0), Implementation::TileSizes_::At(1), Implementation::TileSizes_::At(2), Implementation::K1_, Implementation::K1_, 16, 16, 1, 1, ck::Sequence<4, 8, 1>, ck::Sequence<1, 0, 2>, ck::Sequence<1, 0, 2>, 2, 1, 4, 1, ck::Sequence<4, 8, 1>, ck::Sequence<1, 0, 2>, ck::Sequence<1, 0, 2>, 2, 1, 4, 1, 1, 1, ck::Sequence<1, 32, 1, 8>, 1>; + using DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffleInstance = ck::tensor_operation::device::DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle, ck::Sequence<2, 0, 1>, ck::Sequence<1, 0, 2>, 1, 1, 4, false, ck::Sequence<4, 8, 1>, ck::Sequence<2, 0, 1>, ck::Sequence<1, 0, 2>, 1, 1, 4, false, 1, 1, ck::Sequence<1, 8, 1, 8>, 1>; + + using SelectedInstance = std::conditional_t; + return SelectedInstance{}; + } + + std::string GetInstanceName() const + { + auto str = std::stringstream(); + + std::map KernelToString{ + {Kernel::GroupedConvFwdMultipleABD_Xdl_CShuffle, "GroupedConvFwdMultipleABD_Xdl_CShuffle"}, + {Kernel::GroupedConvBwdWeightTwoStage_Xdl_CShuffle, "GroupedConvBwdWeightTwoStage_Xdl_CShuffle"}}; + + std::map GemmPipelineSchedulerToString{ + {GemmPipelineScheduler::Intrawave, "Intrawave"}, + {GemmPipelineScheduler::Interwave, "Interwave"}}; + + std::map GemmPipelineVersionToString{ + {GemmPipelineVersion::V1, "v1"}, + {GemmPipelineVersion::V2, "v2"}, + {GemmPipelineVersion::V3, "v3"}, + {GemmPipelineVersion::V4, "v4"}, + {GemmPipelineVersion::V5, "v5"}}; + + std::map MFMAInstructionSizeToString{ + {MFMAInstructionSize::M16N16, "16x16"}, + {MFMAInstructionSize::M32N32, "32x32"}}; + + std::map ConvolutionSpecializationToString{ + {ConvolutionSpecialization::Default, "Default"}, + {ConvolutionSpecialization::Filter1x1Pad0, "Filter1x1Pad0"}, + {ConvolutionSpecialization::Filter1x1Stride1Pad0, "Filter1x1Stride1Pad0"}, + {ConvolutionSpecialization::Filter3x3, "Filter3x3"}}; + + std::map MergedGroupsToString{ + {MergedGroups::X16, "16"}, + {MergedGroups::X8, "8"}, + {MergedGroups::X4, "4"}, + {MergedGroups::X2, "2"}, + {MergedGroups::NotSupported, "1"}}; + + // clang-format off + str << KernelToString[GetKernel()] + << "<" + << Implementation::BlockSize_ << ", " + << std::get<0>(Implementation::TileSizes_) << ", " + << std::get<1>(Implementation::TileSizes_) << ", " + << std::get<2>(Implementation::TileSizes_) << ", " + << ConvolutionSpecializationToString[Implementation::ConvolutionSpecialization_] << ", " + << Implementation::K1_ << ", " + << MFMAInstructionSizeToString[Implementation::MFMAInstructionSize_] << ", " + << std::get<0>(Implementation::XdlPerWave_) << ", " + << std::get<1>(Implementation::XdlPerWave_) << ", " + << std::get<0>(Implementation::GlobalTransferVectorSize_) << ", " + << std::get<0>(Implementation::LDSStoreVectorSize_) << ", " + << std::get<1>(Implementation::GlobalTransferVectorSize_) << ", " + << std::get<1>(Implementation::LDSStoreVectorSize_) << ", " + << std::get<2>(Implementation::GlobalTransferVectorSize_) << ", " + << GemmPipelineSchedulerToString[Problem::GemmPipelineScheduler_] << ", " + << GemmPipelineVersionToString[Problem::GemmPipelineVersion_] << ", " + << MergedGroupsToString[Problem::MergedGroups_] << ">"; + // clang-format on + + return str.str(); + } + +private: enum class Kernel { GroupedConvFwdMultipleABD_Xdl_CShuffle, GroupedConvBwdWeightTwoStage_Xdl_CShuffle @@ -31,25 +108,34 @@ struct ConvolutionBuilder { } } - static constexpr auto GetInstance() { - using FwdInstance = ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle< 2, ck::tensor_layout::convolution::NHWGC, ck::tensor_layout::convolution::GKYXC, ck::Tuple<>, ck::tensor_layout::convolution::NHWGK, float, float, float, float, ck::Tuple<>, float, ck::tensor_operation::element_wise::PassThrough, ck::tensor_operation::element_wise::PassThrough, ck::tensor_operation::element_wise::PassThrough, ck::tensor_operation::device::ConvolutionForwardSpecialization::Default, ck::tensor_operation::device::GemmSpecialization::MNKPadding, 1, 256, 128, 256, 32, 8, 8, 32, 32, 2, 4, ck::Sequence<4, 64, 1>, ck::Sequence<1, 0, 2>, ck::Sequence<1, 0, 2>, 2, 8, 8, 1, ck::Sequence<4, 64, 1>, ck::Sequence<1, 0, 2>, ck::Sequence<1, 0, 2>, 2, 8, 8, 1, 1, 1, ck::Sequence<1, 32, 1, 8>, 8>; - using BwdWeightInstance = ck::tensor_operation::device::DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle<2, ck::tensor_layout::convolution::NHWGC, ck::tensor_layout::convolution::GKYXC, ck::tensor_layout::convolution::NHWGK, float, float, float, float, ck::tensor_operation::element_wise::PassThrough, ck::tensor_operation::element_wise::PassThrough, ck::tensor_operation::element_wise::PassThrough, ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization::Default, 64, 16, 16, 32, 8, 16, 16, 1, 1, ck::Sequence<4, 8, 1>, ck::Sequence<2, 0, 1>, ck::Sequence<1, 0, 2>, 1, 1, 4, false, ck::Sequence<4, 8, 1>, ck::Sequence<2, 0, 1>, ck::Sequence<1, 0, 2>, 1, 1, 4, false, 1, 1, ck::Sequence<1, 8, 1, 8>, 1>; - - using SelectedInstance = std::conditional_t; - return SelectedInstance{}; + static constexpr auto GetLayout() { + if constexpr(Implementation::NDimSpatial_ == 2) { + if constexpr(Implementation::ConvolutionLayout_ == ConvolutionLayout::NHWGC_GKYXC_NHWGK) { + return std::tuple{}; + } else { + static_assert("Layout not supported!"); + } + } else { + static_assert("Not supported spatial dim!"); + } } - std::string GetInstanceName() const - { - if constexpr(GetKernel() == Kernel::GroupedConvFwdMultipleABD_Xdl_CShuffle) { - return "GroupedConvFwdMultipleABD_Xdl_CShuffle"; - } else if constexpr (GetKernel() == Kernel::GroupedConvBwdWeightTwoStage_Xdl_CShuffle) { - return "GroupedConvBwdWeightTwoStage_Xdl_CShuffle"; + static constexpr auto GetMultiDLayout() { + return ck::Tuple<>{}; + } + + static constexpr auto GetOutElementwiseOp() { + return ck::tensor_operation::element_wise::PassThrough{}; + } + + static constexpr auto GetConvSpecialization() { + if constexpr(Problem::ConvolutionDirection_ == ConvolutionDirection::Forward) { + return ck::tensor_operation::device::ConvolutionForwardSpecialization::Default; + } else if constexpr(Problem::ConvolutionDirection_ == ConvolutionDirection::BackwardWeight) { + return ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization::Default; } else { - return "Not found."; + static_assert("Specialization not found!"); } - // const auto instance = GetInstance(); - // return instance.GetTypeString(); } }; diff --git a/experimental/convolution_builder/convolution_example.cpp b/experimental/convolution_builder/convolution_example.cpp index c092b8a1f7..96a2955507 100644 --- a/experimental/convolution_builder/convolution_example.cpp +++ b/experimental/convolution_builder/convolution_example.cpp @@ -7,6 +7,7 @@ // Example of problem description for Forward Conv with default settings struct GroupedConvFwdXdlImplicitGemm : public GroupedConvBaseXdlV1 { static constexpr ConvolutionDirection ConvolutionDirection_ = ConvolutionDirection::Forward; + static constexpr ElementwiseOperation ElementwiseOperation_ = ElementwiseOperation::Bias; }; // Example of problem description for Backward Weight Conv with default settings and Split K Two Stage @@ -15,10 +16,20 @@ struct GroupedConvBwdWeightXdlImplicitGemmTwoStage : public GroupedConvBaseXdlV1 static constexpr SplitKSupport SplitKSupport_ = SplitKSupport::SupportedTwoStage; }; +struct ImplementationDescriptor : public NHWCImplementationBaseV1, public BF16ImplementationBaseV1 { + static constexpr ck::index_t BlockSize_ = 64; + static constexpr auto TileSizes_ = std::make_tuple(16, 16, 32); + static constexpr ck::index_t K1_ = 8; + static constexpr MFMAInstructionSize MFMAInstructionSize_ = MFMAInstructionSize::M16N16; + static constexpr auto XdlPerWave_ = std::make_tuple(16, 16); + static constexpr auto GlobalTransferVectorSize_ = std::make_tuple(1, 1, 1); + static constexpr auto LDSStoreVectorSize_ = std::make_tuple(4, 4); +}; + int main () { - ConvolutionBuilder builder_fwd; + ConvolutionBuilder builder_fwd; std::cout << builder_fwd.GetInstanceName() << std::endl; - ConvolutionBuilder builder_bwd_weight_two_stage; + ConvolutionBuilder builder_bwd_weight_two_stage; std::cout << builder_bwd_weight_two_stage.GetInstanceName() << std::endl; return 0; } diff --git a/experimental/convolution_builder/convolution_implementation_descriptor.hpp b/experimental/convolution_builder/convolution_implementation_descriptor.hpp index 5003e9ea84..27eead016d 100644 --- a/experimental/convolution_builder/convolution_implementation_descriptor.hpp +++ b/experimental/convolution_builder/convolution_implementation_descriptor.hpp @@ -1,5 +1,82 @@ #pragma once #include +#include "ck/utility/data_type.hpp" +#include "ck/utility/sequence.hpp" +#include "ck/utility/tuple.hpp" +#include "ck/ck.hpp" + +enum class ImplementationDescriptorVersion +{ + V1 +}; + + +enum class ConvolutionSpecialization { + Default, + Filter1x1Pad0, + Filter1x1Stride1Pad0, + Filter3x3 +}; + +enum class ConvolutionLayout { + NHWGC_GKYXC_NHWGK, + NGCHW_GKCYX_NGKHW +}; + +enum class MFMAInstructionSize { + M16N16, + M32N32 +}; + + template -concept \ No newline at end of file +concept ImplementationDescriptorV1 = requires { + {T::ImplementationDescriptorVersion_} -> std::convertible_to; + {T::NDimSpatial_} -> std::convertible_to; + typename T::DataType; + typename T::ElementwiseOpDataTypes; + {T::ConvolutionSpecialization_} -> std::convertible_to; + {T::ConvolutionLayout_} -> std::convertible_to; + {T::BlockSize_} -> std::convertible_to; + {T::TileSizes_} -> std::convertible_to>; + {T::K1_} -> std::convertible_to; + {T::MFMAInstructionSize_} -> std::convertible_to; + {T::XdlPerWave_} -> std::convertible_to>; + {T::GlobalTransferVectorSize_} -> std::convertible_to>; + {T::LDSStoreVectorSize_} -> std::convertible_to>; +} && (T::ImplementationDescriptorVersion_ == ImplementationDescriptorVersion::V1); + +struct ImplementationBaseV1 { + static constexpr ImplementationDescriptorVersion ImplementationDescriptorVersion_ = ImplementationDescriptorVersion::V1; + using DataType = ck::bhalf_t; + using ElementwiseOpDataTypes = ck::Tuple<>; + static constexpr ConvolutionSpecialization ConvolutionSpecialization_ = ConvolutionSpecialization::Default; +}; + +struct BF16ImplementationBaseV1 : public ImplementationBaseV1 { + using DataType = ck::bhalf_t; +}; + +struct F32ImplementationBaseV1 : public ImplementationBaseV1 { + using DataType = float; +}; + +struct F16ImplementationBaseV1 : public ImplementationBaseV1 { + using DataType = ck::half_t; +}; + +struct NWCImplementationBaseV1 : public ImplementationBaseV1 { + static constexpr int NDimSpatial_ = 1; + static constexpr ConvolutionLayout ConvolutionLayout_ = ConvolutionLayout::NHWGC_GKYXC_NHWGK; +}; + +struct NHWCImplementationBaseV1 : public ImplementationBaseV1 { + static constexpr int NDimSpatial_ = 2; + static constexpr ConvolutionLayout ConvolutionLayout_ = ConvolutionLayout::NHWGC_GKYXC_NHWGK; +}; + +struct NDHWCImplementationBaseV1 : public ImplementationBaseV1 { + static constexpr int NDimSpatial_ = 3; + static constexpr ConvolutionLayout ConvolutionLayout_ = ConvolutionLayout::NHWGC_GKYXC_NHWGK; +};