From c7d08b7c2ae4ad8264c0631addff5b620c6404fb Mon Sep 17 00:00:00 2001 From: coderfeli Date: Sun, 1 Dec 2024 04:50:58 +0000 Subject: [PATCH] use hasmainloop; no spill for 3tail --- example/ck_tile/03_gemm/gemm_basic.cpp | 4 +- .../core/tensor/tile_window_linear.hpp | 54 +---------- .../block/block_gemm_areg_breg_creg_v2.hpp | 6 +- .../gemm_pipeline_agmem_bgmem_creg_v1.hpp | 95 +++++++++++-------- .../gemm/pipeline/gemm_pipeline_problem.hpp | 2 + .../ops/gemm/pipeline/tile_gemm_traits.hpp | 6 +- 6 files changed, 70 insertions(+), 97 deletions(-) diff --git a/example/ck_tile/03_gemm/gemm_basic.cpp b/example/ck_tile/03_gemm/gemm_basic.cpp index 3cbf0c1ed9..4303bf87c9 100644 --- a/example/ck_tile/03_gemm/gemm_basic.cpp +++ b/example/ck_tile/03_gemm/gemm_basic.cpp @@ -33,7 +33,6 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) constexpr ck_tile::index_t M_Warp = 2; constexpr ck_tile::index_t N_Warp = 2; constexpr ck_tile::index_t K_Warp = 1; - constexpr ck_tile::index_t Warp_Size = 64; constexpr ck_tile::index_t M_Warp_Tile = 32; constexpr ck_tile::index_t N_Warp_Tile = 32; @@ -48,6 +47,7 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) ck_tile::sequence>; using TilePartitioner = ck_tile::GemmTilePartitioner; + // constexpr ck_tile::index_t Warp_Size = 64; // using GemmEpilogue = ck_tile::CShuffleEpilogueV2>; using CodegenGemmTraits = - ck_tile::TileGemmTraits; + ck_tile::TileGemmTraits; using CodegenPipelineProblem = ck_tile:: GemmPipelineProblem; using CodegenGemmPolicy = ck_tile::GemmPipelineAGmemBGmemCRegV1DefaultPolicy; diff --git a/include/ck_tile/core/tensor/tile_window_linear.hpp b/include/ck_tile/core/tensor/tile_window_linear.hpp index 70434dabc4..bfa96ae565 100644 --- a/include/ck_tile/core/tensor/tile_window_linear.hpp +++ b/include/ck_tile/core/tensor/tile_window_linear.hpp @@ -454,7 +454,7 @@ struct tile_window_linear CK_TILE_DEVICE constexpr auto get_num_of_access() const { return traits::NumAccess; } template - CK_TILE_DEVICE auto load(DistributedTensor dst_tensor, number = {}, bool_constant = {}) const + CK_TILE_DEVICE auto load(DistributedTensor& dst_tensor, number = {}, bool_constant = {}) const { using vector_t = typename traits::vector_t; using SFC_Ys = typename traits::SFC_Ys; @@ -508,56 +508,8 @@ struct tile_window_linear template CK_TILE_DEVICE auto load(number = {}, bool_constant = {}) const { - using vector_t = typename traits::vector_t; - using SFC_Ys = typename traits::SFC_Ys; - - constexpr auto tile_dstr = TileDstr{}; - - auto dst_tensor = make_static_distributed_tensor(tile_dstr); - - auto issue = [&](auto i_access_) { - constexpr auto IAccess = number{}; - - constexpr auto non_linear_id = number{}; - auto bottom_tensor_thread_coord = cached_coords_[non_linear_id]; - auto bottom_tensor_flag = cached_flags_[IAccess]; - - constexpr auto linear_offset = get_bottom_linear_offset(IAccess); - - // read from bottom tensor - const vector_t vec_value = - get_bottom_tensor_view().template get_vectorized_elements( - bottom_tensor_thread_coord, - linear_offset, - bottom_tensor_flag, - bool_constant{}); -#if 1 - // data index [y0, y1, ...] - constexpr auto idx_diff_ys = SFC_Ys::get_index(IAccess); - // write into distributed tensor - static_for<0, traits::ScalarPerVector, 1>{}([&](auto j) { - constexpr auto idx_ys = generate_tuple( - [&](auto jj) { - return jj == traits::VectorDimY ? (idx_diff_ys[jj] + j) : idx_diff_ys[jj]; - }, - number{}); - - constexpr index_t d = tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys); - - dst_tensor.get_thread_buffer().template at() = - vec_value.template get_as()[j]; - }); -#else - constexpr index_t d = tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys_start); - static_assert(d % traits::ScalarPerVector == 0); - - dst_tensor.get_thread_buffer().template get_as()( - number{}) = bit_cast(vec_value); -#endif - }; - - WINDOW_DISPATCH_ISSUE(); - + auto dst_tensor = make_static_distributed_tensor(TileDstr{}); + load(dst_tensor, number{}, bool_constant{}); return dst_tensor; } diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp index f0f41dda78..40d4bb94c9 100644 --- a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp @@ -222,13 +222,11 @@ struct BlockGemmARegBRegCRegV2 // Prefetch lds template - CK_TILE_DEVICE static auto PrefetchLds(const BlockWindow& block_window, BlockTensor& block_tensor) + CK_TILE_DEVICE static void PrefetchLds(const BlockWindow& block_window, BlockTensor& block_tensor) { auto tileDist = BlockTensor::get_tile_distribution(); - return load_tile(block_tensor, make_tile_window(block_window, tileDist)); - + load_tile(block_tensor, make_tile_window(block_window, tileDist)); // load_tile(block_tensor, make_tile_window_linear(block_window, tileDist)); - // return; } // C = A * B diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp index ed737d1221..bd87ec8bd2 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp @@ -36,6 +36,8 @@ struct GemmPipelineAGmemBGmemCRegV1 static constexpr bool kPadM = Problem::kPadM; static constexpr bool kPadN = Problem::kPadN; static constexpr bool kPadK = Problem::kPadK; + static constexpr bool kHasHotLoop = Problem::kHasHotLoop; + static constexpr auto kTailNum = Problem::kTailNum; // CK_TILE_HOST_DEVICE static constexpr index_t GetStaticLdsSize() // { @@ -131,6 +133,18 @@ struct GemmPipelineAGmemBGmemCRegV1 0x008, num_mfma_inst / num_issue - 3, 0); // MFMA : 5 }); __builtin_amdgcn_sched_barrier(0); + + // static_for<0, 8, 1>{}([&](auto i) { + // ignore = i; + // __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA : 1 + // __builtin_amdgcn_sched_group_barrier(0x100, 2, 0); // DS read : 2 + // __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA: 1 + // __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write : 1 + // __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA : 1 + // __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read :1 + // __builtin_amdgcn_sched_group_barrier(0x008, 5, 0); // MFMA : 5 + // }); + __builtin_amdgcn_sched_barrier(0); } CK_TILE_DEVICE static constexpr auto MakeCBlockSubTile() { @@ -261,60 +275,63 @@ struct GemmPipelineAGmemBGmemCRegV1 ALdsTile a_block_tile1; BLdsTile b_block_tile1; - while(iCounter > 1) - { - // ping + if (kHasHotLoop) { + do + { + // ping + { + block_sync_lds(); + Policy::template BlockGemm::PrefetchLds(a_lds_window1, a_block_tile1); + Policy::template BlockGemm::PrefetchLds(b_lds_window1, b_block_tile1); + LocalPrefill(a_lds_window0, a_global_load_tile, a_element_func); + LocalPrefill(b_lds_window0, b_global_load_tile, b_element_func); + GlobalPrefetch(a_global_load_tile, a_copy_dram_window); + GlobalPrefetch(b_global_load_tile, b_copy_dram_window); + block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + HotLoopScheduler(); + } + // pong + { + block_sync_lds(); + Policy::template BlockGemm::PrefetchLds(a_lds_window0, a_block_tile0); + Policy::template BlockGemm::PrefetchLds(b_lds_window0, b_block_tile0); + LocalPrefill(a_lds_window1, a_global_load_tile, a_element_func); + LocalPrefill(b_lds_window1, b_global_load_tile, b_element_func); + GlobalPrefetch(a_global_load_tile, a_copy_dram_window); + GlobalPrefetch(b_global_load_tile, b_copy_dram_window); + block_gemm(c_block_tile, a_block_tile1, b_block_tile1); + HotLoopScheduler(); + } + iCounter -= 2; + }while(iCounter > 1); + } + + //tail 3 + if (kTailNum == 3) { + // 3 { block_sync_lds(); Policy::template BlockGemm::PrefetchLds(a_lds_window1, a_block_tile1); Policy::template BlockGemm::PrefetchLds(b_lds_window1, b_block_tile1); LocalPrefill(a_lds_window0, a_global_load_tile, a_element_func); LocalPrefill(b_lds_window0, b_global_load_tile, b_element_func); - GlobalPrefetch(a_global_load_tile, a_copy_dram_window); - GlobalPrefetch(b_global_load_tile, b_copy_dram_window); block_gemm(c_block_tile, a_block_tile0, b_block_tile0); - HotLoopScheduler(); } - // pong + // 2 { block_sync_lds(); Policy::template BlockGemm::PrefetchLds(a_lds_window0, a_block_tile0); Policy::template BlockGemm::PrefetchLds(b_lds_window0, b_block_tile0); - LocalPrefill(a_lds_window1, a_global_load_tile, a_element_func); - LocalPrefill(b_lds_window1, b_global_load_tile, b_element_func); - GlobalPrefetch(a_global_load_tile, a_copy_dram_window); - GlobalPrefetch(b_global_load_tile, b_copy_dram_window); block_gemm(c_block_tile, a_block_tile1, b_block_tile1); - HotLoopScheduler(); } - iCounter -= 2; - } - - //tail 3 - // if (iCounter == 1) { - // // 3 - // { - // block_sync_lds(); - // Policy::template BlockGemm::PrefetchLds(a_lds_window1, a_block_tile1); - // Policy::template BlockGemm::PrefetchLds(b_lds_window1, b_block_tile1); - // LocalPrefill(a_lds_window0, a_global_load_tile, a_element_func); - // LocalPrefill(b_lds_window0, b_global_load_tile, b_element_func); - // block_gemm(c_block_tile, a_block_tile0, b_block_tile0); - // } - // // 2 - // { - // block_sync_lds(); - // Policy::template BlockGemm::PrefetchLds(a_lds_window0, a_block_tile0); - // Policy::template BlockGemm::PrefetchLds(b_lds_window0, b_block_tile0); - // block_gemm(c_block_tile, a_block_tile1, b_block_tile1); - // } - // //1 - // { - // block_gemm(c_block_tile, a_block_tile0, b_block_tile0); - // } - // //tail 2 - // } else + //1 + { + block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + } + } + else { + // //tail 2 { block_sync_lds(); Policy::template BlockGemm::PrefetchLds(a_lds_window1, a_block_tile1); diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp index 3c43790bd6..9a403797ff 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp @@ -32,6 +32,8 @@ struct GemmPipelineProblemBase static constexpr bool kPadM = GemmTraits::kPadM; static constexpr bool kPadN = GemmTraits::kPadN; static constexpr bool kPadK = GemmTraits::kPadK; + static constexpr bool kHasHotLoop = GemmTraits::HasHotLoop; + static constexpr auto kTailNum = GemmTraits::TailNum; CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentA() { diff --git a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp index 34756c3ff6..f5aff510bd 100644 --- a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp +++ b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp @@ -12,7 +12,9 @@ template + typename CLayout_, + bool HasHotLoop_, + index_t TailNum_> struct TileGemmTraits { static constexpr bool kPadM = kPadM_; @@ -24,6 +26,8 @@ struct TileGemmTraits using ALayout = ALayout_; using BLayout = BLayout_; using CLayout = CLayout_; + static constexpr bool HasHotLoop = HasHotLoop_; + static constexpr auto TailNum = TailNum_; }; } // namespace ck_tile