diff --git a/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg.hpp b/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg.hpp index a0fc8ed341..041d0cc1a5 100644 --- a/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg.hpp @@ -20,7 +20,124 @@ struct BlockGemmASmemBSmemCReg using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - static constexpr index_t kBlockSize = Problem::kBlockSize; + // using WarpGemmConfig = Policy::template GetWarpGemmMWarpNWarp(); + using WarpGemm = remove_cvref_t().template at<0>())>; + +#if 1 + /* + 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; + }; + */ + + CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode() + { + constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + constexpr index_t MIterPerWarp = BlockGemmShape::kM / (MWarp * WarpGemm::kM); + + constexpr index_t KPerBlock = BlockGemmShape::kK; + constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; + + constexpr auto a_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + 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 auto config = Policy::template GetWarpGemmMWarpNWarp(); + constexpr index_t MWarp = config.template at<1>(); + constexpr index_t NWarp = config.template at<2>(); + constexpr index_t NIterPerWarp = BlockGemmShape::kN / (NWarp * WarpGemm::kN); + + constexpr index_t KPerBlock = BlockGemmShape::kK; + constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; + + constexpr auto b_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + 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; + } +#endif + // C += A * B template diff --git a/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg_default_policy.hpp b/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg_default_policy.hpp index 5c8bda8e15..a6bc4e7563 100644 --- a/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg_default_policy.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/block_gemm_asmem_bsmem_creg_default_policy.hpp @@ -23,13 +23,13 @@ struct BlockGemmASmemBSmemCRegDefaultPolicy std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 2, 2); } else if constexpr(std::is_same_v && std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaBf16Bf16F32M32N32K8TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaBf16Bf16F32M32N32K8TransposedCDistribution{}, 2, 2); } #elif defined(USING_MFMA_32x32x_8x2) #pragma message ("mfma m32 n32 k16") @@ -37,13 +37,13 @@ struct BlockGemmASmemBSmemCRegDefaultPolicy std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 2, 2); } else if constexpr(std::is_same_v && std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaBf16Bf16F32M32N32K16TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaBf16Bf16F32M32N32K16TransposedCDistribution{}, 2, 2); } #elif defined(USING_MFMA_16x16x16) #pragma message ("mfma m16 n16 k16") @@ -51,13 +51,13 @@ struct BlockGemmASmemBSmemCRegDefaultPolicy std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaF16F16F32M16N16K16TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaF16F16F32M16N16K16TransposedCDistribution{}, 2, 2); } else if constexpr(std::is_same_v && std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaBf16Bf16F32M16N16K16TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaBf16Bf16F32M16N16K16TransposedCDistribution{}, 2, 2); } #elif defined(USING_MFMA_16x16x_16x2) #pragma message ("mfma m16 n16 k32") @@ -65,13 +65,13 @@ struct BlockGemmASmemBSmemCRegDefaultPolicy std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaF16F16F32M16N16K32TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaF16F16F32M16N16K32TransposedCDistribution{}, 2, 2); } else if constexpr(std::is_same_v && std::is_same_v && std::is_same_v) { - return make_tuple(WarpGemmMfmaBf16Bf16F32M16N16K32TransposedCDistribution{}, 4, 1); + return make_tuple(WarpGemmMfmaBf16Bf16F32M16N16K32TransposedCDistribution{}, 2, 2); } #endif else diff --git a/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg.hpp b/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg.hpp index 649d56336d..c9ffa1c95b 100644 --- a/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/block_gemm_pipeline_agmem_bgmem_creg.hpp @@ -52,6 +52,7 @@ struct BlockGemmPipelineAGmemBGmemCReg using CLayout = remove_cvref_t; using BlockGemm = remove_cvref_t())>; + using I0 = number<0>; using I1 = number<1>; using I2 = number<2>; @@ -64,8 +65,7 @@ struct BlockGemmPipelineAGmemBGmemCReg static constexpr index_t GetVectorSizeA() { return Policy::template GetVectorSizeA(); } static constexpr index_t GetVectorSizeB() { return Policy::template GetVectorSizeB(); } - static constexpr index_t GetSmemPackA() { return Policy::template GetSmemPackA(); } - static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB(); } + static constexpr index_t GetSmemPack() { return Policy::template GetSmemPack(); } static constexpr bool kPadM = Problem::kPadM; static constexpr bool kPadN = Problem::kPadN; @@ -97,12 +97,7 @@ struct BlockGemmPipelineAGmemBGmemCReg 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 AB_LDS_RW_Width = GetSmemPack(); constexpr index_t A_Buffer_Load_Inst_Num = MPerBlock * KPerBlock / (BlockSize * GetVectorSizeA()); @@ -110,14 +105,14 @@ struct BlockGemmPipelineAGmemBGmemCReg NPerBlock * KPerBlock / (BlockSize * GetVectorSizeB()); constexpr index_t A_LDS_Write_Inst_Num = - MPerBlock * KPerBlock / (BlockSize * A_LDS_Write_Width); + MPerBlock * KPerBlock / (BlockSize * AB_LDS_RW_Width); constexpr index_t B_LDS_Write_Inst_Num = - NPerBlock * KPerBlock / (BlockSize * B_LDS_Write_Width); + NPerBlock * KPerBlock / (BlockSize * AB_LDS_RW_Width); constexpr index_t A_LDS_Read_Inst_Num = - WaveNumN * MPerBlock * KPerBlock / (BlockSize * A_LDS_Read_Width); + WaveNumN * MPerBlock * KPerBlock / (BlockSize * AB_LDS_RW_Width); constexpr index_t B_LDS_Read_Inst_Num = - WaveNumM * NPerBlock * KPerBlock / (BlockSize * B_LDS_Read_Width); + WaveNumM * NPerBlock * KPerBlock / (BlockSize * AB_LDS_RW_Width); constexpr index_t C_MFMA_Inst_Num = MPerBlock * NPerBlock * KPerBlock / (BlockSize / WaveSize) / @@ -126,11 +121,11 @@ struct BlockGemmPipelineAGmemBGmemCReg // A/B split schedule // compiler is likely to use ds_read2 when instruction width smaller than 16bytes constexpr auto num_ds_read_inst_a = - A_LDS_Read_Width * sizeof(ADataType) / APackedSize == 16 ? A_LDS_Read_Inst_Num : - A_LDS_Read_Inst_Num / 2; + AB_LDS_RW_Width * sizeof(ADataType) / APackedSize == 16 ? A_LDS_Read_Inst_Num : + A_LDS_Read_Inst_Num / 2; constexpr auto num_ds_read_inst_b = - B_LDS_Read_Width * sizeof(BDataType) / BPackedSize == 16 ? B_LDS_Read_Inst_Num : - B_LDS_Read_Inst_Num / 2; + AB_LDS_RW_Width * sizeof(BDataType) / BPackedSize == 16 ? B_LDS_Read_Inst_Num : + B_LDS_Read_Inst_Num / 2; constexpr auto num_ds_write_inst_a = A_LDS_Write_Inst_Num; constexpr auto num_ds_write_inst_b = B_LDS_Write_Inst_Num; @@ -142,9 +137,9 @@ struct BlockGemmPipelineAGmemBGmemCReg constexpr auto mfma_cycle = NPerXDL == 16 ? 16 : 32; constexpr auto ds_read_a_issue_cycle = - A_LDS_Read_Width * sizeof(ADataType) / APackedSize == 16 ? 8 : 4; + AB_LDS_RW_Width * sizeof(ADataType) / APackedSize == 16 ? 8 : 4; constexpr auto ds_read_b_issue_cycle = - B_LDS_Read_Width * sizeof(BDataType) / BPackedSize == 16 ? 8 : 4; + AB_LDS_RW_Width * sizeof(BDataType) / BPackedSize == 16 ? 8 : 4; constexpr auto ds_read_a_mfma_rate = (mfma_cycle - 4 + 2 * ds_read_a_issue_cycle - 1) / (2 * ds_read_a_issue_cycle); constexpr auto ds_read_b_mfma_rate = @@ -231,6 +226,38 @@ struct BlockGemmPipelineAGmemBGmemCReg }); } + // template + // CK_TILE_DEVICE constexpr auto GetAWindows(const ADramBlockWindowTmp& a_dram_block_window_tmp, + // const ALdsTensorView& a_lds_block_view, + // const ALdsLoadTileDistr&) const + // { + // constexpr bool is_col_major = std::is_same_v; + + // using YPerTile = std::conditional_t, number>; + // using XPerTile = std::conditional_t, number>; + + // // A DRAM tile window for load + // auto a_copy_dram_window = + // make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), + // make_tuple(YPerTile{}, XPerTile{}), + // a_dram_block_window_tmp.get_window_origin(), + // Policy::template MakeADramTileDistribution()); + + // // A LDS tile window for store + // auto a_copy_lds_window = make_tile_window( + // a_lds_block_view, make_tuple(number{}, number{}), {0, 0}); + + // auto a_lds_gemm_window = + // make_tile_window(a_lds_block_view, + // make_tuple(number{}, number{}), + // {0, 0}, + // ALdsLoadTileDistr{}); + + // return make_tuple(std::move(a_copy_dram_window), + // std::move(a_copy_lds_window), + // std::move(a_lds_gemm_window)); + // } + template CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const AElementFunction& a_element_func, + [[maybe_unused]] const AElementFunction& a_element_func, const BDramBlockWindowTmp& b_dram_block_window_tmp, - const BElementFunction& b_element_func, + [[maybe_unused]] const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { - static_assert( - std::is_same_v> && - std::is_same_v>, - "A/B Dram block window should have the same data type as appropriate " - "([A|B]DataType) defined in Problem definition!"); - // ------------------------------------------------------------------------------------ // Definitions of all needed tiles +#if 1 // A/B tiles in LDS - auto&& [a_lds_block, b_lds_block] = Base::GetABLdsTensorViews(p_smem); + // auto&& [a_lds_block, b_lds_block] = Base::GetABLdsTensorViews(p_smem); + ADataType* p_a_lds = static_cast(p_smem); + constexpr auto a_lds_block_desc = Policy::template MakeALdsBlockDescriptor(); + auto a_lds_block = make_tensor_view(p_a_lds, a_lds_block_desc); + constexpr index_t a_lds_block_space_size_aligned = + integer_divide_ceil(sizeof(ADataType) * a_lds_block_desc.get_element_space_size(), + 16) * 16; + BDataType* p_b_lds = static_cast( + static_cast(static_cast(p_smem) + a_lds_block_space_size_aligned)); + constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor(); + auto b_lds_block = make_tensor_view(p_b_lds, b_lds_block_desc); // Tile distribution for load from lds - constexpr auto a_lds_load_tile_distr = - make_static_tile_distribution(BlockGemm::MakeABlockDistributionEncode()); - constexpr auto b_lds_load_tile_distr = - make_static_tile_distribution(BlockGemm::MakeBBlockDistributionEncode()); + constexpr auto a_lds_load_tile_distr = make_static_tile_distribution(BlockGemm::MakeABlockDistributionEncode()); + constexpr auto b_lds_load_tile_distr = make_static_tile_distribution(BlockGemm::MakeBBlockDistributionEncode()); // A DRAM tile window for load // A LDS tile window for store // A LDS tile for block GEMM - auto&& [a_copy_dram_window, a_copy_lds_window, a_lds_gemm_window] = - Base::GetAWindows(a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr); + auto&& [a_copy_dram_window, a_copy_lds_window, a_lds_gemm_window] = Base::GetAWindows(a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr); // B DRAM tile window for load // B LDS tile window for store // B LDS tile for block GEMM - auto&& [b_copy_dram_window, b_copy_lds_window, b_lds_gemm_window] = - Base::GetBWindows(b_dram_block_window_tmp, b_lds_block, b_lds_load_tile_distr); + auto&& [b_copy_dram_window, b_copy_lds_window, b_lds_gemm_window] = Base::GetBWindows(b_dram_block_window_tmp, b_lds_block, b_lds_load_tile_distr); + +#else + // A tile in LDS + ADataType* p_a_lds = static_cast(p_smem); + + constexpr auto a_lds_block_desc = Policy::template MakeALdsBlockDescriptor(); + + auto a_lds_block = make_tensor_view(p_a_lds, a_lds_block_desc); + + constexpr index_t a_lds_block_space_size_aligned = + integer_divide_ceil(sizeof(ADataType) * a_lds_block_desc.get_element_space_size(), + 16) * 16; + + // B tile in LDS + BDataType* p_b_lds = static_cast( + static_cast(static_cast(p_smem) + a_lds_block_space_size_aligned)); + + constexpr auto b_lds_block_desc = Policy::template MakeBLdsBlockDescriptor(); + + auto b_lds_block = make_tensor_view(p_b_lds, b_lds_block_desc); + + // A DRAM tile window for load + auto a_copy_dram_window = + make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_dram_block_window_tmp.get_window_origin(), + Policy::template MakeADramTileDistribution()); + + // A LDS tile window for store + auto a_copy_lds_window = + make_tile_window(a_lds_block, + make_tuple(number{}, number{}), + {0, 0}, + a_copy_dram_window.get_tile_distribution()); + + // B DRAM tile window for load + auto b_copy_dram_window = + make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_dram_block_window_tmp.get_window_origin(), + Policy::template MakeBDramTileDistribution()); + + // B LDS tile window for store + auto b_copy_lds_window = + make_tile_window(b_lds_block, + make_tuple(number{}, number{}), + {0, 0}, + b_copy_dram_window.get_tile_distribution()); + + // A LDS tile for block GEMM + auto a_lds_gemm_window = make_tile_window( + a_lds_block, make_tuple(number{}, number{}), {0, 0}); + + // B LDS tile for block GEMM + auto b_lds_gemm_window = make_tile_window( + b_lds_block, make_tuple(number{}, number{}), {0, 0}); +#endif // Block GEMM auto block_gemm = BlockGemm(); - auto c_block_tile = block_gemm.MakeCBlockTile(); + // auto c_block_tile = block_gemm.MakeCBlockTile(); + // Acc register tile + auto c_block_tile = decltype(block_gemm(a_lds_gemm_window, b_lds_gemm_window)){}; using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); @@ -299,23 +385,76 @@ struct BlockGemmPipelineAGmemBGmemCReg // ----------------------------------------------------------------------------------------- // Gemm pipeline start + // template + // CK_TILE_DEVICE void GlobalPrefetch(DstBlockTile& dst_block_tile, + // SrcTileWindow& dram_tile_window, + // const DramTileWindowStep& dram_tile_window_step) const + // { + // load_tile(dst_block_tile, dram_tile_window); + // move_tile_window(dram_tile_window, dram_tile_window_step); + // } + + // template + // CK_TILE_DEVICE void LocalPrefill(DstTileWindow& lds_tile_window, + // const SrcBlockTile& src_block_tile, + // const ElementFunction& element_func) const + // { + // const auto block_tile_tmp = tile_elementwise_in(element_func, src_block_tile); + // store_tile(lds_tile_window, block_tile_tmp); + // } + + // template + // CK_TILE_DEVICE void LocalPrefetch(DstBlockTile& dst_block_tile, + // const SrcTileWindow& lds_tile_window) const + // { + // load_tile(dst_block_tile, lds_tile_window); + // } + // prefetch // global read 0 - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + // Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); + load_tile(a_block_tile, a_copy_dram_window); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + // Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + load_tile(b_block_tile, b_copy_dram_window); + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); // LDS write 0 - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + // Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + store_tile(a_copy_lds_window, a_block_tile); + // Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + store_tile(b_copy_lds_window, b_block_tile); - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + // Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); + load_tile(a_block_tile, a_copy_dram_window); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + // Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + load_tile(b_block_tile, b_copy_dram_window); + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); block_sync_lds(); - block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); + // block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); + + // auto a_block_tile = load_tile(a_copy_dram_window); + // auto b_block_tile = load_tile(b_copy_dram_window); + + // static constexpr auto ALdsTileDistr = decltype(make_static_tile_distribution(MakeABlockDistributionEncode())){}; + // static constexpr auto BLdsTileDistr = decltype(make_static_tile_distribution(MakeBBlockDistributionEncode())){}; + // using ALdsTile = decltype(make_static_distributed_tensor(ALdsTileDistr)); + // using BLdsTile = decltype(make_static_distributed_tensor(BLdsTileDistr)); + // ALdsTile a_warp_tile_; + // ALdsTile b_warp_tile_; + // template + // CK_TILE_DEVICE void LocalPrefetch(const ASmemBlockWindow& a_block_window, + // const BSmemBlockWindow& b_block_window) + // { + // load_tile(a_warp_tile_, a_block_window); + // load_tile(b_warp_tile_, b_block_window); + // } + // } __builtin_amdgcn_sched_barrier(0); @@ -326,17 +465,24 @@ struct BlockGemmPipelineAGmemBGmemCReg do { block_sync_lds(); - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + // Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + store_tile(a_copy_lds_window, a_block_tile); + // Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + store_tile(b_copy_lds_window, b_block_tile); + + // Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); + load_tile(a_block_tile, a_copy_dram_window); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + // Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + load_tile(b_block_tile, b_copy_dram_window); + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); block_sync_lds(); - block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); + // block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); HotLoopScheduler(); __builtin_amdgcn_sched_barrier(0); @@ -354,10 +500,13 @@ struct BlockGemmPipelineAGmemBGmemCReg { block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); block_sync_lds(); - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + // Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + store_tile(a_copy_lds_window, a_block_tile); + // Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + store_tile(b_copy_lds_window, b_block_tile); + block_sync_lds(); - block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); + // block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); } // __builtin_amdgcn_sched_barrier(0); 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 6bde5eb939..ced7bce756 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 @@ -343,25 +343,14 @@ struct BlockGemmPipelineAGmemBGmemCRegDefaultPolicy } template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackA() + CK_TILE_HOST_DEVICE static constexpr auto GetSmemPack() { constexpr index_t kKPack = 8; return kKPack; - // using BlockGemm = remove_cvref_t())>; - // constexpr index_t KPack = BlockGemm::Traits::KPack; - // return KPack; - } - - template - CK_TILE_HOST_DEVICE static constexpr auto GetSmemPackB() - { - constexpr index_t kKPack = 8; - return kKPack; - // using BlockGemm = remove_cvref_t())>; - // constexpr index_t KPack = BlockGemm::Traits::KPack; - // return KPack; } +#endif +#if 0 template CK_TILE_HOST_DEVICE static constexpr auto GetBlockGemm() { @@ -385,7 +374,8 @@ struct BlockGemmPipelineAGmemBGmemCRegDefaultPolicy template CK_TILE_HOST_DEVICE static constexpr auto GetBlockGemm() { - return BlockGemmASmemBSmemCReg{}; + using Policy = BlockGemmASmemBSmemCRegDefaultPolicy; + return BlockGemmASmemBSmemCReg{}; } #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 5fffcf7e36..100dd82f61 100644 --- a/example/ck_tile/99_toy_example/02_gemm/gemm.cpp +++ b/example/ck_tile/99_toy_example/02_gemm/gemm.cpp @@ -146,17 +146,17 @@ int main(int argc, char* argv[]) constexpr ck_tile::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock; using gemm_kernel = ck_tile::Gemm; + BDataType, + AccDataType, + CDataType, + CElementFunction, + kAAlignment, + kBAlignment, + kCAlignment, + kBlockSize, + kGemmMPerBlock, + kGemmNPerBlock, + kGemmKPerBlock>; float ave_time = ck_tile::launch_kernel(ck_tile::stream_config{nullptr, true, 0, 5, 1000}, diff --git a/example/ck_tile/99_toy_example/02_gemm/gemm.hpp b/example/ck_tile/99_toy_example/02_gemm/gemm.hpp index 421f442f3d..9b196c96b8 100644 --- a/example/ck_tile/99_toy_example/02_gemm/gemm.hpp +++ b/example/ck_tile/99_toy_example/02_gemm/gemm.hpp @@ -233,11 +233,8 @@ struct Gemm } #if defined(ENABLE_INSTRUCTION_SCH) - static constexpr index_t M_Tile = 128; - static constexpr index_t N_Tile = 128; - static constexpr index_t K_Tile = 64; - static constexpr index_t M_Warp = 2; - static constexpr index_t N_Warp = 2; + static constexpr index_t M_Warp = 4; + static constexpr index_t N_Warp = 1; static constexpr index_t K_Warp = 1; static constexpr index_t M_Warp_Tile = 16; static constexpr index_t N_Warp_Tile = 16; @@ -256,7 +253,7 @@ struct Gemm { #if defined(ENABLE_INSTRUCTION_SCH) // Block GEMM pipeline w/ instruction scheduling - using GemmShape = TileGemmShape, + using GemmShape = TileGemmShape, sequence, sequence, PermuteA,