From cd1c1e0aff88a03780830bc22e2d037be45b9f09 Mon Sep 17 00:00:00 2001 From: John Shumway Date: Fri, 5 Sep 2025 23:14:02 +0000 Subject: [PATCH] Add block GEMM pipeline version to builder. --- .../ck_tile/builder/conv_algorithm.hpp | 15 ++- .../include/ck_tile/builder/conv_factory.hpp | 19 +++- .../builder/test/test_conv_instances.cpp | 93 +++++++++++-------- 3 files changed, 84 insertions(+), 43 deletions(-) diff --git a/experimental/builder/include/ck_tile/builder/conv_algorithm.hpp b/experimental/builder/include/ck_tile/builder/conv_algorithm.hpp index 55c39765aa..f2948fa56f 100644 --- a/experimental/builder/include/ck_tile/builder/conv_algorithm.hpp +++ b/experimental/builder/include/ck_tile/builder/conv_algorithm.hpp @@ -117,7 +117,7 @@ struct BlockCTransferLengthsInfo int n_block; int n_wave_per_xdl; }; -static_assert(BlockBTransferLengths); +static_assert(BlockCTransferLengths); // Concept to check if a struct provides A Block tranfer info. template @@ -137,6 +137,19 @@ concept HasCBlockTransferInfo = requires(T t) { { T::block_transfer.thread_cluster_lengths_c } -> BlockCTransferLengths; }; +enum class BlockGemmPipelineVersion +{ + V3, + V4, + V5 +}; + +// Concept to check if struct provides block_gemm_pipeline_version. +template +concept ProvidesBlockGemmPipelineVersion = requires { + { T::pipeline_version } -> std::convertible_to; +}; + // No requirements yet for a ConvAlogorithm concept. template concept ConvAlgorithm = std::is_class_v; diff --git a/experimental/builder/include/ck_tile/builder/conv_factory.hpp b/experimental/builder/include/ck_tile/builder/conv_factory.hpp index 9971b8b76e..8a01a4bf7c 100644 --- a/experimental/builder/include/ck_tile/builder/conv_factory.hpp +++ b/experimental/builder/include/ck_tile/builder/conv_factory.hpp @@ -217,6 +217,23 @@ constexpr CBlockTransfer SetCBlockTransfer() return block_transfer; } +template +constexpr ck::BlockGemmPipelineVersion SetBlockGemmPipelineVersion() +{ + using AlgorithmType = decltype(ALGORITHM); + if constexpr(ProvidesBlockGemmPipelineVersion) + { + switch(ALGORITHM.pipeline_version) + { + case BlockGemmPipelineVersion::V3: return ck::BlockGemmPipelineVersion::v3; + case BlockGemmPipelineVersion::V4: return ck::BlockGemmPipelineVersion::v4; + case BlockGemmPipelineVersion::V5: return ck::BlockGemmPipelineVersion::v5; + } + } + // Default value. + return ck::BlockGemmPipelineVersion::v4; +} + // Factory builds an instance of a grouped convolution kernel. template requires SupportedVersion @@ -236,7 +253,7 @@ struct GroupedConvForwardXldCShuffleFactoryV3 static constexpr BlockTransfer B_BLOCK_TRANSFER = SetBBlockTransfer(); static constexpr CBlockTransfer C_BLOCK_TRANSFER = SetCBlockTransfer(); static constexpr auto PIPELINE_SCHEDULER = ck::BlockGemmPipelineScheduler::Intrawave; - static constexpr auto PIPELINE_VERSION = ck::BlockGemmPipelineVersion::v4; + static constexpr auto PIPELINE_VERSION = SetBlockGemmPipelineVersion(); // The convlution kernel class instance. using Instance = ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3< // diff --git a/experimental/builder/test/test_conv_instances.cpp b/experimental/builder/test/test_conv_instances.cpp index 67dfefd2df..228ea447f6 100644 --- a/experimental/builder/test/test_conv_instances.cpp +++ b/experimental/builder/test/test_conv_instances.cpp @@ -27,6 +27,7 @@ struct FwdConvAlgorithm ckb::BlockBTransferLengthsInfo thread_cluster_lengths_b; ckb::BlockCTransferLengthsInfo thread_cluster_lengths_c; } block_transfer; + ckb::BlockGemmPipelineVersion pipeline_version; }; static_assert(ckb::ConvAlgorithm); static_assert(ckb::HasThreadBlockInfo); @@ -34,6 +35,7 @@ static_assert(ckb::HasConvTuningInfo); static_assert(ckb::HasABlockTransferInfo); static_assert(ckb::HasBBlockTransferInfo); static_assert(ckb::HasCBlockTransferInfo); +static_assert(ckb::ProvidesBlockGemmPipelineVersion); struct TestCase { @@ -47,62 +49,71 @@ constexpr std::array TEST_CASES = { TestCase{ .name = "ConvFwdXdlBf16CompInstances2x_0", .algorithm = - {.thread_block{.block_size = 256, .sub_matrix = {.m = 256, .n = 128, .k = 64}}, - .tuning_params{.ak1 = 16, .bk1 = 16, .m_xdl_per_wave = 2, .n_xdl_per_wave = 2}, - .block_transfer{ - .thread_cluster_lengths_a = {.k0 = 4, .m = 64, .k1 = 1}, - .thread_cluster_lengths_b = {.k0 = 4, .n = 64, .k1 = 1}, - .thread_cluster_lengths_c = - {.m_block = 1, .m_wave_per_xdl = 32, .n_block = 1, .n_wave_per_xdl = 8}, - }}, + { + .thread_block{.block_size = 256, .sub_matrix = {.m = 256, .n = 128, .k = 64}}, + .tuning_params{.ak1 = 16, .bk1 = 16, .m_xdl_per_wave = 2, .n_xdl_per_wave = 2}, + .block_transfer{ + .thread_cluster_lengths_a = {.k0 = 4, .m = 64, .k1 = 1}, + .thread_cluster_lengths_b = {.k0 = 4, .n = 64, .k1 = 1}, + .thread_cluster_lengths_c = + {.m_block = 1, .m_wave_per_xdl = 32, .n_block = 1, .n_wave_per_xdl = 8}, + }, + .pipeline_version = ckb::BlockGemmPipelineVersion::V4, + }, .expected_type = "DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<256, 256, 128, 64, Default, 32, 32, " "2, 2, 8, 8, 8, 1, 1, BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v4>", }, TestCase{ - .name = "GroupedConvFwdXdlBf16CompInstance0", - .algorithm = {.thread_block{.block_size = 256, .sub_matrix = {.m = 256, .n = 256, .k = 32}}, - .tuning_params{.ak1 = 8, .bk1 = 8, .m_xdl_per_wave = 4, .n_xdl_per_wave = 4}, - .block_transfer{ - .thread_cluster_lengths_a = {.k0 = 4, .m = 64, .k1 = 1}, - .thread_cluster_lengths_b = {.k0 = 4, .n = 64, .k1 = 1}, - .thread_cluster_lengths_c = {.m_block = 1, - .m_wave_per_xdl = 32, - .n_block = 1, - .n_wave_per_xdl = 8}, - }}, + .name = "GroupedConvFwdXdlBf16CompInstance0", + .algorithm = + { + .thread_block{.block_size = 256, .sub_matrix = {.m = 256, .n = 256, .k = 32}}, + .tuning_params{.ak1 = 8, .bk1 = 8, .m_xdl_per_wave = 4, .n_xdl_per_wave = 4}, + .block_transfer{ + .thread_cluster_lengths_a = {.k0 = 4, .m = 64, .k1 = 1}, + .thread_cluster_lengths_b = {.k0 = 4, .n = 64, .k1 = 1}, + .thread_cluster_lengths_c = + {.m_block = 1, .m_wave_per_xdl = 32, .n_block = 1, .n_wave_per_xdl = 8}, + }, + .pipeline_version = ckb::BlockGemmPipelineVersion::V4, + }, .expected_type = "DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<256, 256, 256, 32, Default, 32, 32, " "4, 4, 8, 8, 8, 1, 1, BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v4>", }, TestCase{ - .name = "GroupedConvFwdXdlBf16CompInstance1", - .algorithm = {.thread_block{.block_size = 256, .sub_matrix = {.m = 128, .n = 128, .k = 64}}, - .tuning_params{.ak1 = 8, .bk1 = 8, .m_xdl_per_wave = 2, .n_xdl_per_wave = 2}, - .block_transfer{ - .thread_cluster_lengths_a = {.k0 = 8, .m = 32, .k1 = 1}, - .thread_cluster_lengths_b = {.k0 = 8, .n = 32, .k1 = 1}, - .thread_cluster_lengths_c = {.m_block = 1, - .m_wave_per_xdl = 32, - .n_block = 1, - .n_wave_per_xdl = 8}, - }}, + .name = "GroupedConvFwdXdlBf16CompInstance1", + .algorithm = + { + .thread_block{.block_size = 256, .sub_matrix = {.m = 128, .n = 128, .k = 64}}, + .tuning_params{.ak1 = 8, .bk1 = 8, .m_xdl_per_wave = 2, .n_xdl_per_wave = 2}, + .block_transfer{ + .thread_cluster_lengths_a = {.k0 = 8, .m = 32, .k1 = 1}, + .thread_cluster_lengths_b = {.k0 = 8, .n = 32, .k1 = 1}, + .thread_cluster_lengths_c = + {.m_block = 1, .m_wave_per_xdl = 32, .n_block = 1, .n_wave_per_xdl = 8}, + }, + .pipeline_version = ckb::BlockGemmPipelineVersion::V4, + }, .expected_type = "DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<256, 128, 128, 64, Default, 32, 32, " "2, 2, 8, 8, 8, 1, 1, BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v4>", }, TestCase{ - .name = "GroupedConvFwdXdlBf16CompInstance1", - .algorithm = {.thread_block{.block_size = 256, .sub_matrix = {.m = 128, .n = 128, .k = 32}}, - .tuning_params{.ak1 = 8, .bk1 = 8, .m_xdl_per_wave = 2, .n_xdl_per_wave = 2}, - .block_transfer{ - .thread_cluster_lengths_a = {.k0 = 4, .m = 64, .k1 = 1}, - .thread_cluster_lengths_b = {.k0 = 4, .n = 64, .k1 = 1}, - .thread_cluster_lengths_c = {.m_block = 1, - .m_wave_per_xdl = 32, - .n_block = 1, - .n_wave_per_xdl = 8}, - }}, + .name = "GroupedConvFwdXdlBf16CompInstance1", + .algorithm = + { + .thread_block{.block_size = 256, .sub_matrix = {.m = 128, .n = 128, .k = 32}}, + .tuning_params{.ak1 = 8, .bk1 = 8, .m_xdl_per_wave = 2, .n_xdl_per_wave = 2}, + .block_transfer{ + .thread_cluster_lengths_a = {.k0 = 4, .m = 64, .k1 = 1}, + .thread_cluster_lengths_b = {.k0 = 4, .n = 64, .k1 = 1}, + .thread_cluster_lengths_c = + {.m_block = 1, .m_wave_per_xdl = 32, .n_block = 1, .n_wave_per_xdl = 8}, + }, + .pipeline_version = ckb::BlockGemmPipelineVersion::V4, + }, .expected_type = "DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<256, 128, 128, 32, Default, 32, 32, " "2, 2, 8, 8, 8, 1, 1, BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v4>",