From 85f84a34abcac314d0a9c64b72b49fdb40f81243 Mon Sep 17 00:00:00 2001 From: ThomasNing Date: Mon, 30 Jun 2025 11:27:28 -0500 Subject: [PATCH] finish the feature --- example/ck_tile/03_gemm/universal_gemm.cpp | 6 +- example/ck_tile/36_copy/test_copy.cpp | 8 +- example/ck_tile/36_copy/test_copy.hpp | 13 +- include/ck_tile/core/tensor/buffer_view.hpp | 1 + include/ck_tile/core/tensor/tile_window.hpp | 66 +-- .../core/tensor/tile_window_linear.hpp | 60 +-- include/ck_tile/ops/gemm.hpp | 1 - .../gemm_pipeline_ag_bg_cr_comp_async.hpp | 467 ------------------ ...ine_ag_bg_cr_comp_async_default_policy.hpp | 89 ---- .../gemm_pipeline_ag_bg_cr_comp_v3.hpp | 14 +- ...emm_universal_pipeline_ag_bg_cr_policy.hpp | 219 ++++---- 11 files changed, 137 insertions(+), 807 deletions(-) delete mode 100644 include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp delete mode 100644 include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_default_policy.hpp diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index e3a76450b8..cdf9f5338d 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -249,12 +249,12 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a } else if(a_layout == "C" && b_layout == "R") { - return run_gemm_example_with_layouts( - argc, argv, Col{}, Row{}, Row{}); + return run_gemm_example_with_layouts( + argc, argv, Col{}, Row{}, Row{}); } else if(a_layout == "C" && b_layout == "C") { - return run_gemm_example_with_layouts( + return run_gemm_example_with_layouts( argc, argv, Col{}, Row{}, Row{}); } else diff --git a/example/ck_tile/36_copy/test_copy.cpp b/example/ck_tile/36_copy/test_copy.cpp index ef722bd60f..b235eb852d 100644 --- a/example/ck_tile/36_copy/test_copy.cpp +++ b/example/ck_tile/36_copy/test_copy.cpp @@ -53,10 +53,10 @@ bool run(const ck_tile::ArgParser& arg_parser) x_buf.ToDevice(x_host.data()); - using BlockWaves = ck_tile::sequence<2, 1>; - using BlockTile = ck_tile::sequence<64, 8>; - using WaveTile = ck_tile::sequence<64, 8>; - using Vector = ck_tile::sequence<1, 8>; + using BlockWaves = ck_tile::sequence<2, 1>; + using BlockTile = ck_tile::sequence<64, 8>; + using WaveTile = ck_tile::sequence<64, 8>; + using Vector = ck_tile::sequence<1, 8>; constexpr bool AsyncCopy = true; ck_tile::index_t kGridSize = (m / BlockTile::at(ck_tile::number<0>{})); diff --git a/example/ck_tile/36_copy/test_copy.hpp b/example/ck_tile/36_copy/test_copy.hpp index e694f727bf..1368872f6d 100644 --- a/example/ck_tile/36_copy/test_copy.hpp +++ b/example/ck_tile/36_copy/test_copy.hpp @@ -53,8 +53,8 @@ struct TileCopyShape template struct TileCopyProblem { - using XDataType = remove_cvref_t; - using BlockShape = remove_cvref_t; + using XDataType = remove_cvref_t; + using BlockShape = remove_cvref_t; static constexpr bool AsyncCopy = AsyncCopy_; }; @@ -159,7 +159,8 @@ struct TileCopy if(my_id == warp_id) { - if constexpr (AsyncCopy == false) { + if constexpr(AsyncCopy == false) + { // load from DRAM to registers load_tile(dram_tile, x_block_window); @@ -171,9 +172,11 @@ struct TileCopy // store from registers to DRAM store_tile(y_block_window, dram_tile); - } else { + } + else + { async_load_tile(x_block_lds_window_no_dist, x_block_window); - + load_tile(dram_tile, x_block_lds_window); // store from registers to DRAM diff --git a/include/ck_tile/core/tensor/buffer_view.hpp b/include/ck_tile/core/tensor/buffer_view.hpp index 8e8ec7343c..5cae332007 100644 --- a/include/ck_tile/core/tensor/buffer_view.hpp +++ b/include/ck_tile/core/tensor/buffer_view.hpp @@ -454,6 +454,7 @@ struct buffer_view, t_per_x, Coherence>( smem, src_wave_buffer_resource, diff --git a/include/ck_tile/core/tensor/tile_window.hpp b/include/ck_tile/core/tensor/tile_window.hpp index 3f872d9c17..1a8761ebff 100644 --- a/include/ck_tile/core/tensor/tile_window.hpp +++ b/include/ck_tile/core/tensor/tile_window.hpp @@ -345,8 +345,7 @@ struct tile_window_with_static_distribution #if defined(__HIP_DEVICE_COMPILE__) using LdsTileWindow = remove_cvref_t; using LdsDataType = typename LdsTileWindow::DataType; -#if defined(__gfx950__) - using Traits = typename Base::Traits; + using Traits = typename Base::Traits; using vector_t = typename Traits::vector_t; using SFC_Ys = typename Traits::SFC_Ys; @@ -385,69 +384,6 @@ struct tile_window_with_static_distribution } }); }); -#else - // issues * warps * lanes - static_assert(LdsTileWindow::get_num_of_dimension() == 3); // TODO: hard coded - - // TODO: LDS offset is not good for intrinsic based implementation(compiler can't figure out - // dependency) hence avoid use offset based solution. size_per_buf should be zero how to - // check?) - constexpr index_t size_per_buf = - lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( - make_tuple(number<0>{}, number<0>{}, number<0>{})); - - constexpr index_t size_per_wave = - lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( - make_tuple(number<0>{}, number<1>{}, number<0>{})) - - size_per_buf; - - constexpr index_t size_per_issue = - lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( - make_tuple(number<1>{}, number<0>{}, number<0>{})) - - size_per_buf; - - const index_t m0_init_value = size_per_buf + size_per_wave * get_warp_id(); - - using Traits = typename Base::Traits; - - using vector_t = typename Traits::vector_t; - using SFC_Ys = typename Traits::SFC_Ys; - - // TODO: we force CK_TILE_LDS_ADDR - CK_TILE_LDS_ADDR LdsDataType* smem = - lds_tile.get_bottom_tensor_view().get_buffer_view().p_data_ + m0_init_value; - - // loop over thread tensor space [y0, y1, ...] - static_for<0, NumCoord, 1>{}([&](auto iCoord) { - /// TODO: use structure binding (to be captured later) if compiled in C++20 - auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; - auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; - - static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { - constexpr auto iAccess = number{}; - - // read from bottom tensor - this->get_bottom_tensor_view().template async_get_vectorized_elements( - smem, bottom_tensor_thread_coord, 0, bool_constant{}); - - // move thread coordinate - if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) - { - constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); - - constexpr auto idx_diff_ps_ys = container_concat( - generate_tuple([&](auto) { return number<0>{}; }, number{}), - idx_diff_ys); - - Base::move_window_adaptor_and_bottom_tensor_thread_coordinate( - window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); - - smem += size_per_issue; // Note we manually increase the per-issue - offset - } - }); - }); -#endif #else ignore = lds_tile; #endif diff --git a/include/ck_tile/core/tensor/tile_window_linear.hpp b/include/ck_tile/core/tensor/tile_window_linear.hpp index e142c211cd..9f012a1e9d 100644 --- a/include/ck_tile/core/tensor/tile_window_linear.hpp +++ b/include/ck_tile/core/tensor/tile_window_linear.hpp @@ -186,10 +186,7 @@ struct tile_window_linear const typename Base::WindowLengths& window_lengths, const typename Base::BottomTensorIndex& window_origin, const typename Base::TileDstr& tile_distribution) - : - cached_coords_{}, - cached_window_adaptor_coords_{}, - cached_flags_{} + : cached_coords_{}, cached_window_adaptor_coords_{}, cached_flags_{} { this->bottom_tensor_view_ = bottom_tensor_view; this->window_lengths_ = window_lengths; @@ -566,7 +563,6 @@ struct tile_window_linear static_assert(Base::BottomTensorView::buffer_view::get_address_space() == address_space_enum::global); -#if defined(__gfx950__) // loop over thread tensor space [y0, y1, ...] auto issue = [&](auto i_access_) { constexpr auto IAccess = number{}; @@ -593,55 +589,6 @@ struct tile_window_linear bottom_tensor_flag, bool_constant{}); }; -#else - // issues * warps * lanes - static_assert(LdsTileWindow::get_num_of_dimension() == 3); // TODO: hard coded - - // TODO: LDS offset is not good for intrinsic based implementation(compiler can't figure out - // dependency) hence avoid use offset based solution. size_per_buf should be zero (how to - // check?) - constexpr index_t size_per_buf = - lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( - make_tuple(number<0>{}, number<0>{}, number<0>{})); - - constexpr index_t size_per_wave = - lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( - make_tuple(number<0>{}, number<1>{}, number<0>{})) - - size_per_buf; - - constexpr index_t size_per_issue = - lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset( - make_tuple(number<1>{}, number<0>{}, number<0>{})) - - size_per_buf; - - const index_t m0_init_value = size_per_buf + size_per_wave * get_warp_id(); - - // TODO: we force CK_TILE_LDS_ADDR - CK_TILE_LDS_ADDR LdsDataType* smem = - lds_tile.get_bottom_tensor_view().get_buffer_view().p_data_ + m0_init_value; - - // loop over thread tensor space [y0, y1, ...] - 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]; - - // read from bottom tensor - this->get_bottom_tensor_view().template async_get_vectorized_elements( - smem, - bottom_tensor_thread_coord, - 0, - bottom_tensor_flag, - bool_constant{}); - - // move thread coordinate - if constexpr(i_access_ != (NumAccess - 1)) - { - smem += size_per_issue; // Note we manually increase the per-issue offset - } - }; -#endif WINDOW_DISPATCH_ISSUE(); } @@ -959,7 +906,7 @@ struct tile_window_linear if constexpr(need_save_non_linear_coord) { - cached_coords_(non_linear_id) = bottom_tensor_thread_coord_tmp; + cached_coords_(non_linear_id) = bottom_tensor_thread_coord_tmp; cached_window_adaptor_coords_(non_linear_id) = window_adaptor_thread_coord_tmp; } @@ -980,7 +927,8 @@ struct tile_window_linear // this contains: array cached_coords_; - array cached_window_adaptor_coords_; + array + cached_window_adaptor_coords_; array cached_flags_; }; diff --git a/include/ck_tile/ops/gemm.hpp b/include/ck_tile/ops/gemm.hpp index 80f8f2d2ee..a1d37f0824 100644 --- a/include/ck_tile/ops/gemm.hpp +++ b/include/ck_tile/ops/gemm.hpp @@ -39,7 +39,6 @@ #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp" -#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" #include "ck_tile/ops/gemm/pipeline/tile_gemm_shape.hpp" 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 deleted file mode 100644 index 6cad68b58d..0000000000 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp +++ /dev/null @@ -1,467 +0,0 @@ -// 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_async_default_policy.hpp" - -namespace ck_tile { - -// A Tile Window: global memory -// B Tile Window: global memory -// C Distributed tensor: register -template -struct BaseGemmPipelineAgBgCrCompAsync -{ - 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 BaseGemmPipelineAgBgCrCompAsync -{ - using Base = BaseGemmPipelineAgBgCrCompAsync; - 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 - { - ignore = a_element_func; - ignore = b_element_func; - 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()); - - 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); - - 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"); - 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); - - 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); - block_sync_lds(); - 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 constexpr(TailNum == TailNumber::Three) - { - // 3 - { - 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); - block_gemm(c_block_tile, a_block_tile0, b_block_tile0); - } - // 2 - { - buffer_load_fence(0); - 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 - { - buffer_load_fence(0); - 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 diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_default_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_default_policy.hpp deleted file mode 100644 index fd7b112f93..0000000000 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_default_policy.hpp +++ /dev/null @@ -1,89 +0,0 @@ -// 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/warp/warp_gemm_dispatcher.hpp" -#include "ck_tile/ops/common/tensor_layout.hpp" -#include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" - -namespace ck_tile { -// Default policy for GemmPipelineAGmemBGmemCregComputeV4, except the block gemm method, it shares -// the same vector size implementation, SmemSize, Global memory tile distiribution as the -// UniversalGemm Pipeline Policy. -// Default policy class should not be templated, put template on -// member functions instead. -struct GemmPipelineAgBgCrCompAsyncDefaultPolicy - : public UniversalGemmBasePolicy -{ - static constexpr auto ATileAccessPattern = tile_distribution_pattern::warp_raked; - static constexpr auto BTileAccessPattern = tile_distribution_pattern::warp_raked; - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeALdsBlockDescriptor() - { - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPack = GetSmemPackA(); - - constexpr auto a_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, number{}, number{}), - make_tuple(number{}, number{}, number<1>{}), - number{}, - number<1>{}); - - return transform_tensor_descriptor( - a_lds_block_desc_0, - make_tuple( - make_pass_through_transform(number{}), - make_merge_transform(make_tuple(number{}, number{}))), - make_tuple(sequence<1>{}, sequence<0, 2>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeBLdsBlockDescriptor() - { - constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPack = GetSmemPackB(); - - constexpr auto b_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, number{}, number{}), - make_tuple(number{}, number{}, number<1>{}), - number{}, - number<1>{}); - - return transform_tensor_descriptor( - b_lds_block_desc_0, - make_tuple( - make_pass_through_transform(number{}), - make_merge_transform(make_tuple(number{}, number{}))), - make_tuple(sequence<1>{}, sequence<0, 2>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto GetBlockGemm() - { - using AccDataType = float; - using BlockWarps = typename Problem::BlockGemmShape::BlockWarps; - using WarpTile = typename Problem::BlockGemmShape::WarpTile; - using WarpGemm = WarpGemmMfmaDispatcher; - using BlockGemmPolicy = BlockGemmARegBRegCRegV1CustomPolicy; - - return BlockGemmARegBRegCRegV1{}; - } -}; -} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp index 8b78333123..eb47d9bad6 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp @@ -455,8 +455,8 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 // prefetch // global read 0 - Base::GlobalPrefetchAsync(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetchAsync(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); @@ -485,8 +485,8 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); } - Base::GlobalPrefetchAsync(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetchAsync(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); block_sync_lds(); block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); @@ -524,10 +524,8 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); } - Base::GlobalPrefetchAsync( - a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetchAsync( - b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); + Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp index 43b174a22b..1863274b44 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp @@ -526,115 +526,116 @@ struct UniversalGemmBasePolicy BTileAccessPattern, NumWaveGroups>; } - template - CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledARegTileDistribution() + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledARegTileDistribution() + { + using ALayout = remove_cvref_t; + static_assert(std::is_same_v); + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t VecLoadSize = GetVectorSizeA(); + constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + + using TileEncodingPattern = TileDistributionEncodingPattern2D; + return TileEncodingPattern::MakeShuffled2DStaticTileDistribution(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledBRegTileDistribution() + { + using BLayout = remove_cvref_t; + static_assert(std::is_same_v); + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t VecLoadSize = GetVectorSizeB(); + constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + + using TileEncodingPattern = TileDistributionEncodingPattern2D; + return TileEncodingPattern::MakeShuffled2DStaticTileDistribution(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackA() + { + using BlockGemm = remove_cvref_t())>; + constexpr index_t KPack = BlockGemm::Traits::KPack; + return KPack; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackB() + { + using BlockGemm = remove_cvref_t())>; + constexpr index_t KPack = BlockGemm::Traits::KPack; + return KPack; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeA() + { + constexpr auto a_lds_desc = MakeALdsBlockDescriptor(); + constexpr index_t smem_size_a = integer_least_multiple( + sizeof(typename Problem::ADataType) * a_lds_desc.get_element_space_size(), 16); + return smem_size_a; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeB() + { + constexpr auto b_lds_desc = MakeBLdsBlockDescriptor(); + constexpr index_t smem_size_b = integer_least_multiple( + sizeof(typename Problem::BDataType) * b_lds_desc.get_element_space_size(), 16); + return smem_size_b; + } + + template + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + constexpr index_t smem_size_a = GetSmemSizeA(); + constexpr index_t smem_size_b = GetSmemSizeB(); + + return smem_size_a + smem_size_b; + } + }; + + // UniversalGemm Policy + struct UniversalGemmPipelineAgBgCrPolicy + : public UniversalGemmBasePolicy { - using ALayout = remove_cvref_t; - static_assert(std::is_same_v); - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t VecLoadSize = GetVectorSizeA(); - constexpr index_t NumWaveGroups = Problem::NumWaveGroups; - - using TileEncodingPattern = TileDistributionEncodingPattern2D; - return TileEncodingPattern::MakeShuffled2DStaticTileDistribution(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledBRegTileDistribution() - { - using BLayout = remove_cvref_t; - static_assert(std::is_same_v); - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t VecLoadSize = GetVectorSizeB(); - constexpr index_t NumWaveGroups = Problem::NumWaveGroups; - - using TileEncodingPattern = TileDistributionEncodingPattern2D; - return TileEncodingPattern::MakeShuffled2DStaticTileDistribution(); - } - - template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackA() - { - using BlockGemm = remove_cvref_t())>; - constexpr index_t KPack = BlockGemm::Traits::KPack; - return KPack; - } - - template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackB() - { - using BlockGemm = remove_cvref_t())>; - constexpr index_t KPack = BlockGemm::Traits::KPack; - return KPack; - } - - template - CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeA() - { - constexpr auto a_lds_desc = MakeALdsBlockDescriptor(); - constexpr index_t smem_size_a = integer_least_multiple( - sizeof(typename Problem::ADataType) * a_lds_desc.get_element_space_size(), 16); - return smem_size_a; - } - - template - CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeB() - { - constexpr auto b_lds_desc = MakeBLdsBlockDescriptor(); - constexpr index_t smem_size_b = integer_least_multiple( - sizeof(typename Problem::BDataType) * b_lds_desc.get_element_space_size(), 16); - return smem_size_b; - } - - template - CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() - { - constexpr index_t smem_size_a = GetSmemSizeA(); - constexpr index_t smem_size_b = GetSmemSizeB(); - - return smem_size_a + smem_size_b; - } -}; - -// UniversalGemm Policy -struct UniversalGemmPipelineAgBgCrPolicy - : public UniversalGemmBasePolicy -{ - template - CK_TILE_HOST_DEVICE static constexpr auto GetBlockGemm() - { - using BlockWarps = typename Problem::BlockGemmShape::BlockWarps; - using WarpTile = typename Problem::BlockGemmShape::WarpTile; - using WarpGemm = WarpGemmMfmaDispatcher; - using BlockGemmPolicy = BlockGemmASmemBSmemCRegV1CustomPolicy; - return BlockUniversalGemmAsBsCr{}; - } -}; + template + CK_TILE_HOST_DEVICE static constexpr auto GetBlockGemm() + { + using BlockWarps = typename Problem::BlockGemmShape::BlockWarps; + using WarpTile = typename Problem::BlockGemmShape::WarpTile; + using WarpGemm = WarpGemmMfmaDispatcher; + using BlockGemmPolicy = + BlockGemmASmemBSmemCRegV1CustomPolicy; + return BlockUniversalGemmAsBsCr{}; + } + }; } // namespace ck_tile