diff --git a/include/ck_tile/ops/transpose/block_transpose.hpp b/include/ck_tile/ops/transpose/block_transpose.hpp index e19fda8ede..fca17a6135 100644 --- a/include/ck_tile/ops/transpose/block_transpose.hpp +++ b/include/ck_tile/ops/transpose/block_transpose.hpp @@ -149,25 +149,40 @@ struct BlockTranspose auto y = load_tile_transpose(load_from_lds_window); - constexpr auto lds_load_distr = Policy::template MakeLdsLoadTileDistribution(); - // constexpr auto glb_store_distr = Policy::template MakeOutputDistribution(); - - constexpr auto y_in_desc = lds_load_distr.get_ys_to_d_descriptor(); - // constexpr auto y_out_desc = glb_store_distr{}.get_ys_to_d_descriptor(); - - constexpr index_t NDimY = lds_load_distr.get_num_of_dimension_y(); - - constexpr auto y_in_lengths = to_sequence(y_in_desc.get_lengths()); - constexpr index_t vecLoadSize = y_in_lengths[NDimY - 1]; - // constexpr auto lds_distr_y_indx_zeros = uniform_sequence_gen_t())::NDimY, 0>{}; auto out_tensor = make_static_distributed_tensor( Policy::template MakeOutputDistribution()); - using InVec = array; - using OutVec = array; - out_tensor.get_thread_buffer().template set_as( - number<0>{}, y.get_thread_buffer().template get_as(number<0>{})); + constexpr auto lds_load_distr = Policy::template MakeLdsLoadTileDistribution(); + constexpr auto glb_store_distr = Policy::template MakeOutputDistribution(); + + constexpr auto y_in_desc = lds_load_distr.get_ys_to_d_descriptor(); + constexpr auto y_out_desc = glb_store_distr.get_ys_to_d_descriptor(); + + constexpr index_t NDimYIn = lds_load_distr.get_num_of_dimension_y(); + constexpr index_t NDimYOut = glb_store_distr.get_num_of_dimension_y(); + + constexpr auto y_in_lengths = to_sequence(y_in_desc.get_lengths()); + constexpr auto y_out_lengths = to_sequence(y_out_desc.get_lengths()); + + constexpr auto y_in_element_space_size = y_in_desc.get_element_space_size(); + constexpr auto y_out_element_space_size = y_out_desc.get_element_space_size(); + static_assert(y_in_element_space_size == y_out_element_space_size, + "the element space size is not the same!"); + static_assert(y_in_lengths[NDimYIn - 1] == y_out_lengths[NDimYOut - 1], + "the vector length is not the same!"); + constexpr index_t vecLoadSize = y_in_lengths[NDimYIn - 1]; + constexpr index_t num_of_access = + reduce_on_sequence(y_in_lengths, multiplies{}, number<1>{}) / vecLoadSize; + + // constexpr auto lds_distr_y_indx_zeros = uniform_sequence_gen_t())::NDimY, 0>{}; + + using DataVec = array; + static_for<0, num_of_access, 1>{}([&](auto iAccess) { + out_tensor.get_thread_buffer().template set_as( + number{}, + y.get_thread_buffer().template get_as(number{})); + }); store_tile(output_tile_window, out_tensor); }