diff --git a/example/ck_tile/03_gemm/gemm_basic.cpp b/example/ck_tile/03_gemm/gemm_basic.cpp index e50c83e895..f346b65e35 100644 --- a/example/ck_tile/03_gemm/gemm_basic.cpp +++ b/example/ck_tile/03_gemm/gemm_basic.cpp @@ -12,7 +12,8 @@ #include "ck_tile/host.hpp" #include "gemm_utils.hpp" -template ( + return run_gemm_example_with_layouts( argc, argv, Row{}, Col{}, Row{}); } else if(a_layout == "C" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Col{}, Col{}, Row{}); } else @@ -146,23 +147,23 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a { if(a_layout == "R" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Row{}, Col{}, Row{}); } #if 0 else if(a_layout == "R" && b_layout == "R") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Row{}, Row{}, Row{}); } else if(a_layout == "C" && b_layout == "R") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Col{}, Row{}, Row{}); } else if(a_layout == "C" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Col{}, Col{}, Row{}); } #endif diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp index 45c2513733..78f42e50f3 100644 --- a/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/example/ck_tile/03_gemm/gemm_utils.hpp @@ -15,6 +15,7 @@ #define CK_TILE_PIPELINE_MEMORY 2 #define CK_TILE_PIPELINE_COMPUTE_V4 3 +#if 0 #ifndef CK_TILE_PIPELINE_DEFAULT #define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_COMPUTE_V3 #endif @@ -34,10 +35,29 @@ #else #error "unsupported CK_TILE_PIPELINE_DEFAULT value" #endif +#endif -struct GemmConfig +struct GemmConfigBase +{ + static constexpr bool kPadM = false; + static constexpr bool kPadN = false; + static constexpr bool kPadK = false; + + static constexpr bool PermuteA = false; + static constexpr bool PermuteB = false; + + static constexpr bool TransposeC = false; + static constexpr bool UseStructuredSparsity = false; + + static constexpr int kBlockPerCu = 1; + static constexpr ck_tile::index_t TileParitionerGroupNum = 8; + static constexpr ck_tile::index_t TileParitionerM01 = 4; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; +}; + + +struct GemmConfig_1 : public GemmConfigBase { -#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY) // Memory friendly for Interwave scheduler static constexpr ck_tile::index_t M_Tile = 128; static constexpr ck_tile::index_t N_Tile = 32; @@ -52,8 +72,31 @@ struct GemmConfig static constexpr ck_tile::index_t K_Warp_Tile = 8; static constexpr bool DoubleSmemBuffer = false; -#endif -#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V3) + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY; +}; + +struct GemmConfig_1_1 : public GemmConfigBase +{ + // Memory friendly for Interwave scheduler + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 32; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 4; + static constexpr ck_tile::index_t N_Warp = 1; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 8; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave; +}; + +struct GemmConfig_2 : public GemmConfigBase +{ // Compute friendly for Intrawave scheduler static constexpr ck_tile::index_t M_Tile = 128; static constexpr ck_tile::index_t N_Tile = 128; @@ -68,7 +111,55 @@ struct GemmConfig static constexpr ck_tile::index_t K_Warp_Tile = 32; static constexpr bool DoubleSmemBuffer = false; -#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V4) + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3; + + static constexpr int kBlockPerCu = 2; +}; + +struct GemmConfig_2_1 : public GemmConfigBase +{ + // Compute friendly for Intrawave scheduler + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 128; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 32; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave; + static constexpr int kBlockPerCu = 2; +}; + +struct GemmConfig_2_2 : public GemmConfigBase +{ + // Compute friendly for Intrawave scheduler + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 32; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3; + //static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave; + //static constexpr int kBlockPerCu = 2; +}; + +struct GemmConfig_3 : public GemmConfigBase +{ // Compute friendly for Intrawave scheduler // Using the ping pong reader in the lds level static constexpr ck_tile::index_t M_Tile = 256; @@ -84,21 +175,28 @@ struct GemmConfig static constexpr ck_tile::index_t K_Warp_Tile = 16; static constexpr bool DoubleSmemBuffer = true; -#endif + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4; +}; - static constexpr bool kPadM = false; - static constexpr bool kPadN = false; - static constexpr bool kPadK = false; +struct GemmConfig_3_1 : public GemmConfigBase +{ + // Compute friendly for Intrawave scheduler + // Using the ping pong reader in the lds level + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 32; - static constexpr bool PermuteA = false; - static constexpr bool PermuteB = false; + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; - static constexpr bool TransposeC = false; - static constexpr bool UseStructuredSparsity = false; + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 16; - static constexpr int kBlockPerCu = 2; - static constexpr ck_tile::index_t TileParitionerGroupNum = 8; - static constexpr ck_tile::index_t TileParitionerM01 = 4; + static constexpr bool DoubleSmemBuffer = true; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave; }; template diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index 79ed9ce76b..1ad0f2c139 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -30,7 +30,8 @@ auto calculate_rtol_atol(const ck_tile::index_t K, return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); } -template ; - using GemmPipeline = GEMM_PIPELINE; + constexpr auto GetGemmPipeline = [&]() { + if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_MEMORY) + { + return ck_tile::GemmPipelineAgBgCrMem{}; + } + else if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_COMPUTE_V3) + { + return ck_tile::GemmPipelineAgBgCrCompV3{}; + } + else if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_COMPUTE_V4) + { + return ck_tile::GemmPipelineAgBgCrCompV4{}; + } + else + { + static_assert(0, "Invalid pipeline version"); + } + }; + using GemmPipeline = decltype(GetGemmPipeline()); const ck_tile::index_t K = tensor.get_length(0); const ck_tile::index_t N = tensor.get_length(1); @@ -144,7 +163,8 @@ void permute_vectors_i4x4_b(Tensor& tensor) } } -template ( - args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); + gemm_calc(args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); std::size_t flop = std::size_t(2) * M * N * K; std::size_t num_byte = @@ -199,7 +225,8 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, return ave_time; } -template b_k_n_dev = b_k_n; if constexpr(GemmConfig::PermuteB) { - permute_tensor_b( - a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat); + invoke_gemm(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + kbatch, + n_warmup, + n_repeat); c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); bool pass = true; diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index 6c1e322623..c056ae4c3f 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -25,7 +25,8 @@ void try_run(ck_tile::TailNumber tn) } } -template ; - using BaseGemmPipeline = UNIVERSAL_GEMM_PIPELINE; + constexpr auto GetBaseGemmPipeline = [&]() { + if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_MEMORY) + { + return ck_tile::BaseGemmPipelineAgBgCrMem{}; + } + else if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_COMPUTE_V3) + { + return ck_tile::BaseGemmPipelineAgBgCrCompV3{}; + } + else if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_COMPUTE_V4) + { + return ck_tile::BaseGemmPipelineAgBgCrCompV4{}; + } + else + { + static_assert(0, "Invalid pipeline version"); + } + }; + + using BaseGemmPipeline = decltype(GetBaseGemmPipeline()); const ck_tile::index_t k_grain = args.k_batch * GemmConfig::K_Tile; const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * GemmConfig::K_Tile; @@ -78,7 +98,7 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& const auto memory_operation_) { constexpr bool has_hot_loop_v = has_hot_loop_.value; constexpr auto tail_number_v = tail_number_.value; - constexpr auto scheduler = GEMM_PIPELINE_SCHEDULER; + constexpr auto scheduler = GemmConfig::Scheduler; constexpr auto memory_operation = memory_operation_.value; using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem; - - using GemmPipeline = GEMM_PIPELINE; + constexpr auto GetGemmPipeline = [&]() { + if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_MEMORY) + { + return ck_tile::GemmPipelineAgBgCrMem{}; + } + else if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_COMPUTE_V3) + { + return ck_tile::GemmPipelineAgBgCrCompV3{}; + } + else if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_COMPUTE_V4) + { + return ck_tile::GemmPipelineAgBgCrCompV4{}; + } + else + { + static_assert(0, "Invalid pipeline version"); + } + }; + using GemmPipeline = decltype(GetGemmPipeline()); using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem{}, @@ -179,7 +216,8 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& << "\n File: " << __FILE__ << ":" << __LINE__ << ", in function: " << __func__; throw std::runtime_error(err.str()); } -#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY) + } + if constexpr (GemmConfig::Pipeline == CK_TILE_PIPELINE_MEMORY) { if(tail_num == ck_tile::TailNumber::One) { RunSplitk(ck_tile::bool_constant{}, @@ -201,8 +239,8 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& ck_tile::integral_constant{}, ck_tile::integral_constant{}, ck_tile::integral_constant{}); - -#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V4) + } + if constexpr(GemmConfig::Pipeline == CK_TILE_PIPELINE_COMPUTE_V4) { if(tail_num == ck_tile::TailNumber::Three) { RunSplitk( @@ -214,7 +252,7 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& RunSplitk(ck_tile::bool_constant{}, ck_tile::integral_constant{}); } -#endif + } } else { @@ -250,7 +288,10 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& #include "run_gemm_example.inc" -template +template int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int argc, char* argv[]) { using Row = ck_tile::tensor_layout::gemm::RowMajor; @@ -260,12 +301,12 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a { if(a_layout == "R" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Row{}, Col{}, Row{}); } else if(a_layout == "C" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Col{}, Col{}, Row{}); } else @@ -278,24 +319,24 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a { if(a_layout == "R" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Row{}, Col{}, Row{}); } #if 0 else if(a_layout == "R" && b_layout == "R") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Row{}, Row{}, Row{}); } else if(a_layout == "C" && b_layout == "R") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Col{}, Row{}, Row{}); } else if(a_layout == "C" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Col{}, Col{}, Row{}); } #endif @@ -306,6 +347,7 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a } } +template int run_gemm_example(int argc, char* argv[]) { auto [result, arg_parser] = create_args(argc, argv); @@ -318,21 +360,21 @@ int run_gemm_example(int argc, char* argv[]) if(data_type == "fp16") { - return run_gemm_example_prec_type(a_layout, b_layout, argc, argv); + return run_gemm_example_prec_type(a_layout, b_layout, argc, argv); } #if 0 else if(data_type == "bf16") { - return run_gemm_example_prec_type(a_layout, b_layout, argc, argv); + return run_gemm_example_prec_type(a_layout, b_layout, argc, argv); } else if(data_type == "fp8") { - return run_gemm_example_prec_type( + return run_gemm_example_prec_type( a_layout, b_layout, argc, argv); } else if(data_type == "bf8") { - return run_gemm_example_prec_type( + return run_gemm_example_prec_type( a_layout, b_layout, argc, argv); } @@ -340,7 +382,7 @@ int run_gemm_example(int argc, char* argv[]) else if(data_type == "pk_int4_t") { // TODO: Add support for bhalf_t ADataType - return run_gemm_example_prec_type( + return run_gemm_example_prec_type( a_layout, b_layout, argc, argv); } #endif @@ -355,7 +397,13 @@ int main(int argc, char* argv[]) { try { - run_gemm_example(argc, argv); + run_gemm_example(argc, argv); + run_gemm_example(argc, argv); + run_gemm_example(argc, argv); + //run_gemm_example(argc, argv); + run_gemm_example(argc, argv); + run_gemm_example(argc, argv); + //run_gemm_example(argc, argv); } catch(const std::runtime_error& e) {