diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp index b67058b24a..ae07014695 100644 --- a/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/example/ck_tile/03_gemm/gemm_utils.hpp @@ -179,26 +179,6 @@ struct GemmConfigComputeV4_1 : public GemmConfigBase template struct GemmConfigComputeV5 : public GemmConfigBase -{ - 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 = 64 / sizeof(PrecType); - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 1; - static constexpr ck_tile::index_t K_Warp = 2; - - 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(PrecType) == 2 ? 16 : 64; - - static constexpr bool DoubleSmemBuffer = false; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V5; - static constexpr ck_tile::index_t NumWaNumWaveGroups = 2; -}; - -template -struct GemmConfigComputeV6 : public GemmConfigBase { static constexpr ck_tile::index_t M_Tile = 256; static constexpr ck_tile::index_t N_Tile = 256; @@ -213,7 +193,27 @@ struct GemmConfigComputeV6 : public GemmConfigBase static constexpr ck_tile::index_t K_Warp_Tile = 16; static constexpr bool DoubleSmemBuffer = false; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V6; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V5; +}; + +template +struct GemmConfigComputeV6 : public GemmConfigBase +{ + 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 = 64 / sizeof(PrecType); + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 1; + static constexpr ck_tile::index_t K_Warp = 2; + + 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(PrecType) == 2 ? 16 : 64; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V6; + static constexpr ck_tile::index_t NumWaNumWaveGroups = 2; }; template diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index bb536492a1..da06e50a58 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -324,7 +324,7 @@ int main(int argc, char* argv[]) { try { - return !run_gemm_example(argc, argv); + return !run_gemm_example(argc, argv); } catch(const std::runtime_error& e) { 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 55220730cd..58e7075b0b 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 @@ -1,49 +1,83 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +#pragma once #include "ck_tile/core.hpp" -#include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" -#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp" -#include "ck_tile/host/concat.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp" namespace ck_tile { -// A Tile Window: global memory -// B Tile Window: global memory -// C Distributed Tensor: register +// A Tile Window: global memory +// B Tile Window: global memory +// C Distributed tensor: register template struct BaseGemmPipelineAgBgCrCompV5 { - static constexpr index_t PrefetchStages = 1; + static constexpr index_t PrefetchStages = 3; static constexpr index_t PrefillStages = 1; - static constexpr index_t GlobalBufferNum = 1; + 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_DEVICE static constexpr bool BlockHasHotloop(index_t) { return true; } - - CK_TILE_HOST_DEVICE static constexpr TailNumber GetBlockLoopTailNum(index_t) + CK_TILE_HOST static constexpr bool BlockHasHotloop(index_t num_loop) { - return TailNumber::Empty; + return num_loop > PrefetchStages; + } + + CK_TILE_HOST static constexpr TailNumber GetBlockLoopTailNum(index_t num_loop) + { + if(num_loop % HotloopUnroll == 1) + { + return TailNumber::Odd; + } + else + { + return TailNumber::Even; + } } template - CK_TILE_HOST_DEVICE static auto TailHandler(const RunFunction& run_func, bool, TailNumber) + CK_TILE_HOST_DEVICE static auto TailHandler(const RunFunction& run_func, + [[maybe_unused]] bool has_hot_loop, + TailNumber tail_num) { - return run_func(bool_constant{}, integral_constant{}); + if(tail_num == ck_tile::TailNumber::Three) + { + return run_func( + ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else + { + return run_func( + ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } } }; +/** + * @brief Compute optimized pipeline version 5 + */ template struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 { using Base = BaseGemmPipelineAgBgCrCompV5; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + + static_assert(!std::is_same_v, "Not implemented"); + + static constexpr index_t APackedSize = + ck_tile::numeric_traits>::PackedSize; + static constexpr index_t BPackedSize = + ck_tile::numeric_traits>::PackedSize; using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; @@ -51,10 +85,10 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; - using BlockGemm = remove_cvref_t())>; - using I0 = number<0>; - using I1 = number<1>; - using I2 = number<2>; + using BlockGemm = remove_cvref_t())>; + static constexpr auto I0 = number<0>{}; + static constexpr auto I1 = number<1>{}; + static constexpr auto I2 = number<2>{}; static constexpr index_t BlockSize = Problem::kBlockSize; @@ -66,6 +100,12 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 static constexpr index_t GetVectorSizeB() { return Policy::template GetVectorSizeB(); } static constexpr index_t GetVectorSizeC() { return Policy::template GetVectorSizeC(); } + static constexpr index_t GetSmemPackA() { return Policy::template GetSmemPackA(); } + static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB(); } + + // TODO check KRepeat + static constexpr index_t KRepeat = KPerBlock / GetSmemPackA(); + static constexpr bool kPadM = Problem::kPadM; static constexpr bool kPadN = Problem::kPadN; static constexpr bool kPadK = Problem::kPadK; @@ -76,9 +116,6 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 static constexpr auto TailNum = Problem::TailNum; static constexpr auto Scheduler = Problem::Scheduler; - static constexpr index_t NumWarps = BlockGemmShape::NumWarps; - static constexpr index_t KTileSize = BlockGemmShape::WarpTile::at(I2{}); - [[nodiscard]] CK_TILE_HOST static const std::string GetName() { // clang-format off @@ -106,20 +143,180 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 template <> struct PipelineImpl : public PipelineImplBase { - using Base = PipelineImplBase; + CK_TILE_DEVICE static constexpr auto HotLoopScheduler() + { + constexpr index_t MPerXDL = BlockGemmShape::WarpTile::at(I0); + constexpr index_t NPerXDL = BlockGemmShape::WarpTile::at(I1); + constexpr index_t KPerXDL = BlockGemmShape::WarpTile::at(I2); + + constexpr index_t WaveSize = 64; + constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(I0); + constexpr index_t WaveNumN = BlockGemmShape::BlockWarps::at(I1); + + constexpr index_t A_LDS_Read_Width = KPerXDL; + constexpr index_t B_LDS_Read_Width = KPerXDL; + + constexpr index_t A_Buffer_Load_Inst_Num = + MPerBlock * KPerBlock / (BlockSize * GetVectorSizeA()); + constexpr index_t B_Buffer_Load_Inst_Num = + NPerBlock * KPerBlock / (BlockSize * GetVectorSizeB()); + + constexpr index_t A_LDS_Write_Inst_Num = MPerBlock * KPerBlock / (BlockSize * KPerXDL); + constexpr index_t B_LDS_Write_Inst_Num = NPerBlock * KPerBlock / (BlockSize * KPerXDL); + + constexpr index_t A_LDS_Read_Inst_Num = + WaveNumN * MPerBlock * KPerBlock / (BlockSize * KPerXDL); + constexpr index_t B_LDS_Read_Inst_Num = + WaveNumM * NPerBlock * KPerBlock / (BlockSize * KPerXDL); + + constexpr index_t C_MFMA_Inst_Num = MPerBlock * NPerBlock * KPerBlock / + (BlockSize / WaveSize) / + (MPerXDL * NPerXDL * KPerXDL); + + constexpr auto num_ds_read_inst_a = + A_LDS_Read_Width * sizeof(ADataType) / APackedSize == 16 ? A_LDS_Read_Inst_Num + : A_LDS_Read_Inst_Num / 2; + constexpr auto num_ds_read_inst_b = + B_LDS_Read_Width * sizeof(BDataType) / BPackedSize == 16 ? B_LDS_Read_Inst_Num + : B_LDS_Read_Inst_Num / 2; + + constexpr auto mfma_cycle = NPerXDL == 16 ? 16 : 32; + constexpr auto ds_read_a_issue_cycle = + A_LDS_Read_Width * sizeof(ADataType) == 16 ? 8 : 4; + constexpr auto ds_read_b_issue_cycle = + B_LDS_Read_Width * sizeof(BDataType) == 16 ? 8 : 4; + constexpr auto ds_read_a_mfma_rate = + (mfma_cycle - 4 + 2 * ds_read_a_issue_cycle - 1) / (2 * ds_read_a_issue_cycle); + constexpr auto ds_read_b_mfma_rate = + (mfma_cycle - 4 + 2 * ds_read_b_issue_cycle - 1) / (2 * ds_read_b_issue_cycle); + + constexpr auto num_dsread_stage1_a = num_ds_read_inst_a / KRepeat * (KRepeat - 1); + constexpr auto num_dsread_stage1_b = num_ds_read_inst_b / KRepeat * (KRepeat - 1); + constexpr auto num_dsread_stage3_a = num_ds_read_inst_a / KRepeat; + constexpr auto num_dsread_stage3_b = num_ds_read_inst_b / KRepeat; + + constexpr auto num_dsread_stage1_a_mfma = + (num_dsread_stage1_a + ds_read_a_mfma_rate - 1) / ds_read_a_mfma_rate; + constexpr auto num_dsread_stage1_b_mfma = + (num_dsread_stage1_b + ds_read_b_mfma_rate - 1) / ds_read_b_mfma_rate; + constexpr auto num_dsread_stage3_a_mfma = + (num_dsread_stage3_a + ds_read_a_mfma_rate - 1) / ds_read_a_mfma_rate; + constexpr auto num_dsread_stage3_b_mfma = + (num_dsread_stage3_b + ds_read_b_mfma_rate - 1) / ds_read_b_mfma_rate; + + constexpr auto num_mfma_stage2 = C_MFMA_Inst_Num - + num_ds_read_inst_a / ds_read_a_mfma_rate - + num_ds_read_inst_b / ds_read_b_mfma_rate; + constexpr auto num_mfma_per_issue = + num_mfma_stage2 / (A_Buffer_Load_Inst_Num + B_Buffer_Load_Inst_Num); + constexpr auto num_dswrite_per_issue_a = A_LDS_Write_Inst_Num / A_Buffer_Load_Inst_Num; + constexpr auto num_dswrite_per_issue_b = B_LDS_Write_Inst_Num / B_Buffer_Load_Inst_Num; + + // stage 1 + static_for<0, num_dsread_stage1_a_mfma, 1>{}([&](auto i) { + ignore = i; + if constexpr((num_dsread_stage1_a - (i + 1) * ds_read_a_mfma_rate) >= + ds_read_a_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_a_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier( + 0x100, + num_dsread_stage1_a - (num_dsread_stage1_a_mfma - 1) * ds_read_a_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + static_for<0, num_dsread_stage1_b_mfma, 1>{}([&](auto i) { + ignore = i; + if constexpr((num_dsread_stage1_b - (i + 1) * ds_read_b_mfma_rate) >= + ds_read_b_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_b_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier( + 0x100, + num_dsread_stage1_b - (num_dsread_stage1_b_mfma - 1) * ds_read_b_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + + // stage 2 + static_for<0, A_Buffer_Load_Inst_Num, 1>{}([&](auto i) { + ignore = i; + static_for<0, num_dswrite_per_issue_a, 1>{}([&](auto idswrite) { + ignore = idswrite; + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier( + 0x008, num_mfma_per_issue - num_dswrite_per_issue_a, 0); // MFMA + }); + static_for<0, B_Buffer_Load_Inst_Num, 1>{}([&](auto i) { + ignore = i; + static_for<0, num_dswrite_per_issue_b, 1>{}([&](auto idswrite) { + ignore = idswrite; + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier( + 0x008, num_mfma_per_issue - num_dswrite_per_issue_b, 0); // MFMA + }); + + // stage 3 + static_for<0, num_dsread_stage3_a_mfma, 1>{}([&](auto i) { + ignore = i; + if constexpr((num_dsread_stage3_a - (i + 1) * ds_read_a_mfma_rate) >= + ds_read_a_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_a_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier( + 0x100, + num_dsread_stage3_a - (num_dsread_stage3_a_mfma - 1) * ds_read_a_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + static_for<0, num_dsread_stage3_b_mfma, 1>{}([&](auto i) { + ignore = i; + if constexpr((num_dsread_stage3_b - (i + 1) * ds_read_b_mfma_rate) >= + ds_read_b_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_b_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier( + 0x100, + num_dsread_stage3_b - (num_dsread_stage3_b_mfma - 1) * ds_read_b_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + } template CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, - void* __restrict__ p_smem_0) const + void* p_smem) const { static_assert( std::is_same_v> && @@ -127,216 +324,305 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 remove_cvref_t>, "Data Type conflict on A and B matrix input data type."); - static_assert( - KPerBlock % ((NumWarps / 2) * KTileSize) == 0, - "Ping Pong Warps, TileSize and Block Size for K dimensions does not match."); - constexpr bool is_a_col_major = std::is_same_v; constexpr bool is_b_row_major = std::is_same_v; static_assert(is_a_col_major - ? (KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0{}] && - MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1{}]) - : (MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0{}] && - KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1{}]), + ? (KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0] && + MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1]) + : (MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0] && + KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1]), "A block window has incorrect lengths for defined ALayout!"); static_assert(is_b_row_major - ? (KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0{}] && - NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1{}]) - : (NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0{}] && - KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1{}]), + ? (KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0] && + NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1]) + : (NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0] && + KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1]), "B block window has incorrect lengths for defined BLayout!"); - index_t warp_id = get_warp_id(); - index_t operation_id = - __builtin_amdgcn_readfirstlane(get_warp_id()); // 0 - Memory read, 1 - block-gemm - - auto a_offset = (warp_id == 0) ? make_array(0, 0) : make_array(0, KPerBlock); - auto b_offset = (warp_id == 0) ? make_array(0, 0) : make_array(0, KPerBlock); - - auto tensor_views = - Base::GetABLdsTensorViews(static_cast(static_cast(p_smem_0))); - auto& a_lds_block = tensor_views.get(number<0>{}); - auto& b_lds_block = tensor_views.get(number<1>{}); + ////////////// LDS desc, window & register ///////////////// + using ALdsType = + remove_cvref_t; + using BLdsType = + remove_cvref_t; + auto&& ABLdsTensorViews = PipelineImplBase::GetABLdsTensorViews(p_smem); + ALdsType& a_lds_block = ABLdsTensorViews.at(I0); + BLdsType& b_lds_block = ABLdsTensorViews.at(I1); + // 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()); - auto a_windows = Base::GetAWindows( - a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr, a_offset); - auto& a_copy_dram_window = a_windows.get(number<0>{}); - auto& a_copy_lds_window = a_windows.get(number<1>{}); - auto& a_lds_window = a_windows.get(number<2>{}); + using acopy_dram_type = + remove_cvref_t; + using bcopy_dram_type = + remove_cvref_t; - auto b_windows = Base::GetBWindows( - b_dram_block_window_tmp, b_lds_block, b_lds_load_tile_distr, b_offset); - auto& b_copy_dram_window = b_windows.get(number<0>{}); - auto& b_copy_lds_window = b_windows.get(number<1>{}); - auto& b_lds_window = b_windows.get(number<2>{}); + using a_copy_lds_window_type = + remove_cvref_t; + using b_copy_lds_window_type = + remove_cvref_t; - // DRAM window steps. - using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; - using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; - constexpr ADramTileWindowStep a_dram_tile_window_step = - is_a_col_major ? make_array(KPerBlock * NumWarps, 0) - : make_array(0, KPerBlock * NumWarps); - constexpr BDramTileWindowStep b_dram_tile_window_step = - is_b_row_major ? make_array(KPerBlock * NumWarps, 0) - : make_array(0, KPerBlock * NumWarps); + using a_lds_load_tile_distr_type = + remove_cvref_t; + using b_lds_load_tile_distr_type = + remove_cvref_t; - constexpr auto AGemmTileDistr = decltype(make_static_tile_distribution( - BlockGemm::MakeABlockDistributionEncode())){}; - constexpr auto BGemmTileDistr = decltype(make_static_tile_distribution( - BlockGemm::MakeBBlockDistributionEncode())){}; + 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); - using AGemmTile = decltype(make_static_distributed_tensor(AGemmTileDistr)); - using BGemmTile = decltype(make_static_distributed_tensor(BGemmTileDistr)); - AGemmTile a_tile_0, a_tile_1; - BGemmTile b_tile_0, b_tile_1; + // A DRAM tile window for load + // A LDS tile window for store + // A LDS tile for block GEMM + 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); + + // B DRAM tile window for load + // B LDS tile window for store + // B LDS tile for block GEMM + 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); + + // Block GEMM + auto block_gemm = BlockGemm(); + auto c_block_tile = block_gemm.MakeCBlockTile(); - // Register tile for A and B. 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_global_load_tile; - BBlockTile b_global_load_tile; - // Block GEMM - auto block_gemm = BlockGemm(); - auto c_block_tile_0 = block_gemm.MakeCBlockTile(); - auto c_block_tile_1 = block_gemm.MakeCBlockTile(); + ABlockTile a_block_tile[Base::GlobalBufferNum]; + BBlockTile b_block_tile[Base::GlobalBufferNum]; - CDataType* __restrict__ p_c_lds = static_cast(p_smem_0); - auto c_lds_block_0 = - make_naive_tensor_view(p_c_lds, - make_tuple(MPerBlock, NPerBlock), - make_tuple(NPerBlock, 1), - number{}, - number<1>{}); - auto c_window_0 = make_tile_window(c_lds_block_0, - make_tuple(number{}, number{}), - {0, 0}, - c_block_tile_1.get_tile_distribution()); + using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; + using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; + + constexpr ADramTileWindowStep a_dram_tile_window_step = + is_a_col_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock); + constexpr BDramTileWindowStep b_dram_tile_window_step = + is_b_row_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock); + + constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution( + BlockGemm::MakeABlockDistributionEncode())){}; + constexpr auto BLdsTileDistr = decltype(make_static_tile_distribution( + BlockGemm::MakeBBlockDistributionEncode())){}; + + using ALdsTile = decltype(make_static_distributed_tensor(ALdsTileDistr)); + using BLdsTile = decltype(make_static_distributed_tensor(BLdsTileDistr)); + + ALdsTile a_lds_tile; + BLdsTile b_lds_tile; + // ----------------------------------------------------------------------------------------- + // Gemm pipeline start + + // 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); // initialize C - if(warp_id == 0) + tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + + // Local prefill 1 + if constexpr(is_a_col_major) { - tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile_0); + 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 { - tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile_1); + 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); } - // define ping, pong steps here as lambda functions. - auto MemoryOpsStep = [&](auto idx) { - // Memory read half here. - 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 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 + 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 = [&](auto vmem_buf_idx) { + block_sync_lds(); + + // 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_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); + } + + // 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); + + block_sync_lds(); + + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); + + // 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_global_load_tile); - Base::LocalPrefill(a_copy_lds_window, 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_window, 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_window, 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_window, b_global_load_tile, b_element_func); + PipelineImplBase::LocalPrefill( + b_copy_lds_window, b_block_tile[vmem_buf_idx], b_element_func); } - if(idx == 0) - { - Base::LocalPrefetch(a_tile_0, a_lds_window); - Base::LocalPrefetch(b_tile_0, b_lds_window); - } - else - { - Base::LocalPrefetch(a_tile_1, a_lds_window); - Base::LocalPrefetch(b_tile_1, b_lds_window); - } - }; - - auto ComputeStep = [&](auto idx) { - if(idx == 0) - { - block_gemm(c_block_tile_0, a_tile_0, b_tile_0); - } - else - { - block_gemm(c_block_tile_1, a_tile_1, b_tile_1); - } - }; - - if(operation_id == 0) - { - MemoryOpsStep(warp_id); - } - - index_t num_compute_steps = __builtin_amdgcn_readfirstlane(num_loop); - while(num_compute_steps > 1) - { block_sync_lds(); - operation_id = (operation_id + 1) % NumWaveGroups; + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); - if(operation_id == 0) - { - MemoryOpsStep(warp_id); - } - else - { - ComputeStep(warp_id); - } - num_compute_steps -= 1; - } - block_sync_lds(); + PipelineImplBase::LocalPrefetch(a_lds_tile, a_lds_gemm_window); + PipelineImplBase::LocalPrefetch(b_lds_tile, b_lds_gemm_window); - if(operation_id == 0) + HotLoopScheduler(); + }; + + auto ReadCompFunc = [&]() { + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); + + // Local prefetch 4 + PipelineImplBase::LocalPrefetch(a_lds_tile, a_lds_gemm_window); + PipelineImplBase::LocalPrefetch(b_lds_tile, b_lds_gemm_window); + + block_gemm(c_block_tile, a_lds_tile, b_lds_tile); + + HotLoopScheduler(); + }; + + if constexpr(TailNum == TailNumber::Odd) { - ComputeStep(warp_id); + ReadWriteCompFunc(I0); + ReadWriteCompFunc(I1); + ReadCompFunc(); } - block_sync_lds(); - - if(warp_id == 1) + else if constexpr(TailNum == TailNumber::Even) { - store_tile(c_window_0, c_block_tile_1); + ReadWriteCompFunc(I0); + ReadCompFunc(); } - block_sync_lds(); - if(warp_id == 0) - { - load_tile(c_block_tile_1, c_window_0); - - constexpr auto s_spans = decltype(c_block_tile_0)::get_distributed_spans(); - sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) { - sweep_tile_span(s_spans[number<1>{}], [&](auto idx1) { - auto idx2 = make_tuple(idx0, idx1); - c_block_tile_0(idx2) += c_block_tile_1(idx2); - }); - }); - } - return c_block_tile_0; + return c_block_tile; } }; @@ -349,7 +635,7 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 const BDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, - void* p_smem_0) const + void* p_smem) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, @@ -357,7 +643,7 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 b_dram_block_window_tmp, b_element_func, num_loop, - p_smem_0); + p_smem); } public: @@ -365,7 +651,7 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const BDramBlockWindowTmp& b_dram_block_window_tmp, const index_t num_loop, - void* __restrict__ p_smem_0) const + void* p_smem) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, @@ -373,8 +659,7 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 b_dram_block_window_tmp, [](const BDataType& b) { return b; }, num_loop, - p_smem_0); + p_smem); } }; - } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp index c03db08c3f..74aeccdbe5 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp @@ -38,26 +38,5 @@ struct GemmPipelineAgBgCrCompV5DefaultPolicy return BlockGemmARegBRegCRegV1{}; } - - template - CK_TILE_DEVICE static constexpr index_t GetSmemSizeC() - { - constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - - return integer_least_multiple(sizeof(typename Problem::CDataType) * MPerBlock * NPerBlock, - 16); - } - - template - CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() - { - constexpr index_t smem_size_a = GetSmemSizeA(); - constexpr index_t smem_size_b = GetSmemSizeB(); - constexpr index_t smem_size_c = GetSmemSizeC(); - - return smem_size_a + smem_size_b >= smem_size_c ? (smem_size_a + smem_size_b) - : (smem_size_c); - } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6.hpp index bcb3f2da98..271aa4d7ce 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6.hpp @@ -1,94 +1,60 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -#pragma once #include "ck_tile/core.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp" -#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6_default_policy.hpp" +#include "ck_tile/host/concat.hpp" namespace ck_tile { +// A Tile Window: global memory +// B Tile Window: global memory +// C Distributed Tensor: register -// A Tile Window: global memory -// B Tile Window: global memory -// C Distributed tensor: register template struct BaseGemmPipelineAgBgCrCompV6 { - static constexpr index_t PrefetchStages = 3; + static constexpr index_t PrefetchStages = 1; static constexpr index_t PrefillStages = 1; - static constexpr index_t GlobalBufferNum = 2; - static constexpr index_t HotloopUnroll = 2; + static constexpr index_t GlobalBufferNum = 1; 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; - } + CK_TILE_HOST_DEVICE static constexpr bool BlockHasHotloop(index_t) { return true; } - CK_TILE_HOST static constexpr TailNumber GetBlockLoopTailNum(index_t num_loop) + CK_TILE_HOST_DEVICE static constexpr TailNumber GetBlockLoopTailNum(index_t) { - if(num_loop % HotloopUnroll == 1) - { - return TailNumber::Odd; - } - else - { - return TailNumber::Even; - } + return TailNumber::Empty; } template - CK_TILE_HOST_DEVICE static auto - TailHandler(const RunFunction& run_func, bool has_hot_loop, TailNumber tail_num) + CK_TILE_HOST_DEVICE static auto TailHandler(const RunFunction& run_func, bool, TailNumber) { - if(tail_num == ck_tile::TailNumber::Three) - { - return run_func( - has_hot_loop, - ck_tile::integral_constant{}); - } - else - { - return run_func( - has_hot_loop, - ck_tile::integral_constant{}); - } + return run_func(bool_constant{}, integral_constant{}); } }; -/** - * @brief Compute optimized pipeline version 5 TODO - * - * - * @note TODO - */ template struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 { using Base = BaseGemmPipelineAgBgCrCompV6; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; - - static_assert(!std::is_same_v, "Not implemented"); - - static constexpr index_t APackedSize = - ck_tile::numeric_traits>::PackedSize; - static constexpr index_t BPackedSize = - ck_tile::numeric_traits>::PackedSize; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using ComputeDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; using CLayout = remove_cvref_t; - using BlockGemm = remove_cvref_t())>; - static constexpr auto I0 = number<0>{}; - static constexpr auto I1 = number<1>{}; - static constexpr auto I2 = number<2>{}; + static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + + using BlockGemm = remove_cvref_t())>; + using I0 = number<0>; + using I1 = number<1>; + using I2 = number<2>; static constexpr index_t BlockSize = Problem::kBlockSize; @@ -100,12 +66,6 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 static constexpr index_t GetVectorSizeB() { return Policy::template GetVectorSizeB(); } static constexpr index_t GetVectorSizeC() { return Policy::template GetVectorSizeC(); } - static constexpr index_t GetSmemPackA() { return Policy::template GetSmemPackA(); } - static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB(); } - - // TODO check KRepeat - static constexpr index_t KRepeat = KPerBlock / GetSmemPackA(); - static constexpr bool kPadM = Problem::kPadM; static constexpr bool kPadN = Problem::kPadN; static constexpr bool kPadK = Problem::kPadK; @@ -116,6 +76,18 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 static constexpr auto TailNum = Problem::TailNum; static constexpr auto Scheduler = Problem::Scheduler; + static constexpr index_t NumWarps = BlockGemmShape::NumWarps; + static constexpr index_t KTileSize = BlockGemmShape::WarpTile::at(I2{}); + + [[nodiscard]] CK_TILE_HOST static const std::string GetName() + { + // clang-format off + return concat('_', "pipeline_AgBgCrCompV6", BlockSize, + concat('x', GetVectorSizeA(), GetVectorSizeB(), GetVectorSizeC()), + concat('x', kPadM, kPadN, kPadK)); + // clang-format on + } + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return Policy::template GetSmemSize(); @@ -134,180 +106,20 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 template <> struct PipelineImpl : public PipelineImplBase { - CK_TILE_DEVICE static constexpr auto HotLoopScheduler() - { - constexpr index_t MPerXDL = BlockGemmShape::WarpTile::at(I0); - constexpr index_t NPerXDL = BlockGemmShape::WarpTile::at(I1); - constexpr index_t KPerXDL = BlockGemmShape::WarpTile::at(I2); - - constexpr index_t WaveSize = 64; - constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(I0); - constexpr index_t WaveNumN = BlockGemmShape::BlockWarps::at(I1); - - constexpr index_t A_LDS_Read_Width = KPerXDL; - constexpr index_t B_LDS_Read_Width = KPerXDL; - - constexpr index_t A_Buffer_Load_Inst_Num = - MPerBlock * KPerBlock / (BlockSize * GetVectorSizeA()); - constexpr index_t B_Buffer_Load_Inst_Num = - NPerBlock * KPerBlock / (BlockSize * GetVectorSizeB()); - - constexpr index_t A_LDS_Write_Inst_Num = MPerBlock * KPerBlock / (BlockSize * KPerXDL); - constexpr index_t B_LDS_Write_Inst_Num = NPerBlock * KPerBlock / (BlockSize * KPerXDL); - - constexpr index_t A_LDS_Read_Inst_Num = - WaveNumN * MPerBlock * KPerBlock / (BlockSize * KPerXDL); - constexpr index_t B_LDS_Read_Inst_Num = - WaveNumM * NPerBlock * KPerBlock / (BlockSize * KPerXDL); - - constexpr index_t C_MFMA_Inst_Num = MPerBlock * NPerBlock * KPerBlock / - (BlockSize / WaveSize) / - (MPerXDL * NPerXDL * KPerXDL); - - constexpr auto num_ds_read_inst_a = - A_LDS_Read_Width * sizeof(ADataType) / APackedSize == 16 ? A_LDS_Read_Inst_Num - : A_LDS_Read_Inst_Num / 2; - constexpr auto num_ds_read_inst_b = - B_LDS_Read_Width * sizeof(BDataType) / BPackedSize == 16 ? B_LDS_Read_Inst_Num - : B_LDS_Read_Inst_Num / 2; - - constexpr auto mfma_cycle = NPerXDL == 16 ? 16 : 32; - constexpr auto ds_read_a_issue_cycle = - A_LDS_Read_Width * sizeof(ADataType) == 16 ? 8 : 4; - constexpr auto ds_read_b_issue_cycle = - B_LDS_Read_Width * sizeof(BDataType) == 16 ? 8 : 4; - constexpr auto ds_read_a_mfma_rate = - (mfma_cycle - 4 + 2 * ds_read_a_issue_cycle - 1) / (2 * ds_read_a_issue_cycle); - constexpr auto ds_read_b_mfma_rate = - (mfma_cycle - 4 + 2 * ds_read_b_issue_cycle - 1) / (2 * ds_read_b_issue_cycle); - - constexpr auto num_dsread_stage1_a = num_ds_read_inst_a / KRepeat * (KRepeat - 1); - constexpr auto num_dsread_stage1_b = num_ds_read_inst_b / KRepeat * (KRepeat - 1); - constexpr auto num_dsread_stage3_a = num_ds_read_inst_a / KRepeat; - constexpr auto num_dsread_stage3_b = num_ds_read_inst_b / KRepeat; - - constexpr auto num_dsread_stage1_a_mfma = - (num_dsread_stage1_a + ds_read_a_mfma_rate - 1) / ds_read_a_mfma_rate; - constexpr auto num_dsread_stage1_b_mfma = - (num_dsread_stage1_b + ds_read_b_mfma_rate - 1) / ds_read_b_mfma_rate; - constexpr auto num_dsread_stage3_a_mfma = - (num_dsread_stage3_a + ds_read_a_mfma_rate - 1) / ds_read_a_mfma_rate; - constexpr auto num_dsread_stage3_b_mfma = - (num_dsread_stage3_b + ds_read_b_mfma_rate - 1) / ds_read_b_mfma_rate; - - constexpr auto num_mfma_stage2 = C_MFMA_Inst_Num - - num_ds_read_inst_a / ds_read_a_mfma_rate - - num_ds_read_inst_b / ds_read_b_mfma_rate; - constexpr auto num_mfma_per_issue = - num_mfma_stage2 / (A_Buffer_Load_Inst_Num + B_Buffer_Load_Inst_Num); - constexpr auto num_dswrite_per_issue_a = A_LDS_Write_Inst_Num / A_Buffer_Load_Inst_Num; - constexpr auto num_dswrite_per_issue_b = B_LDS_Write_Inst_Num / B_Buffer_Load_Inst_Num; - - // stage 1 - static_for<0, num_dsread_stage1_a_mfma, 1>{}([&](auto i) { - ignore = i; - if constexpr((num_dsread_stage1_a - (i + 1) * ds_read_a_mfma_rate) >= - ds_read_a_mfma_rate) - { - __builtin_amdgcn_sched_group_barrier(0x100, ds_read_a_mfma_rate, 0); // DS read - } - else - { - __builtin_amdgcn_sched_group_barrier( - 0x100, - num_dsread_stage1_a - (num_dsread_stage1_a_mfma - 1) * ds_read_a_mfma_rate, - 0); // DS read - } - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - }); - static_for<0, num_dsread_stage1_b_mfma, 1>{}([&](auto i) { - ignore = i; - if constexpr((num_dsread_stage1_b - (i + 1) * ds_read_b_mfma_rate) >= - ds_read_b_mfma_rate) - { - __builtin_amdgcn_sched_group_barrier(0x100, ds_read_b_mfma_rate, 0); // DS read - } - else - { - __builtin_amdgcn_sched_group_barrier( - 0x100, - num_dsread_stage1_b - (num_dsread_stage1_b_mfma - 1) * ds_read_b_mfma_rate, - 0); // DS read - } - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - }); - - // stage 2 - static_for<0, A_Buffer_Load_Inst_Num, 1>{}([&](auto i) { - ignore = i; - static_for<0, num_dswrite_per_issue_a, 1>{}([&](auto idswrite) { - ignore = idswrite; - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - }); - __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read - __builtin_amdgcn_sched_group_barrier( - 0x008, num_mfma_per_issue - num_dswrite_per_issue_a, 0); // MFMA - }); - static_for<0, B_Buffer_Load_Inst_Num, 1>{}([&](auto i) { - ignore = i; - static_for<0, num_dswrite_per_issue_b, 1>{}([&](auto idswrite) { - ignore = idswrite; - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - }); - __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read - __builtin_amdgcn_sched_group_barrier( - 0x008, num_mfma_per_issue - num_dswrite_per_issue_b, 0); // MFMA - }); - - // stage 3 - static_for<0, num_dsread_stage3_a_mfma, 1>{}([&](auto i) { - ignore = i; - if constexpr((num_dsread_stage3_a - (i + 1) * ds_read_a_mfma_rate) >= - ds_read_a_mfma_rate) - { - __builtin_amdgcn_sched_group_barrier(0x100, ds_read_a_mfma_rate, 0); // DS read - } - else - { - __builtin_amdgcn_sched_group_barrier( - 0x100, - num_dsread_stage3_a - (num_dsread_stage3_a_mfma - 1) * ds_read_a_mfma_rate, - 0); // DS read - } - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - }); - static_for<0, num_dsread_stage3_b_mfma, 1>{}([&](auto i) { - ignore = i; - if constexpr((num_dsread_stage3_b - (i + 1) * ds_read_b_mfma_rate) >= - ds_read_b_mfma_rate) - { - __builtin_amdgcn_sched_group_barrier(0x100, ds_read_b_mfma_rate, 0); // DS read - } - else - { - __builtin_amdgcn_sched_group_barrier( - 0x100, - num_dsread_stage3_b - (num_dsread_stage3_b_mfma - 1) * ds_read_b_mfma_rate, - 0); // DS read - } - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - }); - } + using Base = PipelineImplBase; template CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, - void* p_smem) const + void* __restrict__ p_smem_0) const { static_assert( std::is_same_v> && @@ -315,305 +127,216 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 remove_cvref_t>, "Data Type conflict on A and B matrix input data type."); + static_assert( + KPerBlock % ((NumWarps / 2) * KTileSize) == 0, + "Ping Pong Warps, TileSize and Block Size for K dimensions does not match."); + constexpr bool is_a_col_major = std::is_same_v; constexpr bool is_b_row_major = std::is_same_v; static_assert(is_a_col_major - ? (KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0] && - MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1]) - : (MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0] && - KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1]), + ? (KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0{}] && + MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1{}]) + : (MPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0{}] && + KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I1{}]), "A block window has incorrect lengths for defined ALayout!"); static_assert(is_b_row_major - ? (KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0] && - NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1]) - : (NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0] && - KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1]), + ? (KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0{}] && + NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1{}]) + : (NPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I0{}] && + KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1{}]), "B block window has incorrect lengths for defined BLayout!"); - ////////////// LDS desc, window & register ///////////////// - using ALdsType = - remove_cvref_t; - using BLdsType = - remove_cvref_t; - auto&& ABLdsTensorViews = PipelineImplBase::GetABLdsTensorViews(p_smem); - ALdsType& a_lds_block = ABLdsTensorViews.at(I0); - BLdsType& b_lds_block = ABLdsTensorViews.at(I1); + index_t warp_id = get_warp_id(); + index_t operation_id = + __builtin_amdgcn_readfirstlane(get_warp_id()); // 0 - Memory read, 1 - block-gemm + + auto a_offset = (warp_id == 0) ? make_array(0, 0) : make_array(0, KPerBlock); + auto b_offset = (warp_id == 0) ? make_array(0, 0) : make_array(0, KPerBlock); + + auto tensor_views = + Base::GetABLdsTensorViews(static_cast(static_cast(p_smem_0))); + auto& a_lds_block = tensor_views.get(number<0>{}); + auto& b_lds_block = tensor_views.get(number<1>{}); - // 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()); - using acopy_dram_type = - remove_cvref_t; - using bcopy_dram_type = - remove_cvref_t; + auto a_windows = Base::GetAWindows( + a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr, a_offset); + auto& a_copy_dram_window = a_windows.get(number<0>{}); + auto& a_copy_lds_window = a_windows.get(number<1>{}); + auto& a_lds_window = a_windows.get(number<2>{}); - using a_copy_lds_window_type = - remove_cvref_t; - using b_copy_lds_window_type = - remove_cvref_t; + auto b_windows = Base::GetBWindows( + b_dram_block_window_tmp, b_lds_block, b_lds_load_tile_distr, b_offset); + auto& b_copy_dram_window = b_windows.get(number<0>{}); + auto& b_copy_lds_window = b_windows.get(number<1>{}); + auto& b_lds_window = b_windows.get(number<2>{}); - using a_lds_load_tile_distr_type = - remove_cvref_t; - using b_lds_load_tile_distr_type = - remove_cvref_t; + // DRAM window steps. + using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; + using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; + constexpr ADramTileWindowStep a_dram_tile_window_step = + is_a_col_major ? make_array(KPerBlock * NumWarps, 0) + : make_array(0, KPerBlock * NumWarps); + constexpr BDramTileWindowStep b_dram_tile_window_step = + is_b_row_major ? make_array(KPerBlock * NumWarps, 0) + : make_array(0, KPerBlock * NumWarps); - 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); + constexpr auto AGemmTileDistr = decltype(make_static_tile_distribution( + BlockGemm::MakeABlockDistributionEncode())){}; + constexpr auto BGemmTileDistr = decltype(make_static_tile_distribution( + BlockGemm::MakeBBlockDistributionEncode())){}; - // A DRAM tile window for load - // A LDS tile window for store - // A LDS tile for block GEMM - 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); - - // B DRAM tile window for load - // B LDS tile window for store - // B LDS tile for block GEMM - 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); - - // Block GEMM - auto block_gemm = BlockGemm(); - auto c_block_tile = block_gemm.MakeCBlockTile(); + using AGemmTile = decltype(make_static_distributed_tensor(AGemmTileDistr)); + using BGemmTile = decltype(make_static_distributed_tensor(BGemmTileDistr)); + AGemmTile a_tile_0, a_tile_1; + BGemmTile b_tile_0, b_tile_1; + // Register tile for A and B. 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_global_load_tile; + BBlockTile b_global_load_tile; - ABlockTile a_block_tile[Base::GlobalBufferNum]; - BBlockTile b_block_tile[Base::GlobalBufferNum]; + // Block GEMM + auto block_gemm = BlockGemm(); + auto c_block_tile_0 = block_gemm.MakeCBlockTile(); + auto c_block_tile_1 = block_gemm.MakeCBlockTile(); - using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; - using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; - - constexpr ADramTileWindowStep a_dram_tile_window_step = - is_a_col_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock); - constexpr BDramTileWindowStep b_dram_tile_window_step = - is_b_row_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock); - - constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution( - BlockGemm::MakeABlockDistributionEncode())){}; - constexpr auto BLdsTileDistr = decltype(make_static_tile_distribution( - BlockGemm::MakeBBlockDistributionEncode())){}; - - using ALdsTile = decltype(make_static_distributed_tensor(ALdsTileDistr)); - using BLdsTile = decltype(make_static_distributed_tensor(BLdsTileDistr)); - - ALdsTile a_lds_tile; - BLdsTile b_lds_tile; - // ----------------------------------------------------------------------------------------- - // Gemm pipeline start - - // 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); + CDataType* __restrict__ p_c_lds = static_cast(p_smem_0); + auto c_lds_block_0 = + make_naive_tensor_view(p_c_lds, + make_tuple(MPerBlock, NPerBlock), + make_tuple(NPerBlock, 1), + number{}, + number<1>{}); + auto c_window_0 = make_tile_window(c_lds_block_0, + make_tuple(number{}, number{}), + {0, 0}, + c_block_tile_1.get_tile_distribution()); // initialize C - tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); - - // Local prefill 1 - if constexpr(is_a_col_major) + if(warp_id == 0) { - 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); + tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile_0); } 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); + tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile_1); } - // 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); + // define ping, pong steps here as lambda functions. + auto MemoryOpsStep = [&](auto idx) { + // Memory read half here. + 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 - 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 - 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 = [&](auto vmem_buf_idx) { - block_sync_lds(); - - // 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_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); - } - - // 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); - - block_sync_lds(); - - block_gemm(c_block_tile, a_lds_tile, b_lds_tile); - - // 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); + transpose_tile2d(a_shuffle_tmp, a_global_load_tile); + Base::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); + Base::LocalPrefill(a_copy_lds_window, 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_block_tile[vmem_buf_idx]); - PipelineImplBase::LocalPrefill( - b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, b_global_load_tile); + Base::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); + Base::LocalPrefill(b_copy_lds_window, b_global_load_tile, b_element_func); } + if(idx == 0) + { + Base::LocalPrefetch(a_tile_0, a_lds_window); + Base::LocalPrefetch(b_tile_0, b_lds_window); + } + else + { + Base::LocalPrefetch(a_tile_1, a_lds_window); + Base::LocalPrefetch(b_tile_1, b_lds_window); + } + }; + + auto ComputeStep = [&](auto idx) { + if(idx == 0) + { + block_gemm(c_block_tile_0, a_tile_0, b_tile_0); + } + else + { + block_gemm(c_block_tile_1, a_tile_1, b_tile_1); + } + }; + + if(operation_id == 0) + { + MemoryOpsStep(warp_id); + } + + index_t num_compute_steps = __builtin_amdgcn_readfirstlane(num_loop); + while(num_compute_steps > 1) + { block_sync_lds(); - block_gemm(c_block_tile, a_lds_tile, b_lds_tile); + operation_id = (operation_id + 1) % NumWaveGroups; - 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_lds_tile, b_lds_tile); - - // Local prefetch 4 - PipelineImplBase::LocalPrefetch(a_lds_tile, a_lds_gemm_window); - PipelineImplBase::LocalPrefetch(b_lds_tile, b_lds_gemm_window); - - block_gemm(c_block_tile, a_lds_tile, b_lds_tile); - - HotLoopScheduler(); - }; - - if constexpr(TailNum == TailNumber::Odd) - { - ReadWriteCompFunc(I0); - ReadWriteCompFunc(I1); - ReadCompFunc(); - } - else if constexpr(TailNum == TailNumber::Even) - { - ReadWriteCompFunc(I0); - ReadCompFunc(); + if(operation_id == 0) + { + MemoryOpsStep(warp_id); + } + else + { + ComputeStep(warp_id); + } + num_compute_steps -= 1; } + block_sync_lds(); - return c_block_tile; + if(operation_id == 0) + { + ComputeStep(warp_id); + } + block_sync_lds(); + + if(warp_id == 1) + { + store_tile(c_window_0, c_block_tile_1); + } + block_sync_lds(); + + if(warp_id == 0) + { + load_tile(c_block_tile_1, c_window_0); + + constexpr auto s_spans = decltype(c_block_tile_0)::get_distributed_spans(); + sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(s_spans[number<1>{}], [&](auto idx1) { + auto idx2 = make_tuple(idx0, idx1); + c_block_tile_0(idx2) += c_block_tile_1(idx2); + }); + }); + } + return c_block_tile_0; } }; @@ -626,7 +349,7 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 const BDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, - void* p_smem) const + void* p_smem_0) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, @@ -634,7 +357,7 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 b_dram_block_window_tmp, b_element_func, num_loop, - p_smem); + p_smem_0); } public: @@ -642,7 +365,7 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const BDramBlockWindowTmp& b_dram_block_window_tmp, const index_t num_loop, - void* p_smem) const + void* __restrict__ p_smem_0) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, @@ -650,7 +373,8 @@ struct GemmPipelineAgBgCrCompV6 : public BaseGemmPipelineAgBgCrCompV6 b_dram_block_window_tmp, [](const BDataType& b) { return b; }, num_loop, - p_smem); + p_smem_0); } }; + } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6_default_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6_default_policy.hpp index 7307981f5f..ccdf7c07ab 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6_default_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6_default_policy.hpp @@ -38,5 +38,26 @@ struct GemmPipelineAgBgCrCompV6DefaultPolicy return BlockGemmARegBRegCRegV1{}; } + + template + CK_TILE_DEVICE static constexpr index_t GetSmemSizeC() + { + constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + + return integer_least_multiple(sizeof(typename Problem::CDataType) * MPerBlock * NPerBlock, + 16); + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + constexpr index_t smem_size_a = GetSmemSizeA(); + constexpr index_t smem_size_b = GetSmemSizeB(); + constexpr index_t smem_size_c = GetSmemSizeC(); + + return smem_size_a + smem_size_b >= smem_size_c ? (smem_size_a + smem_size_b) + : (smem_size_c); + } }; } // namespace ck_tile