diff --git a/example/ck_tile/03_gemm/gemm_basic.cpp b/example/ck_tile/03_gemm/gemm_basic.cpp index ef7b3cac48..0f9f575082 100644 --- a/example/ck_tile/03_gemm/gemm_basic.cpp +++ b/example/ck_tile/03_gemm/gemm_basic.cpp @@ -22,9 +22,6 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) constexpr bool kPadN = false; constexpr bool kPadK = false; - constexpr bool kTilePermute = false; - // The rank and permutation will also be generate out by the CodeGen part. - constexpr ck_tile::index_t kOutputRank = 2; constexpr int kBlockPerCu = 1; @@ -36,6 +33,7 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) constexpr ck_tile::index_t M_Warp = 2; constexpr ck_tile::index_t N_Warp = 2; constexpr ck_tile::index_t K_Warp = 1; + constexpr ck_tile::index_t Warp_Size = 64; constexpr ck_tile::index_t M_Warp_Tile = 32; constexpr ck_tile::index_t N_Warp_Tile = 32; @@ -43,8 +41,6 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) // Whether doing the CShuffle (transpose before the global memory), depending on the output // layout. - constexpr bool CShuffleEpilogue = - std::is_same_v; using CodegenGemmShape = ck_tile::TileGemmShape, @@ -52,21 +48,13 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) ck_tile::sequence>; using TilePartitioner = ck_tile::GemmTilePartitioner; - - using GemmEpilogue = std::conditional_t< - CShuffleEpilogue, - ck_tile::CShuffleEpilogue>, - ck_tile::Default2DEpilogue< - ck_tile::Default2DEpilogueProblem>>; + using GemmEpilogue = ck_tile::CShuffleEpilogueV2>; using CodegenGemmTraits = ck_tile::TileGemmTraits; diff --git a/include/ck_tile/ops/epilogue.hpp b/include/ck_tile/ops/epilogue.hpp index c24744bdbc..cc6b0024d8 100644 --- a/include/ck_tile/ops/epilogue.hpp +++ b/include/ck_tile/ops/epilogue.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/ops/epilogue/cshuffle_epilogue.hpp" +#include "ck_tile/ops/epilogue/cshuffle_epilogue_v2.hpp" #include "ck_tile/ops/epilogue/default_2d_epilogue.hpp" #include "ck_tile/ops/epilogue/dynamic_quant_epilogue.hpp" #include "ck_tile/ops/common/generic_2d_block_shape.hpp" diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue_v2.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue_v2.hpp new file mode 100644 index 0000000000..fa1d5249ba --- /dev/null +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue_v2.hpp @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" + +namespace ck_tile { + +// this epilogue just store out a M*N matrix, row major + +template +struct CShuffleEpilogueV2Problem +{ + using AccDataType = remove_cvref_t; + using ODataType = remove_cvref_t; + // static constexpr bool UseRawStore = UseRawStore_; + static constexpr index_t kBlockSize = kBlockSize_; + static constexpr index_t MPerBlock = kM_; + static constexpr index_t NPerBlock = kN_; + static constexpr bool kPadM = kPadM_; + static constexpr bool kPadN = kPadN_; +}; + + +template +CK_TILE_HOST_DEVICE static constexpr auto MakeOLdsBlockDescriptor() +{ + static constexpr index_t kMPerBlock = Problem::MPerBlock; + static constexpr index_t kNPerBlock = Problem::NPerBlock; + + return make_naive_tensor_descriptor( + make_tuple(number{}, number{}), + make_tuple(number{}, number<1>{}), + number<1>{}, + number<1>{}); +} + + +CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return 65536; } +template +CK_TILE_HOST_DEVICE static constexpr auto MakeODramTileDistribution() +{ + static constexpr index_t kMPerBlock = Problem::MPerBlock; + static constexpr index_t kNPerBlock = Problem::NPerBlock; + static constexpr index_t BlockSize = Problem::kBlockSize; + static constexpr index_t WaveSize = get_warp_size(); + using ODataType = remove_cvref_t; + + // using OLayout = remove_cvref_t; + // if constexpr(std::is_same_v) + // { + // static_assert(0, "not impl"); + // } + constexpr index_t N2 = 8; + constexpr index_t N1 = min(kNPerBlock / N2, WaveSize); + constexpr index_t N0 = integer_divide_ceil(kNPerBlock / N2, WaveSize); + + constexpr index_t M2 = integer_divide_ceil(WaveSize, kNPerBlock / N2); + constexpr index_t M1 = BlockSize / WaveSize; + constexpr index_t M0 = integer_divide_ceil(kMPerBlock, M1 * M2); + + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 1>>, + sequence<1, 2, 2>, + sequence<0, 0, 2>>{}); +} + + +template +struct CShuffleEpilogueV2 +{ + using Problem = remove_cvref_t; + using AccDataType = remove_cvref_t; + using ODataType = remove_cvref_t; + static constexpr bool kPadM = Problem::kPadM; + static constexpr bool kPadN = Problem::kPadN; + static constexpr bool UseRawStore = Problem::UseRawStore; + static constexpr bool kMPerBlock = Problem::MPerBlock; + static constexpr bool kNPerBlock = Problem::NPerBlock; + + // constexpr auto a_warp_y_lengths = + // to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + // constexpr auto c_warp_y_lengths = + // to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + + // constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; + // constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + + // static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + // static_for<0, MIterPerWarp, 1>{}([&](auto mIter) { + // dst_warp_tensor.get_thread_buffer() = ds_in.get_y_sliced_thread_data( + // merge_sequences(sequence{}, c_warp_y_index_zeros), + // merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + + // dst_out.set_y_sliced_thread_data( + // merge_sequences(sequence{}, a_warp_y_index_zeros), + // merge_sequences(sequence<1, 1>{}, a_warp_y_lengths), + // dst_warp_tensor.get_thread_buffer()); + // }); + // }); + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return kMPerBlock * kNPerBlock * sizeof(ODataType); } + + // TODO: this function assume store out vector size is the same as OAccTile last dimension size + // how do we fix this ? + template + CK_TILE_DEVICE auto operator()(ODramWindowTmp& o_dram_window_tmp, const OAccTile& o_acc_tile, void *p_smem) + { + auto o_lds_tile = cast_tile(o_acc_tile); + constexpr auto o_lds_block_desc = MakeOLdsBlockDescriptor(); + auto o_lds_block = make_tensor_view(static_cast(p_smem), o_lds_block_desc); + auto o_lds_window0 = make_tile_window(o_lds_block, make_tuple(number{}, number{}), {0, 0}); + + store_tile(o_lds_window0, o_lds_tile); + block_sync_lds(); + // if (threadIdx.x == 0) { + // printf("%f, %f\n",type_convert(static_cast(p_smem)[32767]), type_convert(static_cast(p_smem)[32768])); + // constexpr auto span_2d = decltype(o_lds_tile)::get_distributed_spans(); + // sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { + // sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { + // constexpr auto i_j_idx = make_tuple(idx0, idx1); + // printf("%f,", type_convert(o_lds_tile(i_j_idx))); + // }); + // printf("\n"); + // }); + // } + auto o_dram_distri = MakeODramTileDistribution(); + auto o_dram_tile = load_tile(make_tile_window(o_lds_window0, o_dram_distri)); + store_tile(o_dram_window_tmp, o_dram_tile); + } +}; +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp index e0c4f4afba..9126fd866d 100644 --- a/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v2.hpp @@ -169,6 +169,23 @@ struct BlockGemmARegBRegCRegV2 return c_block_tensor; } + CK_TILE_DEVICE static constexpr auto MakeCBlockSubTile() + { + constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding< + sequence<>, + tuple, sequence>, + tuple>, + tuple>, + sequence<2>, + sequence<0>>{}; + + constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + c_block_outer_dstr_encoding, typename WG::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; + } + CK_TILE_DEVICE static constexpr auto MakeABlockDistribution() { // M->N Warp diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index 96af6e8260..fed88f1cb3 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -211,8 +211,19 @@ struct GemmKernel c_pad_view, make_tuple(number{}, number{}), {i_m, i_n}); - - EpiloguePipeline{}(CBlockWindow_pad, c_block_tile); + + using CSubTileDistr = decltype(GemmPipeline::MakeCBlockSubTile()); + static_for<0, GemmPipeline::NumCSubTile(), 1>{}([&](auto i_m0) + { + auto c_sub_tile = make_static_distributed_tensor(CSubTileDistr{}); + constexpr auto c_sub_y_index_zeros = uniform_sequence_gen_t{}; + constexpr auto c_sub_y_lengths = to_sequence(CSubTileDistr{}.get_ys_to_d_descriptor().get_lengths()); + c_sub_tile.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( + merge_sequences(sequence{}, c_sub_y_index_zeros), + merge_sequences(sequence<1>{}, c_sub_y_lengths)); + EpiloguePipeline{}(CBlockWindow_pad, c_sub_tile, smem_ptr); + move_tile_window(CBlockWindow_pad, {TilePartitioner::kM / GemmPipeline::NumCSubTile(), 0}); + }); } }; diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp index f14b648dce..8ccb5b1185 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp @@ -11,13 +11,15 @@ namespace ck_tile { // A Tile Window: global memory // B Tile Window: global memory // C Distributed tensor: register -template +template struct GemmPipelineAGmemBGmemCRegV1 { using ADataType = remove_cvref_t; using BDataType = remove_cvref_t; using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; + using Policy = Policy_; + using Problem = Problem; using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; @@ -133,6 +135,16 @@ struct GemmPipelineAGmemBGmemCRegV1 __builtin_amdgcn_sched_barrier(0); } + CK_TILE_DEVICE static constexpr auto MakeCBlockSubTile() { + return Policy::template BlockGemm::MakeCBlockSubTile(); + } + + CK_TILE_DEVICE static constexpr auto NumCSubTile() { + constexpr index_t MPerXDL = BlockGemmShape::WarpTile::at(number<0>{});//32 + constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(number<0>{});//2 + return integer_divide_ceil(kMPerBlock, WaveNumM * MPerXDL); + } + template {}], [&](auto idx0) { - // sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { - // constexpr auto i_j_idx = make_tuple(idx0, idx1); - // printf("%f,", type_convert(a_global_load_tile(i_j_idx))); - // }); - // printf("\n"); - // }); - // printf("bbbbb\n"); - // constexpr auto span_2d2 = decltype(b_global_load_tile)::get_distributed_spans(); - // sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { - // sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { - // constexpr auto i_j_idx = make_tuple(idx0, idx1); - // printf("%f,", type_convert(b_global_load_tile(i_j_idx))); - // }); - // printf("\n"); - // }); - // } ////////////// LDS desc, window & register ///////////////// // AB LDS desc constexpr auto a_lds_block_desc = Policy::template MakeALdsBlockDescriptor(); @@ -363,18 +355,16 @@ struct GemmPipelineAGmemBGmemCRegV1 // 2 { block_gemm(c_block_tile, a_block_tile1, b_block_tile1); - // if (threadIdx.x == 64) { - // constexpr auto span_2d = decltype(a_block_tile0)::get_distributed_spans(); - // sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { - // sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { - // constexpr auto i_j_idx = make_tuple(idx0, idx1); - // printf("%f, %f; %f, %f. ", type_convert(a_block_tile0(i_j_idx)), type_convert(b_block_tile0(i_j_idx)), type_convert(a_block_tile1(i_j_idx)), type_convert(b_block_tile1(i_j_idx))); - // }); - // printf("\n"); - // }); - // } } } + + /// cccccccccc + // constexpr auto c_lds_block_desc = Policy::template MakeCLdsBlockDescriptor(); + // auto c_lds_block = make_tensor_view(reinterpret_cast(p_smem), c_lds_block_desc); + // auto c_lds_window0 = make_tile_window(c_lds_block, make_tuple(number{}, number{}), {0, 0}); + // store_tile(c_lds_window0, c_block_tile); + // block_sync_lds(); + return c_block_tile; }