From d45499f110656617d363053d573d8d13c2ce6833 Mon Sep 17 00:00:00 2001 From: "D.Lejeune" Date: Fri, 4 Jul 2025 10:42:16 -0500 Subject: [PATCH] Skip A/B LDS for universal GEMM --- example/ck_tile/03_gemm/gemm_utils.hpp | 4 +- example/ck_tile/03_gemm/run_gemm_example.inc | 10 + example/ck_tile/03_gemm/universal_gemm.cpp | 3 +- .../block/block_universal_gemm_ar_br_cr.hpp | 265 +++--------------- .../block/block_universal_gemm_ar_bs_cr.hpp | 244 +++------------- .../block/block_universal_gemm_as_br_cr.hpp | 245 +++------------- .../block/block_universal_gemm_as_bs_cr.hpp | 247 +++------------- .../gemm/block/block_universal_gemm_base.hpp | 217 ++++++++++++++ .../ops/gemm/kernel/batched_gemm_kernel.hpp | 4 +- .../ck_tile/ops/gemm/kernel/gemm_kernel.hpp | 19 +- 10 files changed, 411 insertions(+), 847 deletions(-) create mode 100644 include/ck_tile/ops/gemm/block/block_universal_gemm_base.hpp diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp index 0b4361d3b2..750726b37a 100644 --- a/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/example/ck_tile/03_gemm/gemm_utils.hpp @@ -16,7 +16,7 @@ #define CK_TILE_PIPELINE_COMPUTE_V4 3 #ifndef CK_TILE_PIPELINE_DEFAULT -#define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_COMPUTE_V3 +#define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_MEMORY #endif #if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY) @@ -216,6 +216,8 @@ auto create_args(int argc, char* argv[]) .insert("stride_a", "0", "Tensor A stride") .insert("stride_b", "0", "Tensor B stride") .insert("stride_c", "0", "Tensor C stride") + .insert("skip_a_lds", "0", "Flag to skip loading A tensor tile into LDS (0: false, 1: true)") + .insert("skip_b_lds", "0", "Flag to skip loading B tensor tile into LDS (0: false, 1: true)") .insert("v", "2", "0. No validation, 1. Validation on CPU, 2. Validation on GPU") .insert("prec", "fp16", "data type. fp16/bf16/fp8/bf8") .insert("warmup", "50", "number of iterations before benchmark the kernel") diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index bf58dc2f37..b9af203524 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -163,6 +163,8 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, ck_tile::index_t stride_B, ck_tile::index_t stride_C, ck_tile::index_t kbatch, + bool skip_a_lds, + bool skip_b_lds, int n_warmup, int n_repeat) { @@ -177,6 +179,8 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, args.stride_A = stride_A; args.stride_B = stride_B; args.stride_C = stride_C; + args.skip_a_lds = skip_a_lds; + args.skip_b_lds = skip_b_lds; float ave_time = gemm_calc( @@ -194,6 +198,8 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, << " C_Layout=" << CLayout::name << " A_Type=" << DataTypeTraits::name << " B_Type=" << DataTypeTraits::name << " C_Type=" << DataTypeTraits::name + << " SkipALds=" << skip_a_lds + << " SkipBLds=" << skip_b_lds << " StructuredSparsity=" << (GemmConfig::UseStructuredSparsity ? "on" : "off") << " : " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl; @@ -231,6 +237,8 @@ int run_gemm_example_with_layouts(int argc, int n_warmup = arg_parser.get_int("warmup"); int n_repeat = arg_parser.get_int("repeat"); ck_tile::index_t init_method = arg_parser.get_int("init"); + bool skip_a_lds = arg_parser.get_int("skip_a_lds"); + bool skip_b_lds = arg_parser.get_int("skip_b_lds"); stride_A = ck_tile::get_default_stride(M, K, stride_A, is_row_major(a_layout)); stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout)); @@ -317,6 +325,8 @@ int run_gemm_example_with_layouts(int argc, stride_B, stride_C, kbatch, + skip_a_lds, + skip_b_lds, n_warmup, n_repeat); diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index 16e03d2723..283fec1c4a 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -345,6 +345,7 @@ int run_gemm_example(int argc, char* argv[]) { return run_gemm_example_prec_type(a_layout, b_layout, argc, argv); } +#if(CK_TILE_PIPELINE_DEFAULT != CK_TILE_PIPELINE_MEMORY) else if(data_type == "bf16") { return run_gemm_example_prec_type(a_layout, b_layout, argc, argv); @@ -359,7 +360,7 @@ int run_gemm_example(int argc, char* argv[]) return run_gemm_example_prec_type( a_layout, b_layout, argc, argv); } - +#endif #if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V3) else if(data_type == "pk_int4_t") { diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp index 68c7ba9f95..6cea4b02dd 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp @@ -7,6 +7,7 @@ #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" +#include "block_universal_gemm_base.hpp" namespace ck_tile { @@ -14,197 +15,44 @@ namespace ck_tile { // B is block distributed tensor // C is block distributed tensor template -struct BlockUniversalGemmArBrCr +struct BlockUniversalGemmArBrCr : public BlockUniversalGemmBase { private: - // TODO: This should be in Policy - UniversalGemmPolicyBase ? - template - struct GemmTraits_ - { - using Problem = remove_cvref_t; - using Policy = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; - - static constexpr index_t kBlockSize = Problem::kBlockSize; - static constexpr auto Scheduler = Problem::Scheduler; - - static constexpr index_t MPerBlock = BlockGemmShape::kM; - static constexpr index_t NPerBlock = BlockGemmShape::kN; - static constexpr index_t KPerBlock = BlockGemmShape::kK; - - static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); - - using WarpGemm = remove_cvref_t())>; - - static constexpr index_t MWarp = config.template at<1>(); - static constexpr index_t NWarp = config.template at<2>(); - - using I0 = number<0>; - using I1 = number<1>; - - static_assert(MWarp == BlockGemmShape::BlockWarps::at(I0{}), - "Error! WarpGemm's MWarp is not consisten with BlockGemmShape!"); - static_assert(NWarp == BlockGemmShape::BlockWarps::at(I1{}), - "Error! WarpGemm's NWarp is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kM == BlockGemmShape::WarpTile::at(I0{}), - "Error! WarpGemm's M is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kN == BlockGemmShape::WarpTile::at(I1{}), - "Error! WarpGemm's N is not consisten with BlockGemmShape!"); - - static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WarpGemm::kM); - static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); - static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - - static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, - "Error! Warps should cover all Block tile!"); - static_assert(NIterPerWarp * NWarp * WarpGemm::kN == NPerBlock, - "Error! Warps should cover all Block tile!"); - - static constexpr index_t MPerBlockPerIter = MWarp * WarpGemm::kM; - static constexpr index_t NPerBlockPerIter = NWarp * WarpGemm::kN; - static constexpr index_t KPerBlockPerIter = WarpGemm::kK; - - // Controls how many MAC clusters (MFMA blocks) we have per wave - // Ie if - // InterWaveSchedulingMacClusters = 1; - // KPerBlock == 32 - // WarpGemm::kK = 8 - // Then we would group all 4 WarpGemms into single MAC cluster. - // But if we would set InterWaveSchedulingMacClusters = 2, then we would - // split those 4 warp gemms into two groups. - static constexpr index_t InterWaveSchedulingMacClusters = 1; - - // should be at least equal to: WarpGemm::Impl::kABKPerLane - static constexpr index_t KPack = WarpGemm::kKPerThread; - static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread; - }; + using Base = BlockUniversalGemmBase; + using typename Base::ADataType; + using typename Base::BDataType; + using typename Base::CDataType; + using typename Base::WarpGemm; + using typename Base::AWarpTensor; + using typename Base::BWarpTensor; + using typename Base::CWarpTensor; + using Base::MIterPerWarp; + using Base::NIterPerWarp; + using Base::KIterPerWarp; + using Base::a_warp_y_index_zeros; + using Base::b_warp_y_index_zeros; + using Base::c_warp_y_index_zeros; + using Base::a_warp_y_lengths; + using Base::b_warp_y_lengths; + using Base::c_warp_y_lengths; + using Base::Scheduler; + using GemmTraits = typename Base::template GemmTraits_; + public: - using Traits = GemmTraits_; - - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - - using WarpGemm = remove_cvref_t; - - static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; - static constexpr index_t MIterPerWarp = Traits::MIterPerWarp; - static constexpr index_t NIterPerWarp = Traits::NIterPerWarp; - - static constexpr index_t MWarp = Traits::MWarp; - static constexpr index_t NWarp = Traits::NWarp; - - static constexpr auto Scheduler = Traits::Scheduler; - - using AWarpDstr = typename WarpGemm::AWarpDstr; - using BWarpDstr = typename WarpGemm::BWarpDstr; - using CWarpDstr = typename WarpGemm::CWarpDstr; - - using AWarpTensor = typename WarpGemm::AWarpTensor; - using BWarpTensor = typename WarpGemm::BWarpTensor; - using CWarpTensor = typename WarpGemm::CWarpTensor; - - static constexpr auto a_warp_y_lengths = - to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto b_warp_y_lengths = - to_sequence(BWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto c_warp_y_lengths = - to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - - static constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; - - static constexpr index_t APackedSize = - ck_tile::numeric_traits>::PackedSize; - static constexpr index_t BPackedSize = - ck_tile::numeric_traits>::PackedSize; - - using I0 = number<0>; - using I1 = number<1>; - - CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto a_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); - - return a_block_dstr_encode; - } - - CK_TILE_DEVICE static constexpr auto MakeBBlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto b_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); - - return b_block_dstr_encode; - } + using Base::MakeBBlockDistributionEncode; + using Base::MakeABlockDistributionEncode; + using Base::MakeCBlockTile; private: - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - constexpr index_t UnaryOpSize = 8; - const element_wise::PassThroughPack8 elementwise_op{}; - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } - - template + template struct BlockGemmImpl { }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { // C += A * B template @@ -221,7 +69,7 @@ struct BlockUniversalGemmArBrCr "traits should be the same as correspoinding block window data type!"); // hot loop: - static_for<0, GemmTraits::KIterPerWarp, 1>{}([&](auto kIter) { + static_for<0, GemmTraits_::KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { // read A warp tensor from A block tensor AWarpTensor a_warp_tensor; @@ -259,8 +107,8 @@ struct BlockUniversalGemmArBrCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { // C += A * B template @@ -310,11 +158,11 @@ struct BlockUniversalGemmArBrCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { - static constexpr index_t KPerThread = GemmTraits::KPerThread; - static constexpr index_t NumMacClusters = GemmTraits::InterWaveSchedulingMacClusters; + static constexpr index_t KPerThread = GemmTraits_::KPerThread; + static constexpr index_t NumMacClusters = GemmTraits_::InterWaveSchedulingMacClusters; static constexpr index_t KPerInnerLoop = ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); static constexpr index_t KRepeat = KPerThread / KPerInnerLoop; @@ -398,61 +246,38 @@ struct BlockUniversalGemmArBrCr if constexpr(kInnerIter.value == 0 && mIter.value == 0 && nIter.value == 0) { - __builtin_amdgcn_sched_barrier(0); - __builtin_amdgcn_s_setprio(1); - __builtin_amdgcn_sched_barrier(0); } }); }); }); - - __builtin_amdgcn_sched_barrier(0); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_sched_barrier(0); }); } }; public: - CK_TILE_DEVICE static constexpr auto MakeCBlockTile() - { - constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - - constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{}); - constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode); - auto c_block_tensor = make_static_distributed_tensor(c_block_dstr); - - return c_block_tensor; - } - // C += A * B - template + template CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor, const ARegBlockTensor& a_block_tensor, - const BSmemBlockWindow& b_block_window) + const BRegBlockTensor& b_block_tensor) { - block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_window); + block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_tensor); } // C = A * B - template + template CK_TILE_DEVICE auto operator()(const ARegBlockTensor& a_block_tensor, - const BSmemBlockWindow& b_block_window) + const BRegBlockTensor& b_block_tensor) { - auto c_block_tensor = MakeCBlockTile(); - block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_window); + auto c_block_tensor = Base::MakeCBlockTile(); + + block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_tensor); + return c_block_tensor; } private: - BlockGemmImpl block_gemm_impl_{}; + BlockGemmImpl block_gemm_impl_{}; }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp index 05e9b210e7..246285dbac 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp @@ -7,6 +7,7 @@ #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" +#include "block_universal_gemm_base.hpp" namespace ck_tile { @@ -14,197 +15,44 @@ namespace ck_tile { // B is block window on shared memory // C is block distributed tensor template -struct BlockUniversalGemmArBsCr +struct BlockUniversalGemmArBsCr : public BlockUniversalGemmBase { private: - // TODO: This should be in Policy - UniversalGemmPolicyBase ? - template - struct GemmTraits_ - { - using Problem = remove_cvref_t; - using Policy = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; + using Base = BlockUniversalGemmBase; + using typename Base::ADataType; + using typename Base::BDataType; + using typename Base::CDataType; + using typename Base::ComputeDataType; + using typename Base::WarpGemm; + using typename Base::AWarpTensor; + using typename Base::BWarpTensor; + using typename Base::CWarpTensor; + using Base::MIterPerWarp; + using Base::NIterPerWarp; + using Base::a_warp_y_index_zeros; + using Base::b_warp_y_index_zeros; + using Base::c_warp_y_index_zeros; + using Base::a_warp_y_lengths; + using Base::b_warp_y_lengths; + using Base::c_warp_y_lengths; + using Base::Scheduler; - static constexpr index_t kBlockSize = Problem::kBlockSize; - static constexpr auto Scheduler = Problem::Scheduler; - - static constexpr index_t MPerBlock = BlockGemmShape::kM; - static constexpr index_t NPerBlock = BlockGemmShape::kN; - static constexpr index_t KPerBlock = BlockGemmShape::kK; - - static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); - - using WarpGemm = remove_cvref_t())>; - - static constexpr index_t MWarp = config.template at<1>(); - static constexpr index_t NWarp = config.template at<2>(); - - using I0 = number<0>; - using I1 = number<1>; - - static_assert(MWarp == BlockGemmShape::BlockWarps::at(I0{}), - "Error! WarpGemm's MWarp is not consisten with BlockGemmShape!"); - static_assert(NWarp == BlockGemmShape::BlockWarps::at(I1{}), - "Error! WarpGemm's NWarp is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kM == BlockGemmShape::WarpTile::at(I0{}), - "Error! WarpGemm's M is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kN == BlockGemmShape::WarpTile::at(I1{}), - "Error! WarpGemm's N is not consisten with BlockGemmShape!"); - - static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WarpGemm::kM); - static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); - static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - - static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, - "Error! Warps should cover all Block tile!"); - static_assert(NIterPerWarp * NWarp * WarpGemm::kN == NPerBlock, - "Error! Warps should cover all Block tile!"); - - static constexpr index_t MPerBlockPerIter = MWarp * WarpGemm::kM; - static constexpr index_t NPerBlockPerIter = NWarp * WarpGemm::kN; - static constexpr index_t KPerBlockPerIter = WarpGemm::kK; - - // Controls how many MAC clusters (MFMA blocks) we have per wave - // Ie if - // InterWaveSchedulingMacClusters = 1; - // KPerBlock == 32 - // WarpGemm::kK = 8 - // Then we would group all 4 WarpGemms into single MAC cluster. - // But if we would set InterWaveSchedulingMacClusters = 2, then we would - // split those 4 warp gemms into two groups. - static constexpr index_t InterWaveSchedulingMacClusters = 1; - - // should be at least equal to: WarpGemm::Impl::kABKPerLane - static constexpr index_t KPack = WarpGemm::kKPerThread; - static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread; - }; + using Base::KIterPerWarp; + using GemmTraits = typename Base::template GemmTraits_; public: - using Traits = GemmTraits_; - - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - - using WarpGemm = remove_cvref_t; - - static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; - static constexpr index_t MIterPerWarp = Traits::MIterPerWarp; - static constexpr index_t NIterPerWarp = Traits::NIterPerWarp; - - static constexpr index_t MWarp = Traits::MWarp; - static constexpr index_t NWarp = Traits::NWarp; - - static constexpr auto Scheduler = Traits::Scheduler; - - using AWarpDstr = typename WarpGemm::AWarpDstr; - using BWarpDstr = typename WarpGemm::BWarpDstr; - using CWarpDstr = typename WarpGemm::CWarpDstr; - - using AWarpTensor = typename WarpGemm::AWarpTensor; - using BWarpTensor = typename WarpGemm::BWarpTensor; - using CWarpTensor = typename WarpGemm::CWarpTensor; - - static constexpr auto a_warp_y_lengths = - to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto b_warp_y_lengths = - to_sequence(BWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto c_warp_y_lengths = - to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - - static constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; - - static constexpr index_t APackedSize = - ck_tile::numeric_traits>::PackedSize; - static constexpr index_t BPackedSize = - ck_tile::numeric_traits>::PackedSize; - - using I0 = number<0>; - using I1 = number<1>; - - CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto a_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); - - return a_block_dstr_encode; - } - - CK_TILE_DEVICE static constexpr auto MakeBBlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto b_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); - - return b_block_dstr_encode; - } + using Base::MakeBBlockDistributionEncode; + using Base::MakeABlockDistributionEncode; + using Base::MakeCBlockTile; private: - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - constexpr index_t UnaryOpSize = 8; - const element_wise::PassThroughPack8 elementwise_op{}; - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } - - template + template struct BlockGemmImpl { }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { static constexpr auto BLdsTileDistr = decltype(make_static_tile_distribution(MakeBBlockDistributionEncode())){}; @@ -273,8 +121,8 @@ struct BlockUniversalGemmArBsCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { static constexpr auto BLdsTileDistr = decltype(make_static_tile_distribution(MakeBBlockDistributionEncode())){}; @@ -305,7 +153,7 @@ struct BlockUniversalGemmArBsCr "C block tensor data type!"); // hot loop: - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + static_for<0, GemmTraits_::KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { // read A warp tensor from A block tensor AWarpTensor a_warp_tensor; @@ -343,11 +191,11 @@ struct BlockUniversalGemmArBsCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { - static constexpr index_t KPerThread = GemmTraits::KPerThread; - static constexpr index_t NumMacClusters = GemmTraits::InterWaveSchedulingMacClusters; + static constexpr index_t KPerThread = GemmTraits_::KPerThread; + static constexpr index_t NumMacClusters = GemmTraits_::InterWaveSchedulingMacClusters; static constexpr index_t KPerInnerLoop = ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); static constexpr index_t KRepeat = KPerThread / KPerInnerLoop; @@ -476,24 +324,6 @@ struct BlockUniversalGemmArBsCr }; public: - CK_TILE_DEVICE static constexpr auto MakeCBlockTile() - { - constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - - constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{}); - constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode); - auto c_block_tensor = make_static_distributed_tensor(c_block_dstr); - - return c_block_tensor; - } - template CK_TILE_DEVICE void LocalPrefetchB(const BSmemBlockWindow& b_block_window) { @@ -514,13 +344,13 @@ struct BlockUniversalGemmArBsCr CK_TILE_DEVICE auto operator()(const ARegBlockTensor& a_block_tensor, const BSmemBlockWindow& b_block_window) { - auto c_block_tensor = MakeCBlockTile(); + auto c_block_tensor = Base::MakeCBlockTile(); block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_window); return c_block_tensor; } private: - BlockGemmImpl block_gemm_impl_{}; + BlockGemmImpl block_gemm_impl_{}; }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp index 4984999a2c..6001a81c16 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp @@ -7,6 +7,7 @@ #include "ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" +#include "block_universal_gemm_base.hpp" namespace ck_tile { @@ -14,197 +15,45 @@ namespace ck_tile { // B is block distributed tensor // C is block distributed tensor template -struct BlockUniversalGemmAsBrCr +struct BlockUniversalGemmAsBrCr : public BlockUniversalGemmBase { private: - // TODO: This should be in Policy - UniversalGemmPolicyBase ? - template - struct GemmTraits_ - { - using Problem = remove_cvref_t; - using Policy = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; + using Base = BlockUniversalGemmBase; + using typename Base::ADataType; + using typename Base::BDataType; + using typename Base::CDataType; + using typename Base::ComputeDataType; + using typename Base::WarpGemm; + using typename Base::AWarpTensor; + using typename Base::BWarpTensor; + using typename Base::CWarpTensor; + using Base::MIterPerWarp; + using Base::NIterPerWarp; + using Base::KIterPerWarp; + using Base::a_warp_y_index_zeros; + using Base::b_warp_y_index_zeros; + using Base::c_warp_y_index_zeros; + using Base::a_warp_y_lengths; + using Base::b_warp_y_lengths; + using Base::c_warp_y_lengths; + using Base::Scheduler; + using Base::load_interleaved_pk_type; - static constexpr index_t kBlockSize = Problem::kBlockSize; - static constexpr auto Scheduler = Problem::Scheduler; - - static constexpr index_t MPerBlock = BlockGemmShape::kM; - static constexpr index_t NPerBlock = BlockGemmShape::kN; - static constexpr index_t KPerBlock = BlockGemmShape::kK; - - static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); - - using WarpGemm = remove_cvref_t())>; - - static constexpr index_t MWarp = config.template at<1>(); - static constexpr index_t NWarp = config.template at<2>(); - - using I0 = number<0>; - using I1 = number<1>; - - static_assert(MWarp == BlockGemmShape::BlockWarps::at(I0{}), - "Error! WarpGemm's MWarp is not consisten with BlockGemmShape!"); - static_assert(NWarp == BlockGemmShape::BlockWarps::at(I1{}), - "Error! WarpGemm's NWarp is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kM == BlockGemmShape::WarpTile::at(I0{}), - "Error! WarpGemm's M is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kN == BlockGemmShape::WarpTile::at(I1{}), - "Error! WarpGemm's N is not consisten with BlockGemmShape!"); - - static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WarpGemm::kM); - static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); - static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - - static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, - "Error! Warps should cover all Block tile!"); - static_assert(NIterPerWarp * NWarp * WarpGemm::kN == NPerBlock, - "Error! Warps should cover all Block tile!"); - - static constexpr index_t MPerBlockPerIter = MWarp * WarpGemm::kM; - static constexpr index_t NPerBlockPerIter = NWarp * WarpGemm::kN; - static constexpr index_t KPerBlockPerIter = WarpGemm::kK; - - // Controls how many MAC clusters (MFMA blocks) we have per wave - // Ie if - // InterWaveSchedulingMacClusters = 1; - // KPerBlock == 32 - // WarpGemm::kK = 8 - // Then we would group all 4 WarpGemms into single MAC cluster. - // But if we would set InterWaveSchedulingMacClusters = 2, then we would - // split those 4 warp gemms into two groups. - static constexpr index_t InterWaveSchedulingMacClusters = 1; - - // should be at least equal to: WarpGemm::Impl::kABKPerLane - static constexpr index_t KPack = WarpGemm::kKPerThread; - static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread; - }; + using GemmTraits = typename Base::template GemmTraits_; public: - using Traits = GemmTraits_; - - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - - using WarpGemm = remove_cvref_t; - - static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; - static constexpr index_t MIterPerWarp = Traits::MIterPerWarp; - static constexpr index_t NIterPerWarp = Traits::NIterPerWarp; - - static constexpr index_t MWarp = Traits::MWarp; - static constexpr index_t NWarp = Traits::NWarp; - - static constexpr auto Scheduler = Traits::Scheduler; - - using AWarpDstr = typename WarpGemm::AWarpDstr; - using BWarpDstr = typename WarpGemm::BWarpDstr; - using CWarpDstr = typename WarpGemm::CWarpDstr; - - using AWarpTensor = typename WarpGemm::AWarpTensor; - using BWarpTensor = typename WarpGemm::BWarpTensor; - using CWarpTensor = typename WarpGemm::CWarpTensor; - - static constexpr auto a_warp_y_lengths = - to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto b_warp_y_lengths = - to_sequence(BWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto c_warp_y_lengths = - to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - - static constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; - - static constexpr index_t APackedSize = - ck_tile::numeric_traits>::PackedSize; - static constexpr index_t BPackedSize = - ck_tile::numeric_traits>::PackedSize; - - using I0 = number<0>; - using I1 = number<1>; - - CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto a_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); - - return a_block_dstr_encode; - } - - CK_TILE_DEVICE static constexpr auto MakeBBlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto b_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); - - return b_block_dstr_encode; - } + using Base::MakeABlockDistributionEncode; + using Base::MakeBBlockDistributionEncode; + using Base::MakeCBlockTile; private: - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - constexpr index_t UnaryOpSize = 8; - const element_wise::PassThroughPack8 elementwise_op{}; - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } - - template + template struct BlockGemmImpl { }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { static constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution(MakeABlockDistributionEncode())){}; @@ -242,7 +91,7 @@ struct BlockUniversalGemmAsBrCr b_warp_tile_.get_thread_buffer() = b_block_tensor.get_thread_buffer(); // hot loop: - static_for<0, GemmTraits::KIterPerWarp, 1>{}([&](auto kIter) { + static_for<0, GemmTraits_::KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { // read A warp tensor from A block tensor AWarpTensor a_warp_tensor; @@ -280,8 +129,8 @@ struct BlockUniversalGemmAsBrCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { static constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution(MakeABlockDistributionEncode())){}; @@ -357,11 +206,11 @@ struct BlockUniversalGemmAsBrCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { - static constexpr index_t KPerThread = GemmTraits::KPerThread; - static constexpr index_t NumMacClusters = GemmTraits::InterWaveSchedulingMacClusters; + static constexpr index_t KPerThread = GemmTraits_::KPerThread; + static constexpr index_t NumMacClusters = GemmTraits_::InterWaveSchedulingMacClusters; static constexpr index_t KPerInnerLoop = ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); static constexpr index_t KRepeat = KPerThread / KPerInnerLoop; @@ -489,24 +338,6 @@ struct BlockUniversalGemmAsBrCr }; public: - CK_TILE_DEVICE static constexpr auto MakeCBlockTile() - { - constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - - constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{}); - constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode); - auto c_block_tensor = make_static_distributed_tensor(c_block_dstr); - - return c_block_tensor; - } - template CK_TILE_DEVICE void LocalPrefetchA(const ASmemBlockWindow& a_block_window) { @@ -527,13 +358,13 @@ struct BlockUniversalGemmAsBrCr CK_TILE_DEVICE auto operator()(const ASmemBlockWindow& a_block_window, const BRegBlockTensor& b_block_tensor) { - auto c_block_tensor = MakeCBlockTile(); + auto c_block_tensor = Base::MakeCBlockTile(); block_gemm_impl_(c_block_tensor, a_block_window, b_block_tensor); return c_block_tensor; } private: - BlockGemmImpl block_gemm_impl_{}; + BlockGemmImpl block_gemm_impl_{}; }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp index ca660ff234..986dbf45eb 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp @@ -7,6 +7,7 @@ #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" +#include "block_universal_gemm_base.hpp" namespace ck_tile { @@ -14,197 +15,45 @@ namespace ck_tile { // B is block window on shared memory // C is block distributed tensor template -struct BlockUniversalGemmAsBsCr +struct BlockUniversalGemmAsBsCr : public BlockUniversalGemmBase { private: - // TODO: This should be in Policy - UniversalGemmPolicyBase ? - template - struct GemmTraits_ - { - using Problem = remove_cvref_t; - using Policy = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; + using Base = BlockUniversalGemmBase; + using typename Base::ADataType; + using typename Base::BDataType; + using typename Base::CDataType; + using typename Base::ComputeDataType; + using typename Base::WarpGemm; + using typename Base::AWarpTensor; + using typename Base::BWarpTensor; + using typename Base::CWarpTensor; + using Base::MIterPerWarp; + using Base::NIterPerWarp; + using Base::KIterPerWarp; + using Base::a_warp_y_index_zeros; + using Base::b_warp_y_index_zeros; + using Base::c_warp_y_index_zeros; + using Base::a_warp_y_lengths; + using Base::b_warp_y_lengths; + using Base::c_warp_y_lengths; + using Base::Scheduler; + using Base::load_interleaved_pk_type; - static constexpr index_t kBlockSize = Problem::kBlockSize; - static constexpr auto Scheduler = Problem::Scheduler; - - static constexpr index_t MPerBlock = BlockGemmShape::kM; - static constexpr index_t NPerBlock = BlockGemmShape::kN; - static constexpr index_t KPerBlock = BlockGemmShape::kK; - - static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); - - using WarpGemm = remove_cvref_t())>; - - static constexpr index_t MWarp = config.template at<1>(); - static constexpr index_t NWarp = config.template at<2>(); - - using I0 = number<0>; - using I1 = number<1>; - - static_assert(MWarp == BlockGemmShape::BlockWarps::at(I0{}), - "Error! WarpGemm's MWarp is not consisten with BlockGemmShape!"); - static_assert(NWarp == BlockGemmShape::BlockWarps::at(I1{}), - "Error! WarpGemm's NWarp is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kM == BlockGemmShape::WarpTile::at(I0{}), - "Error! WarpGemm's M is not consisten with BlockGemmShape!"); - static_assert(WarpGemm::kN == BlockGemmShape::WarpTile::at(I1{}), - "Error! WarpGemm's N is not consisten with BlockGemmShape!"); - - static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WarpGemm::kM); - static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); - static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - - static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, - "Error! Warps should cover all Block tile!"); - static_assert(NIterPerWarp * NWarp * WarpGemm::kN == NPerBlock, - "Error! Warps should cover all Block tile!"); - - static constexpr index_t MPerBlockPerIter = MWarp * WarpGemm::kM; - static constexpr index_t NPerBlockPerIter = NWarp * WarpGemm::kN; - static constexpr index_t KPerBlockPerIter = WarpGemm::kK; - - // Controls how many MAC clusters (MFMA blocks) we have per wave - // Ie if - // InterWaveSchedulingMacClusters = 1; - // KPerBlock == 32 - // WarpGemm::kK = 8 - // Then we would group all 4 WarpGemms into single MAC cluster. - // But if we would set InterWaveSchedulingMacClusters = 2, then we would - // split those 4 warp gemms into two groups. - static constexpr index_t InterWaveSchedulingMacClusters = 1; - - // should be at least equal to: WarpGemm::Impl::kABKPerLane - static constexpr index_t KPack = WarpGemm::kKPerThread; - static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread; - }; + using GemmTraits = typename Base::template GemmTraits_; public: - using Traits = GemmTraits_; - - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - - using WarpGemm = remove_cvref_t; - - static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; - static constexpr index_t MIterPerWarp = Traits::MIterPerWarp; - static constexpr index_t NIterPerWarp = Traits::NIterPerWarp; - - static constexpr index_t MWarp = Traits::MWarp; - static constexpr index_t NWarp = Traits::NWarp; - - static constexpr auto Scheduler = Traits::Scheduler; - - using AWarpDstr = typename WarpGemm::AWarpDstr; - using BWarpDstr = typename WarpGemm::BWarpDstr; - using CWarpDstr = typename WarpGemm::CWarpDstr; - - using AWarpTensor = typename WarpGemm::AWarpTensor; - using BWarpTensor = typename WarpGemm::BWarpTensor; - using CWarpTensor = typename WarpGemm::CWarpTensor; - - static constexpr auto a_warp_y_lengths = - to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto b_warp_y_lengths = - to_sequence(BWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - static constexpr auto c_warp_y_lengths = - to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); - - static constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t{}; - static constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; - - static constexpr index_t APackedSize = - ck_tile::numeric_traits>::PackedSize; - static constexpr index_t BPackedSize = - ck_tile::numeric_traits>::PackedSize; - - using I0 = number<0>; - using I1 = number<1>; - - CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto a_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); - - return a_block_dstr_encode; - } - - CK_TILE_DEVICE static constexpr auto MakeBBlockDistributionEncode() - { - constexpr index_t KPerThread = Traits::KPerThread; - constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; - constexpr index_t KPerInnerLoop = - ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); - constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; - - using KIterSeq = std::conditional_t, - sequence>; - - constexpr auto b_block_outer_dstr_encoding = - tile_distribution_encoding, - tuple, KIterSeq>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); - - return b_block_dstr_encode; - } + using Base::MakeABlockDistributionEncode; + using Base::MakeBBlockDistributionEncode; + using Base::MakeCBlockTile; private: - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - constexpr index_t UnaryOpSize = 8; - const element_wise::PassThroughPack8 elementwise_op{}; - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } - - template + template struct BlockGemmImpl { }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { static constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution(MakeABlockDistributionEncode())){}; @@ -248,7 +97,7 @@ struct BlockUniversalGemmAsBsCr load_tile(b_warp_tile_, b_block_window); } // hot loop: - static_for<0, GemmTraits::KIterPerWarp, 1>{}([&](auto kIter) { + static_for<0, GemmTraits_::KIterPerWarp, 1>{}([&](auto kIter) { static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { // read A warp tensor from A block tensor AWarpTensor a_warp_tensor; @@ -286,8 +135,8 @@ struct BlockUniversalGemmAsBsCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { static constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution(MakeABlockDistributionEncode())){}; @@ -375,11 +224,11 @@ struct BlockUniversalGemmAsBsCr } }; - template - struct BlockGemmImpl + template + struct BlockGemmImpl { - static constexpr index_t KPerThread = GemmTraits::KPerThread; - static constexpr index_t NumMacClusters = GemmTraits::InterWaveSchedulingMacClusters; + static constexpr index_t KPerThread = GemmTraits_::KPerThread; + static constexpr index_t NumMacClusters = GemmTraits_::InterWaveSchedulingMacClusters; static constexpr index_t KPerInnerLoop = ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); static constexpr index_t KRepeat = KPerThread / KPerInnerLoop; @@ -448,6 +297,8 @@ struct BlockUniversalGemmAsBsCr "The CDataType as defined in traits should be the same as correspoinding " "C block tensor data type!"); + printf("AS_BS_CR\n"); + // hot loop: static_for<0, KRepeat, 1>{}([&](auto kIter) { LocalPrefetchA(a_block_window); @@ -534,24 +385,6 @@ struct BlockUniversalGemmAsBsCr }; public: - CK_TILE_DEVICE static constexpr auto MakeCBlockTile() - { - constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< - sequence<>, - tuple, sequence>, - tuple>, - tuple>, - sequence<1, 2>, - sequence<0, 0>>{}; - - constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( - c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{}); - constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode); - auto c_block_tensor = make_static_distributed_tensor(c_block_dstr); - - return c_block_tensor; - } - template CK_TILE_DEVICE void LocalPrefetch(const ASmemBlockWindow& a_block_window, const BSmemBlockWindow& b_block_window) @@ -586,13 +419,13 @@ struct BlockUniversalGemmAsBsCr CK_TILE_DEVICE auto operator()(const ASmemBlockWindow& a_block_window, const BSmemBlockWindow& b_block_window) { - auto c_block_tensor = MakeCBlockTile(); + auto c_block_tensor = Base::MakeCBlockTile(); block_gemm_impl_(c_block_tensor, a_block_window, b_block_window); return c_block_tensor; } private: - BlockGemmImpl block_gemm_impl_{}; + BlockGemmImpl block_gemm_impl_{}; }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_base.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_base.hpp new file mode 100644 index 0000000000..ae4447275b --- /dev/null +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_base.hpp @@ -0,0 +1,217 @@ +// 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/block/block_gemm_areg_bsmem_creg_v1_default_policy.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" +#include "ck_tile/ops/elementwise.hpp" + +namespace ck_tile { + +template +struct BlockUniversalGemmBase +{ +protected: + // TODO: This should be in Policy - UniversalGemmPolicyBase ? + template + struct GemmTraits_ + { + using Problem = remove_cvref_t; + using Policy = remove_cvref_t; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using ComputeDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + + static constexpr index_t kBlockSize = Problem::kBlockSize; + static constexpr auto Scheduler = Problem::Scheduler; + + static constexpr index_t MPerBlock = BlockGemmShape::kM; + static constexpr index_t NPerBlock = BlockGemmShape::kN; + static constexpr index_t KPerBlock = BlockGemmShape::kK; + + static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); + + using WarpGemm = remove_cvref_t())>; + + static constexpr index_t MWarp = config.template at<1>(); + static constexpr index_t NWarp = config.template at<2>(); + + using I0 = number<0>; + using I1 = number<1>; + + static_assert(MWarp == BlockGemmShape::BlockWarps::at(I0{}), + "Error! WarpGemm's MWarp is not consisten with BlockGemmShape!"); + static_assert(NWarp == BlockGemmShape::BlockWarps::at(I1{}), + "Error! WarpGemm's NWarp is not consisten with BlockGemmShape!"); + static_assert(WarpGemm::kM == BlockGemmShape::WarpTile::at(I0{}), + "Error! WarpGemm's M is not consisten with BlockGemmShape!"); + static_assert(WarpGemm::kN == BlockGemmShape::WarpTile::at(I1{}), + "Error! WarpGemm's N is not consisten with BlockGemmShape!"); + + static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WarpGemm::kM); + static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); + static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; + + static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, + "Error! Warps should cover all Block tile!"); + static_assert(NIterPerWarp * NWarp * WarpGemm::kN == NPerBlock, + "Error! Warps should cover all Block tile!"); + + static constexpr index_t MPerBlockPerIter = MWarp * WarpGemm::kM; + static constexpr index_t NPerBlockPerIter = NWarp * WarpGemm::kN; + static constexpr index_t KPerBlockPerIter = WarpGemm::kK; + + // Controls how many MAC clusters (MFMA blocks) we have per wave + // Ie if + // InterWaveSchedulingMacClusters = 1; + // KPerBlock == 32 + // WarpGemm::kK = 8 + // Then we would group all 4 WarpGemms into single MAC cluster. + // But if we would set InterWaveSchedulingMacClusters = 2, then we would + // split those 4 warp gemms into two groups. + static constexpr index_t InterWaveSchedulingMacClusters = 1; + + // should be at least equal to: WarpGemm::Impl::kABKPerLane + static constexpr index_t KPack = WarpGemm::kKPerThread; + static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread; + }; + +public: + using Traits = GemmTraits_; + + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using ComputeDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using WarpGemm = remove_cvref_t; + + static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; + static constexpr index_t MIterPerWarp = Traits::MIterPerWarp; + static constexpr index_t NIterPerWarp = Traits::NIterPerWarp; + + static constexpr index_t MWarp = Traits::MWarp; + static constexpr index_t NWarp = Traits::NWarp; + + static constexpr auto Scheduler = Traits::Scheduler; + + using AWarpDstr = typename WarpGemm::AWarpDstr; + using BWarpDstr = typename WarpGemm::BWarpDstr; + using CWarpDstr = typename WarpGemm::CWarpDstr; + + using AWarpTensor = typename WarpGemm::AWarpTensor; + using BWarpTensor = typename WarpGemm::BWarpTensor; + using CWarpTensor = typename WarpGemm::CWarpTensor; + + static constexpr auto a_warp_y_lengths = + to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + static constexpr auto b_warp_y_lengths = + to_sequence(BWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + static constexpr auto c_warp_y_lengths = + to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + + static constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; + static constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t{}; + static constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + + static constexpr index_t APackedSize = + ck_tile::numeric_traits>::PackedSize; + static constexpr index_t BPackedSize = + ck_tile::numeric_traits>::PackedSize; + + using I0 = number<0>; + using I1 = number<1>; + + CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode() + { + constexpr index_t KPerThread = Traits::KPerThread; + constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; + constexpr index_t KPerInnerLoop = + ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); + constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; + + using KIterSeq = std::conditional_t, + sequence>; + + constexpr auto a_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, KIterSeq>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); + + return a_block_dstr_encode; + } + + CK_TILE_DEVICE static constexpr auto MakeBBlockDistributionEncode() + { + constexpr index_t KPerThread = Traits::KPerThread; + constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; + constexpr index_t KPerInnerLoop = + ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); + constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; + + using KIterSeq = std::conditional_t, + sequence>; + + constexpr auto b_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, KIterSeq>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{}); + + return b_block_dstr_encode; + } + + CK_TILE_DEVICE static constexpr auto MakeCBlockTile() + { + constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< + sequence<>, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + + constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{}); + constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode); + auto c_block_tensor = make_static_distributed_tensor(c_block_dstr); + + return c_block_tensor; + } + +protected: + template + CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, + const WarpWindow& warp_window) + { + constexpr index_t UnaryOpSize = 8; + const element_wise::PassThroughPack8 elementwise_op{}; + constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; + const auto in_dstr_tensors = load_tile(warp_window); + + static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); + + using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); + static_for<0, thread_buffer_size, 1>{}([&](auto i) { + elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), + in_dstr_tensors.get_thread_buffer().template get_as()[i]); + }); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp index d495c0d950..dccecfc113 100644 --- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp @@ -22,12 +22,14 @@ struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs ck_tile::index_t stride_A_, ck_tile::index_t stride_B_, ck_tile::index_t stride_C_, + bool skip_a_lds_, + bool skip_b_lds_, ck_tile::index_t batch_stride_A_, ck_tile::index_t batch_stride_B_, ck_tile::index_t batch_stride_C_, ck_tile::index_t batch_count_) : GemmHostArgs( - a_ptr_, b_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_), + a_ptr_, b_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_, skip_a_lds_, skip_b_lds_), batch_stride_A(batch_stride_A_), batch_stride_B(batch_stride_B_), batch_stride_C(batch_stride_C_), diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index 9c25104cd7..2b2ec76295 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -53,12 +53,16 @@ struct GemmHostArgs : public GemmProblem index_t K_, index_t stride_A_, index_t stride_B_, - index_t stride_C_) + index_t stride_C_, + bool skip_a_lds_, + bool skip_b_lds_) : GemmProblem(M_, N_, K_, stride_A_, stride_B_, stride_C_), a_ptr(a_ptr_), b_ptr(b_ptr_), c_ptr(c_ptr_), - k_batch(k_batch_) + k_batch(k_batch_), + skip_a_lds(skip_a_lds_), + skip_b_lds(skip_b_lds_) { } @@ -66,6 +70,8 @@ struct GemmHostArgs : public GemmProblem const void* b_ptr; void* c_ptr; index_t k_batch; + bool skip_a_lds; + bool skip_b_lds; }; /// @brief The GEMM kernel device arguments. @@ -93,6 +99,10 @@ struct GemmKernelArgs /// (in memory) of C tensor. index_t stride_C; index_t k_batch; + /// @brief Flag to skip loading A tensor tile into LDS. + bool skip_a_lds; + /// @brief Flag to skip loading B tensor tile into LDS. + bool skip_b_lds; }; /// @brief The GEMM kernel template. @@ -176,7 +186,10 @@ struct GemmKernel hostArgs.stride_A, hostArgs.stride_B, hostArgs.stride_C, - hostArgs.k_batch}; + hostArgs.k_batch, + hostArgs.skip_a_lds, + hostArgs.skip_b_lds + }; } CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()