diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp new file mode 100644 index 0000000000..e445715197 --- /dev/null +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp @@ -0,0 +1,475 @@ +// 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_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_v4_default_policy.hpp" + +namespace ck_tile { + +// A Tile Window: global memory +// B Tile Window: global memory +// C Distributed tensor: register +template +struct BaseGemmPipelineAgBgCrCompV4 +{ + static constexpr index_t PrefetchStages = 2; + static constexpr index_t PrefillStages = 1; + static constexpr index_t GlobalBufferNum = 1; + + CK_TILE_HOST static constexpr bool BlockHasHotloop(index_t num_loop) + { + return num_loop > PrefetchStages; + } + + CK_TILE_HOST static constexpr TailNumber GetBlockLoopTailNum(index_t num_loop) + { + if(num_loop % PrefetchStages == 1) + { + return TailNumber::Three; + } + else + { + return TailNumber::Two; + } + } +}; + +/** + * @brief Compute optimized pipeline version async; which is changed from V4. + * + * This version introduces a dual LDS window mechanism using a ping-pong buffer approach + * for more efficient data handling from global memory. Unlike compute version 3, this method + * allows one LDS to fetch data from global memory while the other LDS executes warps for MFMA + * matrix multiplication. This dual operation helps in keeping the Warp unit continuously busy, + * thereby significantly reducing memory load times and enhancing overall performance. + * + * @note This version shows improved performance over Compute Version 3 with the same block tile. + * It is particularly more efficient for large matrices where M, N, and K are greater than 8K, + * even when Compute Version 3's block size is twice that of Compute Version 4. + */ +template +struct GemmPipelineAgBgCrCompAsync : public BaseGemmPipelineAgBgCrCompV4 +{ + using Base = BaseGemmPipelineAgBgCrCompV4; + 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 ALayout = remove_cvref_t; + using BLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using BlockGemm = remove_cvref_t())>; + using I0 = number<0>; + using I1 = number<1>; + using I2 = number<2>; + + static constexpr index_t BlockSize = Problem::kBlockSize; + + static constexpr index_t MPerBlock = BlockGemmShape::kM; + static constexpr index_t NPerBlock = BlockGemmShape::kN; + static constexpr index_t KPerBlock = BlockGemmShape::kK; + + static constexpr index_t GetVectorSizeA() { return Policy::template GetVectorSizeA(); } + 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(); } + + static constexpr bool kPadM = Problem::kPadM; + static constexpr bool kPadN = Problem::kPadN; + static constexpr bool kPadK = Problem::kPadK; + + static constexpr bool DoubleSmemBuffer = Problem::DoubleSmemBuffer; + + static constexpr bool HasHotLoop = Problem::HasHotLoop; + static constexpr auto TailNum = Problem::TailNum; + static constexpr auto Scheduler = Problem::Scheduler; + + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + return Policy::template GetSmemSize(); + } + + CK_TILE_HOST_DEVICE static constexpr auto IsTransposeC() + { + return Policy::template IsTransposeC(); + } + + template + struct PipelineImpl : public PipelineImplBase + { + }; + + 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 num_ds_read_inst = num_ds_read_inst_a + num_ds_read_inst_b; + constexpr auto num_ds_write_inst = A_LDS_Write_Inst_Num + B_LDS_Write_Inst_Num; + constexpr auto num_buffer_load_inst = A_Buffer_Load_Inst_Num + B_Buffer_Load_Inst_Num; + constexpr auto num_issue = num_buffer_load_inst; + + static_for<0, num_buffer_load_inst, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA : 1 + __builtin_amdgcn_sched_group_barrier( + 0x100, num_ds_read_inst / num_issue, 0); // DS read : 2 + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA: 1 + __builtin_amdgcn_sched_group_barrier( + 0x200, num_ds_write_inst / num_issue, 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, C_MFMA_Inst_Num / num_issue - 3, 0); // MFMA : 5 + }); + __builtin_amdgcn_sched_barrier(0); + } + + 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, + void* __restrict__ p_smem_1) const + { + static_assert( + std::is_same_v> && + std::is_same_v>, + "Data Type conflict on A and B matrix input data type."); + + constexpr bool is_a_col_major = + std::is_same_v; + constexpr bool is_b_row_major = std::is_same_v; + + // TODO currently only support A matrix row major, B matrix col major; if A matrix is + // col major or B is row major, need to combine with transpose load api + static_assert(!(is_a_col_major || is_b_row_major), + "only support A matrix is row major, B matrix is col major!"); + + 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{}]), + "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{}]), + "B block window has incorrect lengths for defined BLayout!"); + + ////////////// global window & register ///////////////// + // 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()); + + // 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()); + + // 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; + auto&& [a_lds_block0, b_lds_block0] = Base::GetABLdsTensorViews(p_smem_0); + auto&& [a_lds_block1, b_lds_block1] = Base::GetABLdsTensorViews(p_smem_1); + + auto a_copy_lds_window0 = make_tile_window( + a_lds_block0, make_tuple(number{}, number{}), {0, 0}); + + auto a_copy_lds_window1 = make_tile_window( + a_lds_block1, make_tuple(number{}, number{}), {0, 0}); + + auto b_copy_lds_window0 = make_tile_window( + b_lds_block0, make_tuple(number{}, number{}), {0, 0}); + + auto b_copy_lds_window1 = make_tile_window( + b_lds_block1, make_tuple(number{}, number{}), {0, 0}); + + 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); + + // the below is used for async cnt calculation + using ACopyDramWindow = remove_cvref_t; + using BCopyDramWindow = remove_cvref_t; + constexpr auto a_number_of_access = ACopyDramWindow{}.get_num_of_access(); + constexpr auto b_number_of_access = BCopyDramWindow{}.get_num_of_access(); + // global prefetch 0 + // global read 0 + Base::GlobalPrefetchAsync( + a_copy_lds_window0, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetchAsync( + b_copy_lds_window0, b_copy_dram_window, b_dram_tile_window_step); + ////////////// LDS desc, window & register ///////////////// + + // 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); + + // global read 1 + Base::GlobalPrefetchAsync( + a_copy_lds_window1, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetchAsync( + b_copy_lds_window1, b_copy_dram_window, b_dram_tile_window_step); + + buffer_load_fence(a_number_of_access + b_number_of_access); + + 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_block_tile0; + ALdsTile a_block_tile1; + + BLdsTile b_block_tile0; + BLdsTile b_block_tile1; + + auto a_lds_ld_window0 = + make_tile_window(a_lds_block0, + make_tuple(number{}, number{}), + {0, 0}, + ALdsTileDistr); + auto a_lds_ld_window1 = + make_tile_window(a_lds_block1, + make_tuple(number{}, number{}), + {0, 0}, + ALdsTileDistr); + auto b_lds_ld_window0 = + make_tile_window(b_lds_block0, + make_tuple(number{}, number{}), + {0, 0}, + BLdsTileDistr); + auto b_lds_ld_window1 = + make_tile_window(b_lds_block1, + make_tuple(number{}, number{}), + {0, 0}, + BLdsTileDistr); + + static_assert( + !(is_tile_window_linear_v)&&!(is_tile_window_linear_v)&&!( + is_tile_window_linear_v< + decltype(b_lds_ld_window0)>)&&!(is_tile_window_linear_v), + "LDS windows must not be linear"); + + Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); + Base::LocalPrefetch(b_block_tile0, b_lds_ld_window0); + + Base::GlobalPrefetchAsync( + a_copy_lds_window0, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetchAsync( + b_copy_lds_window0, b_copy_dram_window, b_dram_tile_window_step); + + if(HasHotLoop) + { + // minus 2 because we have ping-pong double buffer. + index_t iCounter = __builtin_amdgcn_readfirstlane(num_loop - 2); + do + { + // ping + { + buffer_load_fence(a_number_of_access + b_number_of_access); + Base::LocalPrefetch(a_block_tile1, a_lds_ld_window1); + Base::LocalPrefetch(b_block_tile1, b_lds_ld_window1); + + Base::GlobalPrefetchAsync( + a_copy_lds_window1, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetchAsync( + b_copy_lds_window1, b_copy_dram_window, b_dram_tile_window_step); + // gemm + block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + HotLoopScheduler(); + __builtin_amdgcn_sched_barrier(0); + } + // pong + { + buffer_load_fence(a_number_of_access + b_number_of_access); + Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); + Base::LocalPrefetch(b_block_tile0, b_lds_ld_window0); + + block_sync_lds(); + + Base::GlobalPrefetchAsync( + a_copy_lds_window0, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetchAsync( + b_copy_lds_window0, b_copy_dram_window, b_dram_tile_window_step); + // gemm + block_gemm(c_block_tile, a_block_tile1, b_block_tile1); + HotLoopScheduler(); + __builtin_amdgcn_sched_barrier(0); + } + iCounter -= 2; + } while(iCounter > 1); + } + + // tail 3 + if(TailNum == TailNumber::Three) + { + // 3 + { + block_sync_lds(); + Base::LocalPrefetch(a_block_tile1, a_lds_ld_window1); + Base::LocalPrefetch(b_block_tile1, b_lds_ld_window1); + block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + } + // 2 + { + block_sync_lds(); + Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); + Base::LocalPrefetch(a_block_tile0, a_lds_ld_window0); + block_gemm(c_block_tile, a_block_tile1, b_block_tile1); + } + // 1 + { + block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + __builtin_amdgcn_sched_barrier(0); + } + } + else + { + // 2 + { + block_sync_lds(); + Base::LocalPrefetch(a_block_tile1, a_lds_ld_window1); + Base::LocalPrefetch(b_block_tile1, b_lds_ld_window1); + block_gemm(c_block_tile, a_block_tile0, b_block_tile0); + static_for<0, 8, 1>{}([&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x100, 1, 0); // DS read + __builtin_amdgcn_sched_group_barrier(0x008, 8, 0); // MFMA + }); + __builtin_amdgcn_sched_barrier(0); + } + // 1 + { + block_gemm(c_block_tile, a_block_tile1, b_block_tile1); + __builtin_amdgcn_sched_barrier(0); + } + } + return c_block_tile; + } + }; + + 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_0, + void* p_smem_1) const + { + return PipelineImpl{}.template operator()( + a_dram_block_window_tmp, + a_element_func, + b_dram_block_window_tmp, + b_element_func, + num_loop, + p_smem_0, + p_smem_1); + } + + public: + template + 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, + void* __restrict__ p_smem_1) const + { + return PipelineImpl{}.template operator()( + a_dram_block_window_tmp, + [](const ADataType& a) { return a; }, + b_dram_block_window_tmp, + [](const BDataType& b) { return b; }, + num_loop, + p_smem_0, + p_smem_1); + } +}; +} // namespace ck_tile