From 6cc03335e395fe9b45e22333ea3a64d8c5835330 Mon Sep 17 00:00:00 2001 From: Harisankar Sadasivan Date: Tue, 25 Jun 2024 22:37:55 +0000 Subject: [PATCH] added stream-k policies for 3 tile and 4 tile --- .../tensor_operation/gpu/grid/block_to_ctile_map.hpp | 10 ++++++++++ .../profiler/profile_gemm_universal_streamk_impl.hpp | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp b/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp index 6bde0cfb0e..eca523589d 100755 --- a/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp +++ b/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp @@ -1471,6 +1471,16 @@ struct BlockToCTileMap_GemmStreamK_v2 { sk_tiles = bigEnough ? (grid_size + num_tiles % grid_size) : num_tiles; } + else if(streamk_sel == 3) + { + sk_tiles = (num_tiles > (2 * grid_size)) ? (2 * grid_size + num_tiles % grid_size) + : num_tiles; + } + else if(streamk_sel == 4) + { + sk_tiles = (num_tiles > (3 * grid_size)) ? (3 * grid_size + num_tiles % grid_size) + : num_tiles; + } sk_num_blocks = sk_tiles; // if(sk_tiles < sk_num_blocks) // { diff --git a/profiler/include/profiler/profile_gemm_universal_streamk_impl.hpp b/profiler/include/profiler/profile_gemm_universal_streamk_impl.hpp index ba7b09a5b3..f8f56ea5a9 100755 --- a/profiler/include/profiler/profile_gemm_universal_streamk_impl.hpp +++ b/profiler/include/profiler/profile_gemm_universal_streamk_impl.hpp @@ -153,7 +153,7 @@ bool profile_gemm_universal_streamk_impl(int do_verification, for(auto& op_ptr : op_ptrs) { std::vector grid_size_list = {38, 76, 114, 152, 190, 228, 266, 304, 342, 380}; - std::vector streamk_sel_list = {0, 1, 2}; + std::vector streamk_sel_list = {0, 1, 2, 3, 4}; if(Grid_size == -1) {