diff --git a/example/ck_tile/18_flatmm/flatmm_basic.cpp b/example/ck_tile/18_flatmm/flatmm_basic.cpp index 60dced003b..2fecdc1663 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.cpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.cpp @@ -50,35 +50,41 @@ 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]; - constexpr int divisor = FlatmmConfig::N_Warp_Tile == 32 ? 2 : 4; + int n_ = t.get_lengths()[1]; + int k_ = t.get_lengths()[0]; + + constexpr int MaxVecSize = 16 / sizeof(T); + constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane); + ck_tile::HostTensor t_view({n_ / FlatmmConfig::N_Warp_Tile, FlatmmConfig::N_Warp_Tile, - k_ / FlatmmConfig::K_Warp_Tile, - divisor, - FlatmmConfig::K_Warp_Tile / divisor}); + k_ / ItemsPerAccess, + ItemsPerAccess}); std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + return ck_tile::reference_permute(t_view, {0, 2, 1, 3}); } template auto shuffle_b_v1(const ck_tile::HostTensor& t) { assert(t.get_lengths().size() == 2); - int n_ = t.get_lengths()[1]; - int k_ = t.get_lengths()[0]; - constexpr int divisor = FlatmmConfig::N_Warp_Tile == 32 ? 2 : 4; + int n_ = t.get_lengths()[1]; + int k_ = t.get_lengths()[0]; + + constexpr int MaxVecSize = 16 / sizeof(T); + constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane); constexpr int NRepeat = FlatmmConfig::N_Tile / FlatmmConfig::N_Warp_Tile / FlatmmConfig::N_Warp; + ck_tile::HostTensor t_view({n_ / FlatmmConfig::N_Tile, FlatmmConfig::N_Warp, FlatmmConfig::N_Warp_Tile, NRepeat, - k_ / FlatmmConfig::K_Warp_Tile, - divisor, - FlatmmConfig::K_Warp_Tile / divisor}); + k_ / ItemsPerAccess, + ItemsPerAccess}); std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 5, 2, 6}); + return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 2, 5}); } template diff --git a/example/ck_tile/18_flatmm/flatmm_basic.hpp b/example/ck_tile/18_flatmm/flatmm_basic.hpp index 063a981daf..20b9c6ee15 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.hpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.hpp @@ -81,7 +81,7 @@ struct FlatmmConfig16 static constexpr bool DoubleSmemBuffer = false; static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp; - static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; + static constexpr bool TiledMMAPermuteN = N_Repeat % 4 == 0; }; template @@ -94,7 +94,7 @@ struct FlatmmConfig16_950 : public FlatmmConfig16 static constexpr int N_Repeat = N_Tile / FlatmmConfig16::N_Warp_Tile / FlatmmConfig16::N_Warp; - static constexpr bool TiledMMAPermuteN = N_Repeat % 2 == 0; + static constexpr bool TiledMMAPermuteN = N_Repeat % 4 == 0; }; template diff --git a/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp b/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp index d2b9e666de..e002dcc57a 100644 --- a/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp +++ b/example/ck_tile/21_moe_flatmm/moe_flatmm.cpp @@ -35,17 +35,16 @@ auto shuffle_b(const ck_tile::HostTensor& t) int n_ = t.get_lengths()[1]; int k_ = t.get_lengths()[0]; - constexpr int N_Warp_Tile = FlatmmConfig::N_Warp_Tile; - constexpr int N_Warp = FlatmmConfig::N_Warp; - constexpr int KPerLane = FlatmmConfig::K_Warp_Tile / (64 / N_Warp_Tile); + constexpr int MaxVecSize = 16 / sizeof(T); + constexpr int KLane = ck_tile::get_warp_size() / FlatmmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = std::min(MaxVecSize, FlatmmConfig::K_Warp_Tile / KLane); - ck_tile::HostTensor t_view({n_ / N_Warp_Tile, - N_Warp_Tile, - k_ / (64 * KPerLane / N_Warp_Tile), - 64 / N_Warp_Tile, - KPerLane}); + ck_tile::HostTensor t_view({n_ / FlatmmConfig::N_Warp_Tile, + FlatmmConfig::N_Warp_Tile, + k_ / ItemsPerAccess, + ItemsPerAccess}); std::copy(t.begin(), t.end(), t_view.begin()); - return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); + return ck_tile::reference_permute(t_view, {0, 2, 1, 3}); } template 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 ee88be1466..ceb6ef6734 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 @@ -565,13 +565,13 @@ defined(USING_MFMA_32x32x64) && defined(ENABLE_FP4) // mi350 fp4 32c 1*K1 make_tile_window(a_lds_block_ping, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); auto a_warp_window_pong_tmp = make_tile_window(a_lds_block_pong, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); statically_indexed_array< statically_indexed_array, @@ -586,16 +586,10 @@ defined(USING_MFMA_32x32x64) && defined(ENABLE_FP4) // mi350 fp4 32c 1*K1 static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { a_warp_windows_ping(mIter)(kIter) = a_warp_window_ping_tmp; + a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp; move_tile_window(a_warp_windows_ping(mIter)(kIter), {mIter * MPerBlockPerIter, kIter * KPerBlockPerIter}); - }); - }); - - static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - a_warp_windows_pong(mIter)(kIter) = a_warp_window_pong_tmp; - move_tile_window(a_warp_windows_pong(mIter)(kIter), {mIter * MPerBlockPerIter, kIter * KPerBlockPerIter}); }); 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 f546cd7bac..cc882db8d6 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 @@ -250,6 +250,37 @@ struct UniversalFlatmmPipelineAgBgCrPolicy } } + template + CK_TILE_HOST_DEVICE static constexpr auto MakeALDS_WarpTileDistribution() + { + using TileShape = typename Problem::BlockGemmShape; + using ADataType = remove_cvref_t; + using ALayout = remove_cvref_t; + + static_assert(TileShape::BlockWarps::at(I0) == 1, "requires Wave_M == 1"); + + constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); + constexpr index_t KPerXdl = Problem::BlockGemmShape::WarpTile::at(I2); + + constexpr int Repeat = TileShape::BlockWarps::at(number<1>{}); + + constexpr int KLane = get_warp_size() / MPerXdl; + constexpr int KPerThread = KPerXdl / KLane; + + constexpr int MaxVecSize = 16 / sizeof(ADataType); + constexpr int KItemsPerLoad = min(MaxVecSize, KPerThread); + constexpr int KFragment = KPerThread / KItemsPerLoad; + + return make_static_tile_distribution( + tile_distribution_encoding< + sequence, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<1, 0>>, + sequence<2, 2>, + sequence<0, 2>>{}); + } + template CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution() { @@ -303,10 +334,10 @@ struct UniversalFlatmmPipelineAgBgCrPolicy { 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 - if constexpr(get_warp_size() % (M2 * K0) == 0) + if constexpr(get_warp_size() % K0 == 0) { + constexpr index_t M2 = get_warp_size() / K0; constexpr index_t M1 = BlockSize / get_warp_size(); static_assert(M2 != 0, "M2 is zero, which will lead to a division by zero error."); static_assert(M1 != 0, "M1 is zero, which will lead to a division by zero error."); @@ -325,18 +356,18 @@ struct UniversalFlatmmPipelineAgBgCrPolicy } else { - constexpr index_t M0 = BlockSize / get_warp_size(); - constexpr index_t M1 = MPerBlock / (M2 * M0); - static_assert(M0 * M1 * M2 == MPerBlock, - "Incorrect M0, M1, M2 configuration! " - "M0, M1, M2 must cover whole MPerBlock!"); + constexpr index_t KWave = K0 / get_warp_size(); + constexpr index_t M0 = BlockSize / get_warp_size() / KWave; + constexpr index_t M1 = MPerBlock / M0; + return make_static_tile_distribution( - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<1, 2>>, - tuple, sequence<2, 0>>, - sequence<1, 2>, - sequence<1, 1>>{}); + tile_distribution_encoding< + sequence<1>, + tuple, sequence>, + tuple, sequence<2>>, + tuple, sequence<1>>, + sequence<1, 2>, + sequence<1, 2>>{}); } } } @@ -381,11 +412,17 @@ struct UniversalFlatmmPipelineAgBgCrPolicy constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNum = BlockSize / WaveSize; - constexpr index_t KBPerLoad = GetKBPerLoad(); - constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim + constexpr index_t KBPerLoad = GetKBPerLoad(); + + constexpr index_t MaxVecSize = 16 / sizeof(typename Problem::BDataType); + constexpr index_t KItemsPerLoad = min(KBPerLoad, MaxVecSize); + constexpr index_t KFragment = KBPerLoad / KItemsPerLoad; + static_assert(KFragment * KItemsPerLoad == KBPerLoad); + + 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"); + static_assert(TileShape::BlockWarps::at(number<2>{}) == 1, "Requires K_Warp == 1"); constexpr index_t NBPerLoad = 1; constexpr index_t NThdPerWave = 1; @@ -396,9 +433,10 @@ struct UniversalFlatmmPipelineAgBgCrPolicy return make_static_tile_distribution( tile_distribution_encoding< - sequence, // ? - tuple, // second direction - sequence>, // first direction + sequence, // ? + tuple, // second direction + sequence>, // first + // direction // wave in blk, // thd in wave // // tuple, sequence<1, 2>>, // which direction diff --git a/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp b/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp index cd91f33a79..7a47bd7cd2 100644 --- a/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp +++ b/include/ck_tile/ops/moe_flatmm/pipeline/moe_flatmm_pipeline_agmem_bgmem_creg.hpp @@ -508,13 +508,13 @@ struct MoeFlatmmPipelineAGmemBGmemCRegV1 make_tile_window(a_lds_block_ping, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); auto a_warp_window_pong_tmp = make_tile_window(a_lds_block_pong, make_tuple(number{}, number{}), {iMWarp * WG::kM, 0}, - make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + PipelinePolicy::template MakeALDS_WarpTileDistribution()); statically_indexed_array< statically_indexed_array,