From a1e7960552aa0420138dd61d49e55b1fce2beb90 Mon Sep 17 00:00:00 2001 From: Aleksander Dudek Date: Tue, 17 Jun 2025 12:30:18 -0500 Subject: [PATCH] Working --- .../gemm_pipeline_ag_bg_cr_comp_v5.hpp | 410 ++++++++++-------- 1 file changed, 225 insertions(+), 185 deletions(-) diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp index 828d4958ae..e52c8fcc18 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp @@ -19,6 +19,8 @@ struct BaseGemmPipelineAgBgCrCompV5 static constexpr index_t GlobalBufferNum = 2; static constexpr index_t HotloopUnroll = 2; + CK_TILE_HOST_DEVICE static constexpr auto TransposeC() { return Problem::TransposeC; } + CK_TILE_HOST static constexpr bool BlockHasHotloop(index_t num_loop) { return num_loop > PrefetchStages; @@ -84,7 +86,7 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB(); } // TODO check KRepeat - static constexpr index_t KRepeat = 2; + static constexpr index_t KRepeat = KPerBlock / GetSmemPackA(); static constexpr bool kPadM = Problem::kPadM; static constexpr bool kPadN = Problem::kPadN; @@ -114,10 +116,6 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 template <> struct PipelineImpl : public PipelineImplBase { - using Base = PipelineImplBase; - static constexpr index_t PrefetchStages = 3; - static constexpr index_t HotloopUnroll = 2; - CK_TILE_DEVICE static constexpr auto HotLoopScheduler() { constexpr index_t MPerXDL = BlockGemmShape::WarpTile::at(I0{}); @@ -316,28 +314,97 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1{}]), "B block window has incorrect lengths for defined BLayout!"); - ////////////// global window & register ///////////////// + ////////////// LDS desc, window & register ///////////////// + using aldstype = + remove_cvref_t; + using bldstype = + remove_cvref_t; + auto&& xd = PipelineImplBase::GetABLdsTensorViews(p_smem); + aldstype& a_lds_block = xd.at(I0{}); + bldstype& b_lds_block = xd.at(I1{}); + // tie(a_lds_block, b_lds_block) = PipelineImplBase::GetABLdsTensorViews(p_smem); + + // Tile distribution for load from lds + constexpr auto a_lds_load_tile_distr = + make_static_tile_distribution(BlockGemm::MakeABlockDistributionEncode()); + constexpr auto b_lds_load_tile_distr = + make_static_tile_distribution(BlockGemm::MakeBBlockDistributionEncode()); + // A DRAM tile window for load - auto a_copy_dram_window = - make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - a_dram_block_window_tmp.get_window_origin(), - Policy::template MakeADramTileDistribution()); + // A LDS tile window for store + // A LDS tile for block GEMM + + using acopy_dram_type = + remove_cvref_t; + using bcopy_dram_type = + remove_cvref_t; + + using a_copy_lds_window_type = + remove_cvref_t; + using b_copy_lds_window_type = + remove_cvref_t; + + using a_lds_load_tile_distr_type = + remove_cvref_t; + using b_lds_load_tile_distr_type = + remove_cvref_t; + + // tie(a_copy_dram_window, a_copy_lds_window, a_lds_gemm_window) = + // PipelineImplBase::GetAWindows( + // a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr); + + auto&& aWindows = PipelineImplBase::GetAWindows( + a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr); + auto&& bWindows = PipelineImplBase::GetBWindows( + b_dram_block_window_tmp, b_lds_block, b_lds_load_tile_distr); + + acopy_dram_type& a_copy_dram_window = aWindows.at(I0{}); + a_copy_lds_window_type& a_copy_lds_window = aWindows.at(I1{}); + a_lds_load_tile_distr_type& a_lds_gemm_window = aWindows.at(I2{}); + + bcopy_dram_type& b_copy_dram_window = bWindows.at(I0{}); + b_copy_lds_window_type& b_copy_lds_window = bWindows.at(I1{}); + b_lds_load_tile_distr_type& b_lds_gemm_window = bWindows.at(I2{}); // B DRAM tile window for load - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_dram_block_window_tmp.get_window_origin(), - Policy::template MakeBDramTileDistribution()); + // B LDS tile window for store + // B LDS tile for block GEMM + // tie(b_copy_dram_window, b_copy_lds_window, b_lds_gemm_window) = + // PipelineImplBase::GetBWindows( + // b_dram_block_window_tmp, b_lds_block, b_lds_load_tile_distr); - // A register tile for global load - constexpr auto ABlockTileDistr = a_copy_dram_window.get_tile_distribution(); - constexpr auto BBlockTileDistr = b_copy_dram_window.get_tile_distribution(); - using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr)); - using BBlockTile = decltype(make_static_distributed_tensor(BBlockTileDistr)); - ABlockTile a_global_load_tile; - BBlockTile b_global_load_tile; + // Block GEMM + auto block_gemm = BlockGemm(); + auto c_block_tile = block_gemm.MakeCBlockTile(); + + using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); + using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); + + using ABlockTile = + decltype(make_static_distributed_tensor(ABlockTileDistr{})); + using BBlockTile = + decltype(make_static_distributed_tensor(BBlockTileDistr{})); + + ABlockTile a_block_tile[Base::GlobalBufferNum]; // TODO GlobalBufferNum + BBlockTile b_block_tile[Base::GlobalBufferNum]; // TODO GlobalBufferNum using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; @@ -347,61 +414,6 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 constexpr BDramTileWindowStep b_dram_tile_window_step = is_b_row_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock); - // Global prefetch 1 - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); - - ////////////// LDS desc, window & register ///////////////// - auto&& [a_lds_block0, b_lds_block0] = Base::GetABLdsTensorViews(p_smem); - - auto a_copy_lds_window0 = make_tile_window( - a_lds_block0, make_tuple(number{}, number{}), {0, 0}); - - auto b_copy_lds_window0 = make_tile_window( - b_lds_block0, make_tuple(number{}, number{}), {0, 0}); - - // Local prefill 1 - if constexpr(is_a_col_major) - { - auto a_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); - } - else - { - Base::LocalPrefill(a_copy_lds_window0, a_global_load_tile, a_element_func); - } - if constexpr(is_b_row_major) - { - auto b_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); - } - else - { - Base::LocalPrefill(b_copy_lds_window0, b_global_load_tile, b_element_func); - } - - // Global prefetch 2 - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); - - // Global prefetch 3 - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); - - // Block GEMM - auto block_gemm = BlockGemm(); - auto c_block_tile = block_gemm.MakeCBlockTile(); - - // Initialize C - tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); - - block_sync_lds(); - - // Prepare block tiles and windows constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution( BlockGemm::MakeABlockDistributionEncode())){}; constexpr auto BLdsTileDistr = decltype(make_static_tile_distribution( @@ -410,161 +422,189 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 using ALdsTile = decltype(make_static_distributed_tensor(ALdsTileDistr)); using BLdsTile = decltype(make_static_distributed_tensor(BLdsTileDistr)); - ALdsTile a_block_tile0; - BLdsTile b_block_tile0; + ALdsTile a_lds_tile; + BLdsTile b_lds_tile; + // ----------------------------------------------------------------------------------------- + // Gemm pipeline start - auto a_lds_ld_window0 = - make_tile_window(a_lds_block0, - make_tuple(number{}, number{}), - {0, 0}, - ALdsTileDistr); + // Global prefetch 1 + PipelineImplBase::GlobalPrefetch( + a_block_tile[I0{}], a_copy_dram_window, a_dram_tile_window_step); + PipelineImplBase::GlobalPrefetch( + b_block_tile[I0{}], b_copy_dram_window, b_dram_tile_window_step); - auto b_lds_ld_window0 = - make_tile_window(b_lds_block0, - make_tuple(number{}, number{}), - {0, 0}, - BLdsTileDistr); + // initialize C + tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + + // Local prefill 1 + if constexpr(is_a_col_major) + { + auto a_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledARegTileDistribution()); + transpose_tile2d(a_shuffle_tmp, a_block_tile[I0{}]); + PipelineImplBase::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + } + else + { + PipelineImplBase::LocalPrefill( + a_copy_lds_window, a_block_tile[I0{}], a_element_func); + } + if constexpr(is_b_row_major) + { + auto b_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledBRegTileDistribution()); + transpose_tile2d(b_shuffle_tmp, b_block_tile[I0{}]); + PipelineImplBase::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + } + else + { + PipelineImplBase::LocalPrefill( + b_copy_lds_window, b_block_tile[I0{}], b_element_func); + } + + // Global prefetch 2 + PipelineImplBase::GlobalPrefetch( + a_block_tile[I0{}], a_copy_dram_window, a_dram_tile_window_step); + PipelineImplBase::GlobalPrefetch( + b_block_tile[I0{}], b_copy_dram_window, b_dram_tile_window_step); + + // Global prefetch 3 + PipelineImplBase::GlobalPrefetch( + a_block_tile[I1{}], a_copy_dram_window, a_dram_tile_window_step); + PipelineImplBase::GlobalPrefetch( + b_block_tile[I1{}], b_copy_dram_window, b_dram_tile_window_step); + + block_sync_lds(); // Local prefetch 1 - Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); - Base::LocalPrefetch(b_block_tile0, b_lds_ld_window0); + PipelineImplBase::LocalPrefetch(a_lds_tile, a_lds_gemm_window); + PipelineImplBase::LocalPrefetch(b_lds_tile, b_lds_gemm_window); if(HasHotLoop) { index_t i = 0; do { - auto LoopFunc = [&]() { - static_for<0, KRepeat, 1>{}([&](auto k0) { - if constexpr(k0 == (KRepeat - 1)) - { - block_sync_lds(); - - // Local prefill 2 - if constexpr(is_a_col_major) - { - auto a_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledARegTileDistribution< - Problem>()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill( - a_copy_lds_window0, a_shuffle_tmp, a_element_func); - } - else - { - Base::LocalPrefill( - a_copy_lds_window0, a_global_load_tile, a_element_func); - } - if constexpr(is_b_row_major) - { - auto b_shuffle_tmp = make_static_distributed_tensor( - Policy::template MakeShuffledBRegTileDistribution< - Problem>()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill( - b_copy_lds_window0, b_shuffle_tmp, b_element_func); - } - else - { - Base::LocalPrefill( - b_copy_lds_window0, b_global_load_tile, b_element_func); - } - - // Global prefetch 4 - Base::GlobalPrefetch(a_global_load_tile, - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, - b_copy_dram_window, - b_dram_tile_window_step); - - block_sync_lds(); - - block_gemm(c_block_tile, a_block_tile0, b_block_tile0); - - // Local prefetch 2 - Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); - Base::LocalPrefetch(b_block_tile0, b_lds_ld_window0); - } - }); - - HotLoopScheduler(); - }; - - LoopFunc(); - LoopFunc(); - - i += HotloopUnroll; - } while(i < (num_loop - PrefetchStages)); - } - - // tail - auto ReadWriteCompFunc = [&]() { - static_for<0, KRepeat, 1>{}([&](auto k0) { - if constexpr(k0 == (KRepeat - 1)) - { + auto LoopFunc = [&](auto vmem_buf_idx) { block_sync_lds(); - // Local prefill 3 + // Local prefill 2 if constexpr(is_a_col_major) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, a_block_tile[vmem_buf_idx]); + PipelineImplBase::LocalPrefill( + a_copy_lds_window, a_shuffle_tmp, a_element_func); } else { - Base::LocalPrefill( - a_copy_lds_window0, a_global_load_tile, a_element_func); + PipelineImplBase::LocalPrefill( + a_copy_lds_window, a_block_tile[vmem_buf_idx], a_element_func); } if constexpr(is_b_row_major) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, b_block_tile[vmem_buf_idx]); + PipelineImplBase::LocalPrefill( + b_copy_lds_window, b_shuffle_tmp, b_element_func); } else { - Base::LocalPrefill( - b_copy_lds_window0, b_global_load_tile, b_element_func); + PipelineImplBase::LocalPrefill( + b_copy_lds_window, b_block_tile[vmem_buf_idx], b_element_func); } - block_sync_lds(); - } - block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + // Global prefetch 4 + PipelineImplBase::GlobalPrefetch(a_block_tile[vmem_buf_idx], + a_copy_dram_window, + a_dram_tile_window_step); + PipelineImplBase::GlobalPrefetch(b_block_tile[vmem_buf_idx], + b_copy_dram_window, + b_dram_tile_window_step); - // Local prefetch 3 - Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); - Base::LocalPrefetch(b_block_tile0, b_lds_ld_window0); - }); + block_sync_lds(); + + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); + // Policy::template GetVectorSizeA(); + + // Local prefetch 2 + PipelineImplBase::LocalPrefetch(a_lds_tile, a_lds_gemm_window); + PipelineImplBase::LocalPrefetch(b_lds_tile, b_lds_gemm_window); + + HotLoopScheduler(); + __builtin_amdgcn_sched_barrier(0); + }; + + LoopFunc(I0{}); + LoopFunc(I1{}); + + i += Base::HotloopUnroll; + } while(i < (num_loop - Base::PrefetchStages)); + } + + auto ReadWriteCompFunc = [&](auto vmem_buf_idx) { + block_sync_lds(); + + // Local prefill 3 + if constexpr(is_a_col_major) + { + auto a_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledARegTileDistribution()); + transpose_tile2d(a_shuffle_tmp, a_block_tile[vmem_buf_idx]); + PipelineImplBase::LocalPrefill( + a_copy_lds_window, a_shuffle_tmp, a_element_func); + } + else + { + PipelineImplBase::LocalPrefill( + a_copy_lds_window, a_block_tile[vmem_buf_idx], a_element_func); + } + if constexpr(is_b_row_major) + { + auto b_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledBRegTileDistribution()); + transpose_tile2d(b_shuffle_tmp, b_block_tile[vmem_buf_idx]); + PipelineImplBase::LocalPrefill( + b_copy_lds_window, b_shuffle_tmp, b_element_func); + } + else + { + PipelineImplBase::LocalPrefill( + b_copy_lds_window, b_block_tile[vmem_buf_idx], b_element_func); + } + + block_sync_lds(); + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); + + PipelineImplBase::LocalPrefetch(a_lds_tile, a_lds_gemm_window); + PipelineImplBase::LocalPrefetch(b_lds_tile, b_lds_gemm_window); HotLoopScheduler(); }; auto ReadCompFunc = [&]() { - block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); // Local prefetch 4 - Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); - Base::LocalPrefetch(b_block_tile0, b_lds_ld_window0); + PipelineImplBase::LocalPrefetch(a_lds_tile, a_lds_gemm_window); + PipelineImplBase::LocalPrefetch(b_lds_tile, b_lds_gemm_window); - // TODO verify - block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); HotLoopScheduler(); }; if constexpr(TailNum == TailNumber::Odd) { - ReadWriteCompFunc(); - ReadWriteCompFunc(); + ReadWriteCompFunc(I0{}); + ReadWriteCompFunc(I1{}); ReadCompFunc(); } else if constexpr(TailNum == TailNumber::Even) { - ReadWriteCompFunc(); + ReadWriteCompFunc(I0{}); ReadCompFunc(); }