From acf28d2c416a782ff5a71ca750e2dc36211531e3 Mon Sep 17 00:00:00 2001 From: Qun Lin Date: Sat, 10 May 2025 16:40:57 +0800 Subject: [PATCH] rebase flatmm fp8 --- example/ck_tile/18_flatmm/CMakeLists.txt | 5 +- example/ck_tile/18_flatmm/flatmm_basic.cpp | 77 +----- example/ck_tile/18_flatmm/flatmm_basic.hpp | 143 ++++++++++- .../ck_tile/18_flatmm/run_flatmm_example.inc | 225 ++++++++++++++---- .../flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 5 + ...mm_pipeline_agmem_bgmem_creg_v1_policy.hpp | 25 +- include/ck_tile/ops/gemm/warp/warp_gemm.hpp | 41 +++- 7 files changed, 390 insertions(+), 131 deletions(-) diff --git a/example/ck_tile/18_flatmm/CMakeLists.txt b/example/ck_tile/18_flatmm/CMakeLists.txt index f4d823e91a..37c8431dcd 100644 --- a/example/ck_tile/18_flatmm/CMakeLists.txt +++ b/example/ck_tile/18_flatmm/CMakeLists.txt @@ -3,6 +3,7 @@ add_executable(tile_example_flatmm_basic EXCLUDE_FROM_ALL flatmm_basic.cpp) set(EXAMPLE_FLATMM_COMPILE_OPTIONS) # list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal) # list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -Wno-unused-variable -Wno-unused-parameter) -list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -DUSING_MFMA_16x16x32=1 -DENABLE_FP8=1 -Wno-unused-local-typedef) -#list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -DUSING_MFMA_32x32x16=1 -DENABLE_FP8=1 -Wno-unused-local-typedef) +#list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -DUSING_MFMA_16x16x32=1 -DENABLE_FP8=1 -Wno-unused-local-typedef) +list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -DUSING_MFMA_32x32x16=1 -DENABLE_FP8=1 -Wno-unused-local-typedef) target_compile_options(tile_example_flatmm_basic PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS}) +target_compile_options(tile_example_flatmm_basic PRIVATE -save-temps=obj -Wno-gnu-line-marker) \ No newline at end of file diff --git a/example/ck_tile/18_flatmm/flatmm_basic.cpp b/example/ck_tile/18_flatmm/flatmm_basic.cpp index 5f2c2a5aab..219c7ecb8f 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.cpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.cpp @@ -16,6 +16,7 @@ template @@ -29,32 +30,19 @@ float flatmm_calc(const ck_tile::FlatmmHostArgs& args, const ck_tile::stream_con constexpr int kBlockPerCu = 2; // This part comes from the Codegen -#if defined(USING_MFMA_16x16x32) || defined(ENABLE_FP16) - constexpr ck_tile::index_t M_Tile = 128; - constexpr ck_tile::index_t N_Tile = 128; - constexpr ck_tile::index_t K_Tile = 128; + static_assert(sizeof(ADataType) == 2 || sizeof(ADataType) == 1); + constexpr ck_tile::index_t M_Tile = GemmConfig::M_Tile; + constexpr ck_tile::index_t N_Tile = GemmConfig::N_Tile; + constexpr ck_tile::index_t K_Tile = GemmConfig::K_Tile; - constexpr ck_tile::index_t M_Warp = 1; - constexpr ck_tile::index_t N_Warp = 4; - constexpr ck_tile::index_t K_Warp = 1; + constexpr ck_tile::index_t M_Warp = GemmConfig::M_Warp; + constexpr ck_tile::index_t N_Warp = GemmConfig::N_Warp; + constexpr ck_tile::index_t K_Warp = GemmConfig::K_Warp; - constexpr ck_tile::index_t M_Warp_Tile = is_8bit_type::value ? 16 : 32; - constexpr ck_tile::index_t N_Warp_Tile = is_8bit_type::value ? 16 : 32; - constexpr ck_tile::index_t K_Warp_Tile = is_8bit_type::value ? 64 : 16; + constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; + constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; + constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; -#elif defined(USING_MFMA_32x32x16) && defined(ENABLE_FP8) - constexpr ck_tile::index_t M_Tile = 128; - constexpr ck_tile::index_t N_Tile = 256; - constexpr ck_tile::index_t K_Tile = 128; - - constexpr ck_tile::index_t M_Warp = 1; - constexpr ck_tile::index_t N_Warp = 8; - constexpr ck_tile::index_t K_Warp = 1; - - constexpr ck_tile::index_t M_Warp_Tile = is_8bit_type::value ? 32 : 32; - constexpr ck_tile::index_t N_Warp_Tile = is_8bit_type::value ? 32 : 32; - constexpr ck_tile::index_t K_Warp_Tile = is_8bit_type::value ? 32 : 16; -#endif using CodegenFlatmmShape = ck_tile::TileFlatmmShape, ck_tile::sequence, @@ -134,47 +122,4 @@ float flatmm_calc(const ck_tile::FlatmmHostArgs& args, const ck_tile::stream_con #include "run_flatmm_example.inc" -int run_flatmm_example(int argc, char* argv[]) -{ - auto [result, arg_parser] = create_args(argc, argv); - if(!result) - return -1; - - using Row = ck_tile::tensor_layout::gemm::RowMajor; - using Col = ck_tile::tensor_layout::gemm::ColumnMajor; - - std::string data_type = arg_parser.get_str("prec"); - std::string a_layout = arg_parser.get_str("a_layout"); - std::string b_layout = arg_parser.get_str("b_layout"); - - if(a_layout == "R" && b_layout == "C") - { - if(data_type == "fp16") - { - run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); - } - else if(data_type == "bf16") - { - run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); - } - else if(data_type == "fp8") - { - run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); - } - else if(data_type == "bf8") - { - run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); - } - else - { - throw std::runtime_error("Unsupported data_type!"); - } - } - else - { - throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!"); - } - return -1; -} - int main(int argc, char* argv[]) { return !run_flatmm_example(argc, argv); } diff --git a/example/ck_tile/18_flatmm/flatmm_basic.hpp b/example/ck_tile/18_flatmm/flatmm_basic.hpp index bbce978724..9da3c4c143 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.hpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.hpp @@ -31,7 +31,137 @@ #error "unsupported CK_TILE_PIPELINE_DEFAULT value" #endif -template +template +struct GemmConfig +{ + 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 = 128 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + 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 = sizeof(DataType) == 2 ? 16 : 64; +}; + +template +struct GemmConfig +{ + 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 = 256 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + 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 = sizeof(DataType) == 2 ? 16 : 64; +}; + +template +struct GemmConfig +{ + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 128; + static constexpr ck_tile::index_t K_Tile = 128 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + 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 = sizeof(DataType) == 2 ? 16 : 64; +}; + +template +struct GemmConfig +{ + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 128; + static constexpr ck_tile::index_t K_Tile = 256 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + 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 = sizeof(DataType) == 2 ? 16 : 64; +}; + +template +struct GemmConfig +{ + 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 = 256 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 8; + 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 = sizeof(DataType) == 2 ? 16 : 64; +}; + +template +struct GemmConfig +{ + 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 = 128 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 8; + 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 = sizeof(DataType) == 2 ? 16 : 64; +}; + +template +struct GemmConfig +{ + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 128 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 8; + 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 = sizeof(DataType) == 2 ? 16 : 64; +}; + +template +struct GemmConfig +{ + 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 = 128 / sizeof(DataType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + 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 = sizeof(DataType) == 2 ? 32 : 128; +}; + +constexpr ck_tile::index_t GemmConfigNumVer = 8; + +template struct GemmBasicTypeConfig; template <> @@ -103,10 +233,10 @@ struct DataTypeTraits static constexpr const char* name = "fp16"; }; -template -struct is_8bit_type - : std::bool_constant || std::is_same_v> +template <> +struct DataTypeTraits { + static constexpr const char* name = "bf16"; }; auto create_args(int argc, char* argv[]) @@ -126,8 +256,9 @@ auto create_args(int argc, char* argv[]) .insert("warmup", "50", "number of iterations before benchmark the kernel") .insert("repeat", "100", "number of iterations to benchmark the kernel") .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer") - .insert("split_k", "1", "splitK value"); - + .insert("split_k", "1", "splitK value") + .insert("cfg_ver", "0", "gemm config version") + .insert("init", "0", "0:random, 1:linear, 2:constant(1)"); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); } diff --git a/example/ck_tile/18_flatmm/run_flatmm_example.inc b/example/ck_tile/18_flatmm/run_flatmm_example.inc index c191fff7d0..0f8c409115 100644 --- a/example/ck_tile/18_flatmm/run_flatmm_example.inc +++ b/example/ck_tile/18_flatmm/run_flatmm_example.inc @@ -32,38 +32,34 @@ static constexpr inline auto is_row_major(Layout layout_) } // mfma_type, 0:32x32, 1:16x16 -template -auto shuffle_b(const ck_tile::HostTensor& t, std::string mfma_dtype, int mfma_type) +template +auto shuffle_b(const ck_tile::HostTensor& t) { assert(t.get_lengths().size() == 2); int n_ = t.get_lengths()[1]; int k_ = t.get_lengths()[0]; - if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 0) + if constexpr(GemmConfig::N_Warp_Tile == 32) { - ck_tile::HostTensor t_view({n_ / 32, 32, k_ / 16, 2, 8}); + ck_tile::HostTensor t_view({n_ / 32, + 32, + k_ / GemmConfig::K_Warp_Tile, + 2, + GemmConfig::K_Warp_Tile / 2}); std::copy(t.begin(), t.end(), t_view.begin()); return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); } - else if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 1) + else { - ck_tile::HostTensor t_view({n_ / 16, 16, k_ / 32, 4, 8}); + static_assert(GemmConfig::N_Warp_Tile == 16); + ck_tile::HostTensor t_view({n_ / 16, + 16, + k_ / GemmConfig::K_Warp_Tile, + 4, + GemmConfig::K_Warp_Tile / 4}); std::copy(t.begin(), t.end(), t_view.begin()); return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); } - else if((mfma_dtype == "int8" || mfma_dtype == "fp8" || mfma_dtype == "bf8") && mfma_type == 0) - { - ck_tile::HostTensor t_view({n_ / 32, 32, k_ / 32, 2, 16}); - std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); - } - else if((mfma_dtype == "int8" || mfma_dtype == "fp8" || mfma_dtype == "bf8") && mfma_type == 1) - { - ck_tile::HostTensor t_view({n_ / 16, 16, k_ / 64, 4, 16}); - std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); - } - return t; } template @@ -91,6 +87,7 @@ template @@ -121,8 +118,14 @@ float invoke_flatmm(ck_tile::DeviceMem& a_dev_buf, args.stride_C = stride_C; float ave_time = - flatmm_calc( - args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); + flatmm_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 = @@ -138,7 +141,14 @@ float invoke_flatmm(ck_tile::DeviceMem& a_dev_buf, return ave_time; } -template +template int run_flatmm_example_with_layouts(int argc, char* argv[], const ALayout a_layout = ALayout{}, @@ -149,11 +159,6 @@ int run_flatmm_example_with_layouts(int argc, if(!result) return -1; - using ADataType = typename GemmBasicTypeConfig::ADataType; - using BDataType = typename GemmBasicTypeConfig::BDataType; - using CDataType = typename GemmBasicTypeConfig::CDataType; - using AccDataType = typename GemmBasicTypeConfig::AccDataType; - ck_tile::index_t M = arg_parser.get_int("m"); ck_tile::index_t N = arg_parser.get_int("n"); ck_tile::index_t K = arg_parser.get_int("k"); @@ -165,6 +170,7 @@ int run_flatmm_example_with_layouts(int argc, ck_tile::index_t kbatch = arg_parser.get_int("split_k"); int n_warmup = arg_parser.get_int("warmup"); int n_repeat = arg_parser.get_int("repeat"); + ck_tile::index_t init_method = arg_parser.get_int("init"); stride_A = ck_tile::get_default_stride(M, K, stride_A, is_row_major(a_layout)); stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout)); @@ -178,8 +184,26 @@ int run_flatmm_example_with_layouts(int argc, ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); // TODO: add different init types - ck_tile::FillUniformDistribution{-.5f, .5f}(a_host); - ck_tile::FillUniformDistribution{-.5f, .5f}(b_origin_host); + if(init_method == 0) + { + ck_tile::FillUniformDistribution{-.5f, .5f}(a_host); + ck_tile::FillUniformDistribution{-.5f, .5f}(b_origin_host); + } + else if(init_method == 1) + { + ck_tile::FillMonotonicSeq{}(a_host); + ck_tile::FillMonotonicSeq{}(b_origin_host); + } + else if(init_method == 2) + { + ck_tile::FillUniformDistribution{1.f, 1.f}(a_host); + ck_tile::FillUniformDistribution{1.f, 1.f}(b_origin_host); + } + else + { + a_host.SetZero(); + b_origin_host.SetZero(); + } ck_tile::DeviceMem a_dev_buf(a_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem c_dev_buf(c_rslt_host.get_element_space_size_in_bytes()); @@ -188,29 +212,29 @@ int run_flatmm_example_with_layouts(int argc, c_rslt_host.SetZero(); // do pre-shuffle - std::string mfma = arg_parser.get_str("prec"); -#if defined(USING_MFMA_16x16x32) && defined(ENABLE_FP8) - ck_tile::index_t mfma_type = 1; -#else - ck_tile::index_t mfma_type = 0; -#endif - ck_tile::HostTensor b_shuffle_host = shuffle_b(b_origin_host, mfma, mfma_type); + ck_tile::HostTensor b_shuffle_host = shuffle_b(b_origin_host); ck_tile::DeviceMem b_shuffle_dev_buf(b_shuffle_host.get_element_space_size_in_bytes()); b_shuffle_dev_buf.ToDevice(b_shuffle_host.data()); - invoke_flatmm( - a_dev_buf, - b_shuffle_dev_buf, - c_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat); + invoke_flatmm(a_dev_buf, + b_shuffle_dev_buf, + c_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + kbatch, + n_warmup, + n_repeat); c_dev_buf.FromDevice(c_rslt_host.data()); bool pass = true; @@ -300,3 +324,106 @@ int run_flatmm_example_with_layouts(int argc, return pass; } + +int run_flatmm_example(int argc, char* argv[]) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + return -1; + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + + std::string a_layout = arg_parser.get_str("a_layout"); + std::string b_layout = arg_parser.get_str("b_layout"); + std::string data_type = arg_parser.get_str("prec"); + ck_tile::index_t cfg_ver = arg_parser.get_int("cfg_ver"); + + if(a_layout == "R" && b_layout == "C") + { + if(cfg_ver >= GemmConfigNumVer) + { + throw std::runtime_error("Unsupported config version for this operation !!!"); + } + int ret = 0; + ck_tile::static_for<0, GemmConfigNumVer, 1>{}([&](auto ver) { + if(cfg_ver == ver) + { + if(data_type == "fp16") + { + using Types = GemmBasicTypeConfig; + using ADataType = Types::ADataType; + using BDataType = Types::BDataType; + using AccDataType = Types::AccDataType; + using CDataType = Types::CDataType; + ret = run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); + } + else if(data_type == "bf16") + { + using Types = GemmBasicTypeConfig; + using ADataType = Types::ADataType; + using BDataType = Types::BDataType; + using AccDataType = Types::AccDataType; + using CDataType = Types::CDataType; + ret = run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); + } + else if(data_type == "fp8") + { + + using Types = GemmBasicTypeConfig; + using ADataType = Types::ADataType; + using BDataType = Types::BDataType; + using AccDataType = Types::AccDataType; + using CDataType = Types::CDataType; + ret = run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); + } + else if(data_type == "bf8") + { + using Types = GemmBasicTypeConfig; + using ADataType = Types::ADataType; + using BDataType = Types::BDataType; + using AccDataType = Types::AccDataType; + using CDataType = Types::CDataType; + ret = run_flatmm_example_with_layouts(argc, argv, Row{}, Col{}, Row{}); + } + else + { + throw std::runtime_error("Unsupported data type for this operation !!!"); + } + } + }); + return ret; + } + else + { + throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!"); + } +} diff --git a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp index cbd20a6ea3..b767c53c26 100644 --- a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -116,6 +116,10 @@ struct FlatmmPipelineAGmemBGmemCRegV1 }); #elif defined(USING_MFMA_32x32x16) + ignore = A_Buffer_Load_Inst_Num; + ignore = A_LDS_Read_Inst_Num; + ignore = B_Buffer_Load_Inst_Num; +#if 0 static_for<0, A_LDS_Read_Inst_Num / 2 - A_Buffer_Load_Inst_Num - B_Buffer_Load_Inst_Num, 1>{}([&](auto i) { @@ -148,6 +152,7 @@ struct FlatmmPipelineAGmemBGmemCRegV1 __builtin_amdgcn_sched_group_barrier(0x008, 3, 0); // MFMA }); __builtin_amdgcn_sched_group_barrier(0x008, 4, 0); // MFMA +#endif #endif } diff --git a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index 1a1b729394..8e0f056fe8 100644 --- a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -138,6 +138,21 @@ struct UniversalFlatmmPipelineAgBgCrPolicy return Problem::VectorLoadSize / sizeof(typename Problem::ADataType); } + template + CK_TILE_HOST_DEVICE static constexpr auto GetK1() + { + using TileShape = typename Problem::BlockGemmShape; + if constexpr(TileShape::WarpTile::at(TileShape::idxN) == 32) + { + return TileShape::WarpTile::at(TileShape::idxK) / 2; + } + else + { + static_assert(TileShape::WarpTile::at(TileShape::idxN) == 16); + return TileShape::WarpTile::at(TileShape::idxK) / 4; + } + } + template CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution() { @@ -189,7 +204,7 @@ struct UniversalFlatmmPipelineAgBgCrPolicy } else { - constexpr index_t K1 = 16 / sizeof(ADataType); + constexpr index_t K1 = Problem::VectorLoadSize / sizeof(ADataType); constexpr index_t K0 = KPerBlock / K1; constexpr index_t M2 = get_warp_size() / K0; // coalesce reading for each blocks @@ -232,19 +247,17 @@ struct UniversalFlatmmPipelineAgBgCrPolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeBFlatDramTileDistribution() { - using BDataType = remove_cvref_t; - using TileShape = typename Problem::BlockGemmShape; // ck_tile::TileFlatmmShape constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNum = BlockSize / WaveSize; - constexpr index_t KBPerLoad = - Problem::VectorLoadSize / sizeof(BDataType); // dwordx4 load B elem cnt - constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim + constexpr index_t KBPerLoad = GetK1(); + constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim constexpr index_t KWavePerBlk = 1; constexpr index_t KRepeat = 1; + static_assert(TileShape::flatKPerWarp == KThdPerWave * KBPerLoad, "wrong"); constexpr index_t NBPerLoad = 1; constexpr index_t NThdPerWave = 1; diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp index be5d5690ff..c3247c298d 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp @@ -235,8 +235,16 @@ using WarpGemmMfma_f32_16x16x64_bf8_bf8 = WarpGemmImpl, 2>>; -using WarpGemmMfma_f32_16x16x128_fp8_fp8 = WarpGemmImpl>>; +using WarpGemmMfma_f32_16x16x32_bf8_bf8 = WarpGemmImpl< + WarpGemmAtrributeMfma>>; + +using WarpGemmMfma_f32_32x32x32_fp8_fp8 = WarpGemmImpl, + 2>>; + +using WarpGemmMfma_f32_32x32x32_bf8_bf8 = WarpGemmImpl, + 2>>; using WarpGemmMfma_f32_16x16x128_fp8_bf8 = WarpGemmImpl>>; @@ -244,6 +252,10 @@ using WarpGemmMfma_f32_16x16x128_fp8_bf8 = WarpGemmImpl>>; +#if defined(__gfx950__) +using WarpGemmMfma_f32_16x16x128_fp8_fp8 = WarpGemmImpl>>; + using WarpGemmMfma_f32_16x16x128_bf8_bf8 = WarpGemmImpl>>; @@ -258,6 +270,31 @@ using WarpGemmMfma_f32_32x32x64_bf8_fp8 = WarpGemmImpl< using WarpGemmMfma_f32_32x32x64_bf8_bf8 = WarpGemmImpl< WarpGemmAtrributeMfma>>; +#else +using WarpGemmMfma_f32_16x16x128_fp8_fp8 = WarpGemmImpl, + 4>>; + +using WarpGemmMfma_f32_16x16x128_bf8_bf8 = WarpGemmImpl, + 4>>; + +using WarpGemmMfma_f32_32x32x64_fp8_fp8 = WarpGemmImpl, + 4>>; + +using WarpGemmMfma_f32_32x32x64_fp8_bf8 = WarpGemmImpl, + 4>>; + +using WarpGemmMfma_f32_32x32x64_bf8_fp8 = WarpGemmImpl, + 4>>; + +using WarpGemmMfma_f32_32x32x64_bf8_bf8 = WarpGemmImpl, + 4>>; +#endif using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed = WarpGemmImpl