diff --git a/example/ck_tile/18_flatmm/flatmm_basic.cpp b/example/ck_tile/18_flatmm/flatmm_basic.cpp index 8929c170e3..8a94d749f3 100644 --- a/example/ck_tile/18_flatmm/flatmm_basic.cpp +++ b/example/ck_tile/18_flatmm/flatmm_basic.cpp @@ -34,6 +34,16 @@ float flatmm_calc(const ck_tile::FlatmmHostArgs<>& args, const ck_tile::stream_c FlatmmConfig::N_Warp_Tile, FlatmmConfig::K_Warp_Tile>>; + // static constexpr index_t kM = BlockTile::at(number<0>{}); + // static constexpr index_t kN = BlockTile::at(number<1>{}); + // static constexpr index_t kK = BlockTile::at(number<2>{}); + + // static constexpr bool PermuteA = PermuteA_; + // static constexpr bool PermuteB = PermuteB_; + + // static constexpr index_t flatNPerWarp = BlockWarps::at(number<1>{}); // 4 + // static constexpr index_t flatKPerWarp = WarpTile::at(number<2>{}) * WarpTile::at(number<1>{});// 16 * 64 + // static constexpr index_t flatKPerBlock = flatKPerWarp * kK / WarpTile::at(number<2>{}); // 16 * 128 using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner{}; constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t{}; - + constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t{}; + static_assert(BWarpDstr::NDimY==1); __builtin_amdgcn_sched_barrier(0); // A tile in LDS @@ -647,28 +652,38 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV // B flat DRAM window for load auto b_flat_distribution = PipelinePolicy::template MakeBFlatDramTileDistribution(); + auto b_flat_distribution1 = + PipelinePolicy::template MakeBFlatDramFullTileDistribution(); + + // auto b_flat_dram_window = // tile_window_with_static_distribution + // make_tile_window( + // b_flat_dram_block_window_tmp.get_bottom_tensor_view(), // from kernel gemm_pad_views + // make_tuple(number{}, number{}), + // b_flat_dram_block_window_tmp.get_window_origin(), + // b_flat_distribution); auto b_flat_dram_window = // tile_window_with_static_distribution - make_tile_window( + make_tile_window_linear( b_flat_dram_block_window_tmp.get_bottom_tensor_view(), // from kernel gemm_pad_views make_tuple(number{}, number{}), b_flat_dram_block_window_tmp.get_window_origin(), - b_flat_distribution); + b_flat_distribution1); // pingpong buffer for B - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> - b_flat_dram_windows; + // statically_indexed_array< + // statically_indexed_array, + // NIterPerWarp> + // b_flat_dram_windows; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> - b_warp_tensor_ping; - - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> - b_warp_tensor_pong; + // statically_indexed_array< + // statically_indexed_array, + // NIterPerWarp> + // b_warp_tensor_ping; + // statically_indexed_array< + // statically_indexed_array, + // NIterPerWarp> + // b_warp_tensor_pong; + decltype(load_tile(b_flat_dram_window)) b_warp_tensor_ping; + decltype(load_tile(b_flat_dram_window)) b_warp_tensor_pong; // Prefetch A0 @@ -677,16 +692,17 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV move_tile_window(a_copy_dram_window, {0, kKPerBlock}); // prefetch B - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; + // static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { + // static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + // b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; - move_tile_window(b_flat_dram_windows(nIter)(kIter), - {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); + // move_tile_window(b_flat_dram_windows(nIter)(kIter), + // {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); - }); - }); + // b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + // }); + // }); + b_warp_tensor_ping = load_tile(b_flat_dram_window); // move B window to next flat K move_tile_window(b_flat_dram_window, {0, BlockGemmShape::flatKPerBlock}); @@ -745,16 +761,18 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV while(iCounter > 0) { // prefetch B(2i+1) - static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; + // static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + // static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { + // b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; - move_tile_window(b_flat_dram_windows(nIter)(kIter), - {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); + // move_tile_window(b_flat_dram_windows(nIter)(kIter), + // {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); - }); - }); + // b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + // }); + // }); + + b_warp_tensor_pong = load_tile(b_flat_dram_window); // Prefill A(2i+1) a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile); @@ -777,10 +795,16 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV a_warp_tensor.get_thread_buffer() = a_warp_tensor_ping(AwarpIter).get_y_sliced_thread_data( merge_sequences(sequence{}, a_warp_y_index_zeros), merge_sequences(sequence<1>{}, a_warp_y_lengths)); + BWarpTensor b_warp_tensor; + b_warp_tensor.get_thread_buffer() = b_warp_tensor_ping.get_y_sliced_thread_data( + merge_sequences(sequence{}, b_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, b_warp_y_lengths)); + static_assert(b_warp_tensor.get_thread_buffer_size() == 16); + static_assert(b_warp_tensor_ping.get_thread_buffer_size() == 64); c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - + // if constexpr(mIter==0 && nIter ==0) // if(threadIdx.x % 16== 0 && threadIdx.x<64){ // for(int i=0;i{}([&](auto kIter) { - static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { - b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; + // static_for<0, KIterPerWarp, 1>{}([&](auto kIter) { + // static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { + // b_flat_dram_windows(nIter)(kIter) = b_flat_dram_window; - move_tile_window(b_flat_dram_windows(nIter)(kIter), - {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); + // move_tile_window(b_flat_dram_windows(nIter)(kIter), + // {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); - }); - }); + // b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + // }); + // }); + + b_warp_tensor_ping = load_tile(b_flat_dram_window); // Prefill A(2i+2) a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile); @@ -873,7 +900,12 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV // } // } - WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor_pong(nIter)(kIter)); + + BWarpTensor b_warp_tensor; + b_warp_tensor.get_thread_buffer() = b_warp_tensor_pong.get_y_sliced_thread_data( + merge_sequences(sequence{}, b_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, b_warp_y_lengths)); + WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor); // write C warp tensor into C block tensor c_block_tile.set_y_sliced_thread_data( @@ -1029,7 +1061,11 @@ struct FlatmmPipelineAGmemBGmemCRegV1 : public BaseFlatmmPipelineAGmemBGmemCRegV c_warp_tensor.get_thread_buffer() = c_block_tile.get_y_sliced_thread_data( merge_sequences(sequence{}, c_warp_y_index_zeros), merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); - WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor_ping(nIter)(kIter)); + BWarpTensor b_warp_tensor; + b_warp_tensor.get_thread_buffer() = b_warp_tensor_ping.get_y_sliced_thread_data( + merge_sequences(sequence{}, b_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, b_warp_y_lengths)); + WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor); // write C warp tensor into C block tensor c_block_tile.set_y_sliced_thread_data( diff --git a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp index 9e1daa6f16..cba7bd780d 100644 --- a/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp @@ -23,38 +23,38 @@ struct UniversalFlatmmPipelineAgBgCrPolicy constexpr index_t MPerXdl = Problem::BlockGemmShape::WarpTile::at(I0); constexpr index_t NPerXdl = Problem::BlockGemmShape::WarpTile::at(I1); - if constexpr(MPerXdl == 16 && NPerXdl == 16) - { - /*reduce transform layers,compare with old ck*/ - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPack = GetSmemPackA(); + // if constexpr(MPerXdl == 16 && NPerXdl == 16) + // { + // /*reduce transform layers,compare with old ck*/ + // constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + // constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + // constexpr index_t KPack = GetSmemPackA(); - constexpr auto a_lds_block_desc_0 = make_naive_tensor_descriptor( - make_tuple(number{}, number{}, number{}), - make_tuple(number{}, number{}, number<1>{}), - number{}, - number<1>{}); + // 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_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 = transform_tensor_descriptor( - a_lds_block_desc_permuted, - make_tuple(make_pass_through_transform(number{}), - make_merge_transform_v3_division_mod( - make_tuple(number{}, number{}))), - make_tuple(sequence<1>{}, sequence<0, 2>{}), - make_tuple(sequence<0>{}, sequence<1>{})); + // constexpr auto a_lds_block_desc = transform_tensor_descriptor( + // a_lds_block_desc_permuted, + // make_tuple(make_pass_through_transform(number{}), + // make_merge_transform_v3_division_mod( + // make_tuple(number{}, number{}))), + // make_tuple(sequence<1>{}, sequence<0, 2>{}), + // make_tuple(sequence<0>{}, sequence<1>{})); - return a_lds_block_desc; - } - else + // return a_lds_block_desc; + // } + // else { constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; @@ -406,7 +406,58 @@ struct UniversalFlatmmPipelineAgBgCrPolicy sequence<1, 1, 2, 2>, sequence<0, 3, 0, 3>>{}); } + template + CK_TILE_HOST_DEVICE static constexpr auto MakeBFlatDramFullTileDistribution() + { + using TileShape = typename Problem::BlockGemmShape; // ck_tile::TileFlatmmShape + // constexpr index_t BlockSize = Problem::kBlockSize; + // constexpr index_t WaveSize = get_warp_size(); + // constexpr index_t WaveNum = BlockSize / WaveSize; + // constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + // constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + + // constexpr index_t KBPerLoad = GetKBPerLoad(); + // constexpr index_t KThdPerWave = WaveSize; // threads cnt in K dim + // constexpr index_t KWavePerBlk = 1; + // constexpr index_t KRepeat = 2; + // static_assert(TileShape::flatKPerWarp == KThdPerWave * KBPerLoad, "wrong"); + + // constexpr index_t NBPerLoad = 1; + // constexpr index_t NThdPerWave = 1; + // constexpr index_t NWavePerBlk = TileShape::BlockWarps::at(number<1>{}); // N_Warp + // constexpr index_t NRepeat = 2; + + // constexpr index_t WaveRepeat = WaveNum / TileShape::flatNPerWarp; + + // // static constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN); + // // static constexpr index_t KIterPerWarp = kKPerBlock / WG::kK; + // return make_static_tile_distribution( + // tile_distribution_encoding< + // sequence, // ? + // tuple, // second direction + // sequence>, // first direction + // // wave in blk, // thd in wave + // // // + // tuple, sequence<1, 2>>, // which direction + // tuple, sequence<2, 2>>, // which index + // // + // sequence<1, 1, 2, 2>, + // sequence<0, 3, 0, 3>>{}); + + return make_static_tile_distribution( + tile_distribution_encoding< + sequence<1>, // ? + tuple, // second direction + sequence<2, 64, 16>>, // first direction + // wave in blk, // thd in wave + // // + tuple, sequence<2>>, // which direction + tuple, sequence<1>>, // which index + // + sequence<1, 2, 2>, + sequence<0, 0, 2>>{}); + } template CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledARegBlockDistribution() {