From 5a2d93d436d8ce13dbd7d549c8685d0078dadca1 Mon Sep 17 00:00:00 2001 From: coderfeli Date: Thu, 28 Nov 2024 10:33:42 +0000 Subject: [PATCH] revert code --- include/ck_tile/core/tensor/load_tile.hpp | 16 ++++++ .../core/tensor/tile_window_linear.hpp | 52 +++++++++++++++++ .../block/block_gemm_areg_breg_creg_v2.hpp | 2 +- .../gemm_pipeline_agmem_bgmem_creg_v1.hpp | 56 +++++++++++++++---- 4 files changed, 113 insertions(+), 13 deletions(-) diff --git a/include/ck_tile/core/tensor/load_tile.hpp b/include/ck_tile/core/tensor/load_tile.hpp index f150fc54ca..6c54aa9e0c 100644 --- a/include/ck_tile/core/tensor/load_tile.hpp +++ b/include/ck_tile/core/tensor/load_tile.hpp @@ -46,6 +46,22 @@ CK_TILE_DEVICE auto load_tile(const tile_window_linear{}, bool_constant{}); } +template +CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile, + const tile_window_linear& tile_window, + bool_constant = {}) +{ + return tile_window.load(dst_tile, number<-1>{}, bool_constant{}); +} + template + CK_TILE_DEVICE auto load(DistributedTensor dst_tensor, number = {}, bool_constant = {}) const + { + using vector_t = typename traits::vector_t; + using SFC_Ys = typename traits::SFC_Ys; + + constexpr auto tile_dstr = TileDstr{}; + + auto issue = [&](auto i_access_) { + constexpr auto IAccess = number{}; + + constexpr auto non_linear_id = number{}; + auto bottom_tensor_thread_coord = cached_coords_[non_linear_id]; + auto bottom_tensor_flag = cached_flags_[IAccess]; + + constexpr auto linear_offset = get_bottom_linear_offset(IAccess); + + // read from bottom tensor + const vector_t vec_value = + get_bottom_tensor_view().template get_vectorized_elements( + bottom_tensor_thread_coord, + linear_offset, + bottom_tensor_flag, + bool_constant{}); +#if 1 + // data index [y0, y1, ...] + constexpr auto idx_diff_ys = SFC_Ys::get_index(IAccess); + // write into distributed tensor + static_for<0, traits::ScalarPerVector, 1>{}([&](auto j) { + constexpr auto idx_ys = generate_tuple( + [&](auto jj) { + return jj == traits::VectorDimY ? (idx_diff_ys[jj] + j) : idx_diff_ys[jj]; + }, + number{}); + + constexpr index_t d = tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys); + + dst_tensor.get_thread_buffer().template at() = + vec_value.template get_as()[j]; + }); +#else + constexpr index_t d = tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys_start); + static_assert(d % traits::ScalarPerVector == 0); + + dst_tensor.get_thread_buffer().template get_as()( + number{}) = bit_cast(vec_value); +#endif + }; + + WINDOW_DISPATCH_ISSUE(); + } + template CK_TILE_DEVICE auto load(number = {}, bool_constant = {}) const { 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 be7cc23c88..e0c4f4afba 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 @@ -210,7 +210,7 @@ struct BlockGemmARegBRegCRegV2 auto tileDist = BlockTensor::get_tile_distribution(); return load_tile(block_tensor, make_tile_window(block_window, tileDist)); - // load_tile_raw(block_tensor, make_tile_window_linear_raw(block_window, tileDist)); + // load_tile(block_tensor, make_tile_window_linear(block_window, tileDist)); // return; } 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 6155c87868..f14b648dce 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 @@ -256,8 +256,40 @@ struct GemmPipelineAGmemBGmemCRegV1 using BLdsTile = decltype(make_static_distributed_tensor(BLdsTileDistr{})); ALdsTile a_block_tile0; BLdsTile b_block_tile0; - load_tile(a_block_tile0, make_tile_window(a_lds_window0, ALdsTileDistr{})); - load_tile(b_block_tile0, make_tile_window(b_lds_window0, BLdsTileDistr{})); + Policy::template BlockGemm::PrefetchLds(a_lds_window0, a_block_tile0); + Policy::template BlockGemm::PrefetchLds(b_lds_window0, b_block_tile0); + + // 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; ", type_convert(a_block_tile0(i_j_idx)), type_convert(b_block_tile0(i_j_idx))); + // }); + // printf("\n"); + // }); + // } + // if (threadIdx.x == 0) { + // printf("aalds\n"); + // 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,", type_convert(a_block_tile0(i_j_idx))); + // }); + // printf("\n"); + // }); + // printf("bbbbblds\n"); + // constexpr auto span_2d2 = decltype(b_block_tile0)::get_distributed_spans(); + // sweep_tile_span(span_2d2[number<0>{}], [&](auto idx0) { + // sweep_tile_span(span_2d2[number<1>{}], [&](auto idx1) { + // constexpr auto i_j_idx = make_tuple(idx0, idx1); + // printf("%f,", type_convert(b_block_tile0(i_j_idx))); + // }); + // printf("\n"); + // }); + // } + // LDS write 1 LocalPrefill(a_lds_window1, a_global_load_tile, a_element_func); LocalPrefill(b_lds_window1, b_global_load_tile, b_element_func); @@ -274,8 +306,8 @@ struct GemmPipelineAGmemBGmemCRegV1 // ping { block_sync_lds(); - load_tile(a_block_tile1, make_tile_window(a_lds_window1, ALdsTileDistr{})); - load_tile(b_block_tile1, make_tile_window(b_lds_window1, BLdsTileDistr{})); + Policy::template BlockGemm::PrefetchLds(a_lds_window1, a_block_tile1); + Policy::template BlockGemm::PrefetchLds(b_lds_window1, b_block_tile1); LocalPrefill(a_lds_window0, a_global_load_tile, a_element_func); LocalPrefill(b_lds_window0, b_global_load_tile, b_element_func); GlobalPrefetch(a_global_load_tile, a_copy_dram_window); @@ -286,8 +318,8 @@ struct GemmPipelineAGmemBGmemCRegV1 // pong { block_sync_lds(); - load_tile(a_block_tile0, make_tile_window(a_lds_window0, ALdsTileDistr{})); - load_tile(b_block_tile0, make_tile_window(b_lds_window0, BLdsTileDistr{})); + Policy::template BlockGemm::PrefetchLds(a_lds_window0, a_block_tile0); + Policy::template BlockGemm::PrefetchLds(b_lds_window0, b_block_tile0); LocalPrefill(a_lds_window1, a_global_load_tile, a_element_func); LocalPrefill(b_lds_window1, b_global_load_tile, b_element_func); GlobalPrefetch(a_global_load_tile, a_copy_dram_window); @@ -303,8 +335,8 @@ struct GemmPipelineAGmemBGmemCRegV1 // 3 { block_sync_lds(); - load_tile(a_block_tile1, make_tile_window(a_lds_window1, ALdsTileDistr{})); - load_tile(b_block_tile1, make_tile_window(b_lds_window1, BLdsTileDistr{})); + Policy::template BlockGemm::PrefetchLds(a_lds_window1, a_block_tile1); + Policy::template BlockGemm::PrefetchLds(b_lds_window1, b_block_tile1); LocalPrefill(a_lds_window0, a_global_load_tile, a_element_func); LocalPrefill(b_lds_window0, b_global_load_tile, b_element_func); block_gemm(c_block_tile, a_block_tile0, b_block_tile0); @@ -312,8 +344,8 @@ struct GemmPipelineAGmemBGmemCRegV1 // 2 { block_sync_lds(); - load_tile(a_block_tile0, make_tile_window(a_lds_window0, ALdsTileDistr{})); - load_tile(b_block_tile0, make_tile_window(b_lds_window0, BLdsTileDistr{})); + Policy::template BlockGemm::PrefetchLds(a_lds_window0, a_block_tile0); + Policy::template BlockGemm::PrefetchLds(b_lds_window0, b_block_tile0); block_gemm(c_block_tile, a_block_tile1, b_block_tile1); } //1 @@ -324,8 +356,8 @@ struct GemmPipelineAGmemBGmemCRegV1 } else { { block_sync_lds(); - load_tile(a_block_tile1, make_tile_window(a_lds_window1, ALdsTileDistr{})); - load_tile(b_block_tile1, make_tile_window(b_lds_window1, BLdsTileDistr{})); + Policy::template BlockGemm::PrefetchLds(a_lds_window1, a_block_tile1); + Policy::template BlockGemm::PrefetchLds(b_lds_window1, b_block_tile1); block_gemm(c_block_tile, a_block_tile0, b_block_tile0); } // 2