diff --git a/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp b/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp index fdf6c95bff..5a342c7dfa 100644 --- a/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp @@ -5,6 +5,10 @@ #include "ck_tile/core.hpp" #include "ck_tile/core/tensor/tile_distribution.hpp" +#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_custom_policy.hpp" +#include "ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp" +#include "ck_tile/ops/common/tensor_layout.hpp" #include "block_gemm_asmem_bsmem_creg.hpp" #include "config.h" @@ -106,7 +110,6 @@ struct BlockGemmPipelineAGmemBGmemCRegDefaultPolicy make_tuple(sequence<1, 0>{}, sequence<2, 3>{}), make_tuple(sequence<0>{}, sequence<1>{})); #endif - return a_lds_block_desc; } @@ -260,9 +263,330 @@ struct BlockGemmPipelineAGmemBGmemCRegDefaultPolicy template CK_TILE_HOST_DEVICE static constexpr auto GetBlockGemm() { - return BlockGemmASmemBSmemCReg{}; } }; +#if 0 +// UniversalGemm Policy +struct UniversalGemmPipelineAgBgCrPolicy +{ + static constexpr auto I0 = number<0>{}; + static constexpr auto I1 = number<1>{}; + static constexpr auto I2 = number<2>{}; + + // static constexpr auto ATileAccessPattern = tile_distribution_pattern::thread_raked; + // static constexpr auto BTileAccessPattern = tile_distribution_pattern::thread_raked; + + template + CK_TILE_HOST_DEVICE static constexpr auto GetGlobalVectorLoadSize() + { + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t elements_per_thread = MNPerBlock * KPerBlock / BlockSize; // 32 = 128 * 64 / 256 + constexpr index_t PackedSize = + ck_tile::numeric_traits>::PackedSize; + + // Assume DataType is even! + if constexpr(XPerTile % (PackedSize * 32 / sizeof(DataType)) == 0 && + elements_per_thread % (PackedSize * 32 / sizeof(DataType)) == 0 && + PackedSize == 2) + { + return (PackedSize * 32 / sizeof(DataType)); + } + else if constexpr(XPerTile % (PackedSize * 16 / sizeof(DataType)) == 0 && + elements_per_thread % (PackedSize * 16 / sizeof(DataType)) == 0) + { + return (PackedSize * 16 / sizeof(DataType)); + } + else if constexpr(XPerTile % (PackedSize * 8 / sizeof(DataType)) == 0 && + elements_per_thread % (PackedSize * 8 / sizeof(DataType)) == 0) + { + return (PackedSize * 8 / sizeof(DataType)); + } + else if constexpr(sizeof(DataType) >= PackedSize * 4 && + XPerTile % (PackedSize * 4 / sizeof(DataType)) == 0 && + elements_per_thread % (PackedSize * 4 / sizeof(DataType)) == 0) + { + return (PackedSize * 4 / sizeof(DataType)); + } + else if constexpr(sizeof(DataType) >= PackedSize * 2 && + XPerTile % (PackedSize * 2 / sizeof(DataType)) == 0 && + elements_per_thread % (PackedSize * 2 / sizeof(DataType)) == 0) + { + return (PackedSize * 2 / sizeof(DataType)); + } + else + { + return PackedSize; + } + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeA() + { + using ADataType = remove_cvref_t; + using ALayout = remove_cvref_t; + static_assert(std::is_same_v); + + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + + return GetGlobalVectorLoadSize(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeB() + { + using BDataType = remove_cvref_t; + using BLayout = remove_cvref_t; + static_assert(std::is_same_v); + + constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + + return GetGlobalVectorLoadSize(); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeC() + { + using BlockGemm = remove_cvref_t())>; + using WG = typename BlockGemm::WarpGemm; + + // constexpr bool TransposeC = Problem::TransposeC; + // using CLayout = typename Problem::CLayout; + using CWarpDstr = typename WG::CWarpDstr; + + // In this case each thread has multiple consecutive elements in + // N dimension, however consecutive threads' elements have stride. + constexpr index_t NDimY = CWarpDstr::NDimY; + constexpr auto c_warp_y_lengths = + CWarpDstr{}.get_ys_to_d_descriptor().get_lengths(); + static_assert(WG::WarpGemmAttribute::Impl::kCM1PerLane == + c_warp_y_lengths.get(number{})); + return c_warp_y_lengths.get(number{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto IsTransposeC() + { + return Problem::TransposeC; + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution() + { + using ADataType = remove_cvref_t; + + constexpr index_t kBlockSize = Problem::kBlockSize; + + constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; + + constexpr index_t K1 = 16 / sizeof(ADataType); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t M2 = get_warp_size() / K0; + // coalesce reading for each blocks + constexpr index_t M1 = kBlockSize / get_warp_size(); + constexpr index_t M0 = kMPerBlock / (M2 * M1); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{}); + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeBDramTileDistribution() + { + using BDataType = remove_cvref_t; + + constexpr index_t kBlockSize = Problem::kBlockSize; + + constexpr index_t kNPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; + + constexpr index_t K1 = 16 / sizeof(BDataType); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t N2 = get_warp_size() / K0; + // coalesce reading for each blocks + constexpr index_t N1 = kBlockSize / get_warp_size(); + constexpr index_t N0 = kNPerBlock / (N2 * N1); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{}); + } + + 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; + } + + // 3d + padding + template + CK_TILE_HOST_DEVICE static constexpr auto MakeALdsBlockDescriptor() + { + constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t kKPack = 8; + using ADataType = remove_cvref_t; + + constexpr auto DataTypeSize = sizeof(ADataType); + constexpr auto MLdsLayer = + (32 * 4 / kKPerBlock / DataTypeSize) < 1 ? 1 : (32 * 4 / kKPerBlock / DataTypeSize); + + 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>{}); + + constexpr auto a_lds_block_desc_permuted = transform_tensor_descriptor( + a_lds_block_desc_0, + make_tuple(make_xor_transform(make_tuple(number{}, + number{})), + make_pass_through_transform(number{})), + make_tuple(sequence<1, 0>{}, sequence<2>{}), + make_tuple(sequence<1, 0>{}, sequence<2>{})); + + constexpr auto a_lds_block_desc_xk0_mnldslayer_mn_xk1 = transform_tensor_descriptor( + a_lds_block_desc_permuted, + make_tuple(make_unmerge_transform( + make_tuple(number{}, number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}), + make_tuple(sequence<0, 2>{}, sequence<1>{}, sequence<3>{})); + + constexpr auto a_lds_block_desc = transform_tensor_descriptor( + a_lds_block_desc_xk0_mnldslayer_mn_xk1, + make_tuple(make_merge_transform( + make_tuple(number{}, number{})), + make_merge_transform( + make_tuple(number{}, number{}))), + make_tuple(sequence<1, 0>{}, sequence<2, 3>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + return a_lds_block_desc; + } + + // 3d + padding + template + CK_TILE_HOST_DEVICE static constexpr auto MakeBLdsBlockDescriptor() + { + constexpr index_t kNPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t kKPack = 8; + using BDataType = remove_cvref_t; + + constexpr auto DataTypeSize = sizeof(BDataType); + constexpr auto NLdsLayer = + (32 * 4 / kKPerBlock / DataTypeSize) < 1 ? 1 : (32 * 4 / kKPerBlock / DataTypeSize); + + 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>{}); + + constexpr auto b_lds_block_desc_permuted = transform_tensor_descriptor( + b_lds_block_desc_0, + make_tuple(make_xor_transform(make_tuple(number{}, + number{})), + make_pass_through_transform(number{})), + make_tuple(sequence<1, 0>{}, sequence<2>{}), + make_tuple(sequence<1, 0>{}, sequence<2>{})); + + constexpr auto b_lds_block_desc_xk0_mnldslayer_mn_xk1 = transform_tensor_descriptor( + b_lds_block_desc_permuted, + make_tuple(make_unmerge_transform( + make_tuple(number{}, number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}), + make_tuple(sequence<0, 2>{}, sequence<1>{}, sequence<3>{})); + + constexpr auto b_lds_block_desc = transform_tensor_descriptor( + b_lds_block_desc_xk0_mnldslayer_mn_xk1, + make_tuple(make_merge_transform( + make_tuple(number{}, number{})), + make_merge_transform( + make_tuple(number{}, number{}))), + make_tuple(sequence<1, 0>{}, sequence<2, 3>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + return b_lds_block_desc; + } + + 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{}; + } +}; +#endif + } // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/02_gemm/config.h b/example/ck_tile/99_toy_example/02_gemm/config.h index 92ad729b4b..fc70599520 100644 --- a/example/ck_tile/99_toy_example/02_gemm/config.h +++ b/example/ck_tile/99_toy_example/02_gemm/config.h @@ -34,6 +34,6 @@ #define ENABLE_INSTRUCTION_SCH #define ENABLE_CACHE_AWARE_WG_SCH #else - #define NAIVE_IMPLEMENTATION + #define NAIVE_IMPLEMENTATION #endif diff --git a/example/ck_tile/99_toy_example/02_gemm/gemm.cpp b/example/ck_tile/99_toy_example/02_gemm/gemm.cpp index 51dd5c59f9..e7f8b3b9f3 100644 --- a/example/ck_tile/99_toy_example/02_gemm/gemm.cpp +++ b/example/ck_tile/99_toy_example/02_gemm/gemm.cpp @@ -49,38 +49,38 @@ int main(int argc, char* argv[]) } #if defined(KERNEL_A) - printf("*** KernelA test *** \n"); + printf("*** Kernel A test *** \n"); printf(" --> Using mfma_32x32x(8x2)\n"); #elif defined(KERNEL_B) - printf("*** KernelB test *** \n"); + printf("*** Kernel B test *** \n"); printf(" --> Using mfma_16x16x16\n"); #elif defined(KERNEL_C) - printf("*** KernelC test *** \n"); + printf("*** Kernel C test *** \n"); printf(" --> Using mfma_16x16x(16x2)\n"); #elif defined(KERNEL_D) - printf("*** KernelD test *** \n"); + printf("*** Kernel D test *** \n"); printf(" --> Using mfma_16x16x(16x2)\n"); printf(" --> XOR-based bank-conflict-free\n"); #elif defined(KERNEL_E) - printf("*** KernelE test ***\n"); + printf("*** Kernel E test ***\n"); printf(" --> Using mfma_16x16x(16x2)\n"); printf(" --> XOR-based bank-conflict-free\n"); printf(" --> Adjust block tile shape\n"); #elif defined(KERNEL_F) - printf("*** KernelF test ***\n"); + printf("*** Kernel F test ***\n"); printf(" --> Using mfma_16x16x(16x2)\n"); printf(" --> XOR-based bank-conflict-free\n"); printf(" --> Adjust block tile shape\n"); printf(" --> Enable prefetch\n"); #elif defined(KERNEL_G) - printf("*** KernelG test ***\n"); + printf("*** Kernel G test ***\n"); printf(" --> Using mfma_16x16x(16x2)\n"); printf(" --> XOR-based bank-conflict-free\n"); printf(" --> Adjust block tile shape\n"); printf(" --> Enable prefetch\n"); printf(" --> Enable instruction schedule\n"); #elif defined(KERNEL_H) - printf("*** KernelH test ***\n"); + printf("*** Kernel H test ***\n"); printf(" --> Using mfma_16x16x(16x2)\n"); printf(" --> XOR-based bank-conflict-free\n"); printf(" --> Adjust block tile shape\n"); diff --git a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_ag_bg_cr_comp_v3.hpp b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_ag_bg_cr_comp_v3.hpp index 12b2ec550d..fc9fc232fe 100644 --- a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_ag_bg_cr_comp_v3.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_ag_bg_cr_comp_v3.hpp @@ -16,52 +16,9 @@ namespace ck_tile { -// A Tile Window: global memory -// B Tile Window: global memory -// C Distributed tensor: register -template -struct BaseGemmPipelineAgBgCrCompV3 -{ - static constexpr index_t PrefetchStages = 2; - static constexpr index_t PrefillStages = 1; - static constexpr index_t GlobalBufferNum = 1; - - CK_TILE_HOST_DEVICE static constexpr auto TransposeC() { return Problem::TransposeC; } - - CK_TILE_HOST static constexpr bool BlockHasHotloop(index_t num_loop) - { - return num_loop > PrefetchStages; - } - - CK_TILE_HOST static constexpr TailNumber GetBlockLoopTailNum(index_t num_loop) - { - if(BlockHasHotloop(num_loop)) - { - return TailNumber::Full; - } - else - { - if(num_loop == 1) - { - return TailNumber::Odd; - } - else - { - return TailNumber::Even; - } - } - } -}; - -// Compute optimized pipeline -// GlobalPrefetchStages: 2 -// LocalPreFillStages: 1 -// LocalPreFetchStages: 1 -// LocalSharedMemoryBuffer: 1 template -struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 +struct GemmPipelineAgBgCrCompV3 { - using Base = BaseGemmPipelineAgBgCrCompV3; using PipelineImplBase = GemmPipelineAgBgCrImplBase; using ADataType = remove_cvref_t; @@ -105,8 +62,6 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 static constexpr auto TailNum = Problem::TailNum; static constexpr auto Scheduler = Problem::Scheduler; - using Base::PrefetchStages; - [[nodiscard]] CK_TILE_HOST static const std::string GetName() { // clang-format off @@ -121,56 +76,6 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 return Policy::template GetSmemSize(); } - CK_TILE_HOST static std::string Print() - { - constexpr index_t MPerXDL = BlockGemm::WarpGemm::kM; - constexpr index_t NPerXDL = BlockGemm::WarpGemm::kN; - constexpr index_t KPerXDL = BlockGemm::WarpGemm::WarpGemmAttribute::Impl::kK; - - constexpr index_t WaveSize = 64; - constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(I0{}); - constexpr index_t WaveNumN = BlockGemmShape::BlockWarps::at(I1{}); - - // Below should be equal to AK1|BK1 - constexpr index_t A_LDS_Read_Width = GetSmemPackA(); - constexpr index_t B_LDS_Read_Width = GetSmemPackB(); - - constexpr index_t A_LDS_Write_Width = GetSmemPackA(); - constexpr index_t B_LDS_Write_Width = GetSmemPackB(); - - 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 * A_LDS_Write_Width); - constexpr index_t B_LDS_Write_Inst_Num = - NPerBlock * KPerBlock / (BlockSize * B_LDS_Write_Width); - - constexpr index_t A_LDS_Read_Inst_Num = - WaveNumN * MPerBlock * KPerBlock / (BlockSize * A_LDS_Read_Width); - constexpr index_t B_LDS_Read_Inst_Num = - WaveNumM * NPerBlock * KPerBlock / (BlockSize * B_LDS_Read_Width); - - constexpr index_t C_MFMA_Inst_Num = MPerBlock * NPerBlock * KPerBlock / - (BlockSize / WaveSize) / (MPerXDL * NPerXDL * KPerXDL); - - auto str = std::stringstream{}; - - str << "A/B vector size: " << GetVectorSizeA() << ", " << GetVectorSizeB() << "\n" - << "A/B LDS read/write width: " << A_LDS_Read_Width << ", " << B_LDS_Read_Width << "\n" - << "A/B buffer load inst: " << A_Buffer_Load_Inst_Num << ", " << B_Buffer_Load_Inst_Num - << "\n" - << "A/B LDS write inst: " << A_LDS_Write_Inst_Num << ", " << B_LDS_Write_Inst_Num - << "\n" - << "A/B LDS read inst: " << A_LDS_Read_Inst_Num << ", " << B_LDS_Read_Inst_Num << "\n" - << "C MFMA inst: " << C_MFMA_Inst_Num << "\n" - << "KPack: " << BlockGemm::Traits::KPack << "\n" - << "PrefetchStages: " << PrefetchStages << "\n"; - return str.str(); - } - template struct PipelineImpl : public PipelineImplBase { diff --git a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_problem.hpp b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_problem.hpp index f833ccc849..34e8fa18c8 100644 --- a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_problem.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_pipeline_problem.hpp @@ -43,16 +43,6 @@ struct GemmPipelineProblemBase static constexpr auto Scheduler = GemmPipelineScheduler::Default; static constexpr index_t VectorLoadSize = Traits::_VectorSize; - [[nodiscard]] CK_TILE_HOST static const std::string GetName() - { - // clang-format off - return concat('_', "gemm_problem", - concat('x', VectorLoadSize, kBlockSize), - concat('x', kPadM, kPadN, kPadK), - Scheduler); - // clang-format on - } - CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentA() { constexpr index_t PackedSize = diff --git a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_universal_pipeline_ag_bg_cr_policy.hpp b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_universal_pipeline_ag_bg_cr_policy.hpp index 313821d11c..b147e8f178 100644 --- a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_universal_pipeline_ag_bg_cr_policy.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/gemm_universal_pipeline_ag_bg_cr_policy.hpp @@ -11,25 +11,13 @@ namespace ck_tile { -template -struct UniversalGemmBasePolicy +// UniversalGemm Policy +struct UniversalGemmPipelineAgBgCrPolicy { static constexpr auto I0 = number<0>{}; static constexpr auto I1 = number<1>{}; static constexpr auto I2 = number<2>{}; - static constexpr auto ATileAccessPattern = tile_distribution_pattern::thread_raked; - static constexpr auto BTileAccessPattern = tile_distribution_pattern::thread_raked; - - /** - * @brief Get the maximum global memory vector load size. - * - * @tparam Problem The UniversalGemmPipelineProblem object. - * @tparam DataType The tensor data type we're considering. - * @tparam MNPerBlock The MPerBlock or NPerBlock value depending on tensor (A/B). - * @tparam XPerTile The contiguous Tile dimension size. - * @return Maximum DRAM vector load size. - */ template CK_TILE_HOST_DEVICE static constexpr auto GetGlobalVectorLoadSize() { @@ -77,105 +65,44 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeA() { - using ALayout = remove_cvref_t; - using ADataType = remove_cvref_t; + using ADataType = remove_cvref_t; + using ALayout = remove_cvref_t; + static_assert(std::is_same_v); + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - if constexpr(std::is_same_v) - { - return GetGlobalVectorLoadSize(); - } - else - { - return GetGlobalVectorLoadSize(); - } + return GetGlobalVectorLoadSize(); } template CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeB() { - using BLayout = remove_cvref_t; - using BDataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using BLayout = remove_cvref_t; + static_assert(std::is_same_v); + constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - if constexpr(std::is_same_v) - { - return GetGlobalVectorLoadSize(); - } - else - { - return GetGlobalVectorLoadSize(); - } + return GetGlobalVectorLoadSize(); } - /** - * @brief Get the vector store size for C tensor. - * - * @tparam Problem - Gemm pipeline problem class. - * - * @note The vector store size for output C tensor would depend on multiple factors - * like its data layout and warp gemm C transposition. In general it would - * be the number of consecutive elements in contiguous C dimension hold by - * single thread. - * - * @return The vector store size for C tensor. - */ template CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeC() { - using BlockGemm = remove_cvref_t())>; + using BlockGemm = remove_cvref_t())>; using WG = typename BlockGemm::WarpGemm; + using CWarpDstr = typename WG::CWarpDstr; - constexpr bool TransposeC = Problem::TransposeC; - using CLayout = typename Problem::CLayout; - using CWarpDstr = typename WG::CWarpDstr; - - // N is contiguous dimension - if constexpr(std::is_same_v) - { - if constexpr(TransposeC) - { - // In this case each thread has multiple consecutive elements in - // N dimension, however consecutive threads' elements have stride. - constexpr index_t NDimY = CWarpDstr::NDimY; - constexpr auto c_warp_y_lengths = - CWarpDstr{}.get_ys_to_d_descriptor().get_lengths(); - static_assert(WG::WarpGemmAttribute::Impl::kCM1PerLane == - c_warp_y_lengths.get(number{})); - return c_warp_y_lengths.get(number{}); - } - else - { - // In this case each thread has just a single item in Ndim - return WG::WarpGemmAttribute::Impl::kCNLane / WG::kN; - } - } - // M is contiguous dimension - else if constexpr(std::is_same_v) - { - if constexpr(TransposeC) - { - // In this case each thread has just a single item in Mdim - return WG::WarpGemmAttribute::Impl::kCNLane / WG::kN; - } - else - { - // In this case each thread has multiple consecutive elements in - // M dimension, however consecutive threads' elements have stride. - constexpr index_t NDimY = CWarpDstr::NDimY; - constexpr auto c_warp_y_lengths = - CWarpDstr{}.get_ys_to_d_descriptor().get_lengths(); - static_assert(WG::WarpGemmAttribute::Impl::kCM1PerLane == - c_warp_y_lengths.get(number{})); - return c_warp_y_lengths.get(number{}); - } - } - else - { - static_assert(false, "Unsupported CLayout!"); - } + // In this case each thread has multiple consecutive elements in + // N dimension, however consecutive threads' elements have stride. + constexpr index_t NDimY = CWarpDstr::NDimY; + constexpr auto c_warp_y_lengths = + CWarpDstr{}.get_ys_to_d_descriptor().get_lengths(); + static_assert(WG::WarpGemmAttribute::Impl::kCM1PerLane == + c_warp_y_lengths.get(number{})); + return c_warp_y_lengths.get(number{}); } template @@ -187,107 +114,59 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution() { - using ALayout = remove_cvref_t; + using ADataType = remove_cvref_t; - 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 kBlockSize = Problem::kBlockSize; - // Tile: MPerBlock X KPerBlock - if constexpr(std::is_same_v) - { - using TileEncodingPattern = TileDistributionEncodingPattern2D; - return TileEncodingPattern::Make2DStaticTileDistribution(); - } - // Tile: KPerBlock X MPerBlock - else - { - using TileEncodingPattern = TileDistributionEncodingPattern2D; - return TileEncodingPattern::Make2DStaticTileDistribution(); - } + constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; + + constexpr index_t K1 = 16 / sizeof(ADataType); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t M2 = get_warp_size() / K0; + // coalesce reading for each blocks + constexpr index_t M1 = kBlockSize / get_warp_size(); + constexpr index_t M0 = kMPerBlock / (M2 * M1); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{}); } template CK_TILE_HOST_DEVICE static constexpr auto MakeBDramTileDistribution() { - using BLayout = remove_cvref_t; + using BDataType = remove_cvref_t; - 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 kBlockSize = Problem::kBlockSize; - // Tile: KPerBlock X NPerBlock - if constexpr(std::is_same_v) - { - using TileEncodingPattern = TileDistributionEncodingPattern2D; - return TileEncodingPattern::Make2DStaticTileDistribution(); - } - // Tile: NPerBlock X KPerBlock - else - { - using TileEncodingPattern = TileDistributionEncodingPattern2D; - return TileEncodingPattern::Make2DStaticTileDistribution(); - } - } + constexpr index_t kNPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; - 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 K1 = 16 / sizeof(BDataType); + constexpr index_t K0 = kKPerBlock / K1; + constexpr index_t N2 = get_warp_size() / K0; + // coalesce reading for each blocks + constexpr index_t N1 = kBlockSize / get_warp_size(); + constexpr index_t N0 = kNPerBlock / (N2 * N1); - 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(); - - using TileEncodingPattern = TileDistributionEncodingPattern2D; - return TileEncodingPattern::MakeShuffled2DStaticTileDistribution(); + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{}); } template CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackA() { - using BlockGemm = remove_cvref_t())>; + using BlockGemm = remove_cvref_t())>; constexpr index_t KPack = BlockGemm::Traits::KPack; return KPack; } @@ -295,7 +174,7 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackB() { - using BlockGemm = remove_cvref_t())>; + using BlockGemm = remove_cvref_t())>; constexpr index_t KPack = BlockGemm::Traits::KPack; return KPack; } @@ -303,7 +182,7 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeA() { - constexpr auto a_lds_desc = Derived::template MakeALdsBlockDescriptor(); + 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; @@ -312,7 +191,7 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeB() { - constexpr auto b_lds_desc = Derived::template MakeBLdsBlockDescriptor(); + 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; @@ -326,249 +205,103 @@ struct UniversalGemmBasePolicy return smem_size_a + smem_size_b; } -}; -// UniversalGemm Policy -struct UniversalGemmPipelineAgBgCrPolicy - : public UniversalGemmBasePolicy -{ + // 3d + padding template CK_TILE_HOST_DEVICE static constexpr auto MakeALdsBlockDescriptor() { + constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t kKPack = 8; using ADataType = remove_cvref_t; - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPack = GetSmemPackA(); - constexpr auto DataTypeSize = sizeof(ADataType); constexpr auto MLdsLayer = - (32 * 4 / KPerBlock / DataTypeSize) < 1 ? 1 : (32 * 4 / KPerBlock / DataTypeSize); + (32 * 4 / kKPerBlock / DataTypeSize) < 1 ? 1 : (32 * 4 / kKPerBlock / DataTypeSize); constexpr auto a_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, - number{}, - number{}), - make_tuple(number{}, number{}, number<1>{}), - number{}, + make_tuple(number{}, + number{}, + number{}), + make_tuple(number{}, number{}, number<1>{}), + number{}, number<1>{}); constexpr auto a_lds_block_desc_permuted = transform_tensor_descriptor( a_lds_block_desc_0, - make_tuple(make_xor_transform(make_tuple(number{}, - number{})), - make_pass_through_transform(number{})), + make_tuple(make_xor_transform(make_tuple(number{}, + number{})), + make_pass_through_transform(number{})), make_tuple(sequence<1, 0>{}, sequence<2>{}), make_tuple(sequence<1, 0>{}, sequence<2>{})); constexpr auto a_lds_block_desc_xk0_mnldslayer_mn_xk1 = transform_tensor_descriptor( a_lds_block_desc_permuted, make_tuple(make_unmerge_transform( - make_tuple(number{}, number{})), - make_pass_through_transform(number{}), - make_pass_through_transform(number{})), + make_tuple(number{}, number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(number{})), make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}), make_tuple(sequence<0, 2>{}, sequence<1>{}, sequence<3>{})); constexpr auto a_lds_block_desc = transform_tensor_descriptor( a_lds_block_desc_xk0_mnldslayer_mn_xk1, - make_tuple(make_merge_transform_v3_division_mod( - make_tuple(number{}, number{})), - make_merge_transform_v3_division_mod( - make_tuple(number{}, number{}))), - make_tuple(sequence<1, 2>{}, sequence<0, 3>{}), + make_tuple(make_merge_transform( + make_tuple(number{}, number{})), + make_merge_transform( + make_tuple(number{}, number{}))), + make_tuple(sequence<1, 0>{}, sequence<2, 3>{}), make_tuple(sequence<0>{}, sequence<1>{})); - return a_lds_block_desc; } - /** - * @brief Create LDS block descriptor for B tensor. - * - * @tparam Problem Gemm pipeline problem. - * @return B tensor LDS block descriptor. - */ + // 3d + padding template CK_TILE_HOST_DEVICE static constexpr auto MakeBLdsBlockDescriptor() { - // using BLayout = remove_cvref_t; + constexpr index_t kNPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t kKPack = 8; using BDataType = remove_cvref_t; - constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr auto DataTypeSize = sizeof(BDataType); + constexpr auto NLdsLayer = + (32 * 4 / kKPerBlock / DataTypeSize) < 1 ? 1 : (32 * 4 / kKPerBlock / DataTypeSize); -#if 1 - // if constexpr(std::is_same_v) - { - constexpr index_t KPack = GetSmemPackB(); - constexpr auto BK0 = number{}; - constexpr auto DataTypeSize = sizeof(BDataType); - constexpr auto NLdsLayer = - (32 * 4 / KPerBlock / DataTypeSize) < 1 ? 1 : (32 * 4 / KPerBlock / DataTypeSize); + 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>{}); - constexpr auto b_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple( - BK0 * number{}, number{}, number{}), - make_tuple(number{}, number{}, number<1>{}), - number{}, - number<1>{}); + constexpr auto b_lds_block_desc_permuted = transform_tensor_descriptor( + b_lds_block_desc_0, + make_tuple(make_xor_transform(make_tuple(number{}, + number{})), + make_pass_through_transform(number{})), + make_tuple(sequence<1, 0>{}, sequence<2>{}), + make_tuple(sequence<1, 0>{}, sequence<2>{})); - constexpr auto b_lds_block_desc_permuted = transform_tensor_descriptor( - b_lds_block_desc_0, - make_tuple(make_xor_transform(make_tuple(number{}, - BK0 * number{})), - make_pass_through_transform(number{})), - make_tuple(sequence<1, 0>{}, sequence<2>{}), - make_tuple(sequence<1, 0>{}, sequence<2>{})); + constexpr auto b_lds_block_desc_xk0_mnldslayer_mn_xk1 = transform_tensor_descriptor( + b_lds_block_desc_permuted, + make_tuple(make_unmerge_transform( + make_tuple(number{}, number{})), + make_pass_through_transform(number{}), + make_pass_through_transform(number{})), + make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}), + make_tuple(sequence<0, 2>{}, sequence<1>{}, sequence<3>{})); - constexpr auto b_lds_block_desc_bk0_nldslayer_n_bk1 = transform_tensor_descriptor( - b_lds_block_desc_permuted, - make_tuple(make_unmerge_transform(make_tuple(BK0, number{})), - make_pass_through_transform(number{}), - make_pass_through_transform(number{})), - make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}), - make_tuple(sequence<0, 2>{}, sequence<1>{}, sequence<3>{})); - - constexpr auto b_lds_block_desc = transform_tensor_descriptor( - b_lds_block_desc_bk0_nldslayer_n_bk1, - make_tuple(make_merge_transform_v3_division_mod( - make_tuple(number{}, number{})), - make_merge_transform_v3_division_mod(make_tuple(BK0, number{}))), - make_tuple(sequence<1, 2>{}, sequence<0, 3>{}), - make_tuple(sequence<0>{}, sequence<1>{})); - return b_lds_block_desc; - } -#else - else // B is Row Major - { - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t VecLoadSize = GetVectorSizeB(); - using TileEncodingPattern = TileDistributionEncodingPattern2D; - - constexpr auto BK0 = number{}; - constexpr auto BK1 = number{}; - // constexpr auto N0 = BBlockTransferThreadClusterLengths_BK0_N_BK1{}.At(I1); - constexpr auto N0 = TileEncodingPattern::X0; - constexpr auto N1 = NPerBlock / N0; - - using WarpTile = typename Problem::BlockGemmShape::WarpTile; - constexpr auto NPerXdl = number{}; - - // constexpr auto KThreadWrite = - // BBlockTransferThreadClusterLengths_BK0_N_BK1{}.At(I0); - constexpr auto KThreadWrite = TileEncodingPattern::Y2; - constexpr auto K0PerThreadWrite = BK0 / KThreadWrite; - constexpr auto KThreadRead = 64 / NPerXdl; - constexpr auto K0PerThreadRead = BK0 / KThreadRead; - - constexpr auto kfold = - (BK1 * N0 * sizeof(BDataType) > 128) ? 1 : 128 / (BK1 * N0 * sizeof(BDataType)); - constexpr auto KThreadReadPerm = - (kfold * K0PerThreadWrite / K0PerThreadRead) > 1 - ? KThreadRead / (kfold * K0PerThreadWrite / K0PerThreadRead) - : KThreadRead; - - // 1<=npair<=n0 - constexpr auto npair = (BK1 * NPerXdl * sizeof(BDataType) > 128) - ? 1 - : ((128 / (BK1 * NPerXdl * sizeof(BDataType))) > N0 - ? N0 - : 128 / (BK1 * NPerXdl * sizeof(BDataType))); - - constexpr auto b_lds_block_desc = make_naive_tensor_descriptor_packed( - make_tuple(number{}, - number{}, - number{}, - number{}, - number{}, - BK1)); - - constexpr auto b_lds_block_desc_permuted = transform_tensor_descriptor( - b_lds_block_desc, - make_tuple( - make_pass_through_transform(number{}), - make_pass_through_transform(number{}), - make_xor_transform( - make_tuple(number{}, number{})), - make_pass_through_transform(number{}), - make_pass_through_transform(BK1)), - make_tuple( - sequence<0>{}, sequence<1>{}, sequence<2, 3>{}, sequence<4>{}, sequence<5>{}), - make_tuple( - sequence<0>{}, sequence<1>{}, sequence<2, 3>{}, sequence<4>{}, sequence<5>{})); - - constexpr auto b_lds_block_desc_unmerged = transform_tensor_descriptor( - b_lds_block_desc_permuted, - make_tuple( - make_pass_through_transform(number{}), - make_pass_through_transform(number{}), - make_unmerge_transform(make_tuple(number{}, number{})), - make_unmerge_transform(make_tuple(number{}, number{})), - make_pass_through_transform(number{}), - make_pass_through_transform(BK1)), - make_tuple(sequence<0>{}, - sequence<1>{}, - sequence<2>{}, - sequence<3>{}, - sequence<4>{}, - sequence<5>{}), - make_tuple(sequence<1>{}, - sequence<2>{}, - sequence<0, 3>{}, - sequence<4, 5>{}, - sequence<6>{}, - sequence<7>{})); - - // constexpr auto b_lds_block_desc_bk0_n_bk1 = transform_tensor_descriptor( - // b_lds_block_desc_unmerged, - // make_tuple(make_merge_transform_v3_division_mod( - // make_tuple(number{}, - // number{}, - // number{}, - // number{})), - // make_merge_transform_v3_division_mod( - // make_tuple(number{}, number{}, number{})), - // make_pass_through_transform(BK1)), - // make_tuple(sequence<0, 1, 4, 2>{}, sequence<5, 6, 3>{}, sequence<7>{}), - // make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{})); - - constexpr auto b_lds_block_desc_kn = transform_tensor_descriptor( - b_lds_block_desc_unmerged, - make_tuple(make_merge_transform_v3_division_mod( - make_tuple(number{}, - number{}, - number{}, - number{}, - BK1)), - make_merge_transform_v3_division_mod( - make_tuple(number{}, number{}, number{}))), - make_tuple(sequence<0, 1, 4, 2, 7>{}, sequence<5, 6, 3>{}), - make_tuple(sequence<1>{}, sequence<0>{})); - - // return b_lds_block_desc_bk0_n_bk1; - return b_lds_block_desc_kn; - - // constexpr auto b_lds_block_desc_bk0_n_bk1 = make_naive_tensor_descriptor( - // make_tuple(BK0, number{}, number{}), - // make_tuple(number{}, number{}, number<1>{}), - // number{}, - // number<1>{}); - - // constexpr auto b_lds_block_desc = transform_tensor_descriptor( - // b_lds_block_desc_bk0_n_bk1, - // make_tuple(make_pass_through_transform(number{}), - // make_merge_transform_v3_division_mod(make_tuple(BK0, - // number{}))), - // make_tuple(sequence<1>{}, sequence<0, 2>{}), - // make_tuple(sequence<0>{}, sequence<1>{})); - - // return b_lds_block_desc; - } -#endif + constexpr auto b_lds_block_desc = transform_tensor_descriptor( + b_lds_block_desc_xk0_mnldslayer_mn_xk1, + make_tuple(make_merge_transform( + make_tuple(number{}, number{})), + make_merge_transform( + make_tuple(number{}, number{}))), + make_tuple(sequence<1, 0>{}, sequence<2, 3>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + return b_lds_block_desc; } template diff --git a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/tile_gemm_shape.hpp b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/tile_gemm_shape.hpp index f0aa4472e1..0212c1c9f8 100644 --- a/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/tile_gemm_shape.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/instruction_schedule/tile_gemm_shape.hpp @@ -27,16 +27,6 @@ struct TileGemmShape static constexpr bool PermuteA = PermuteA_; static constexpr bool PermuteB = PermuteB_; - - CK_TILE_HOST static std::string GetName() - { - // clang-format off - return concat('_', "tile_gemm_shape", - concat('x', kM, kN, kK, NumWarps), - concat('x', BlockWarps::at(number<0>{}), BlockWarps::at(number<1>{}), BlockWarps::at(number<2>{})), - concat('x', (WarpTile::at(number<0>{})), WarpTile::at(number<1>{}), WarpTile::at(number<2>{}))); - // clang-format on - } }; } // namespace ck_tile