diff --git a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp index ed6e45846dfb..a03c7b8d82e1 100644 --- a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp +++ b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp @@ -396,22 +396,28 @@ struct tile_window_with_static_distribution constexpr auto iAccess = number{}; if constexpr(Begin <= iAccess && iAccess < End) { - constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - constexpr auto lds_access_offset = [&]() { + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + const index_t lds_access_offset = [&]() { constexpr auto idx_off_ys = SFC_Ys::get_step_between(number<0>{}, iAccess); constexpr auto adapter_ys_offset = make_tensor_adaptor_coordinate( tile_dstr.get_ps_ys_to_xs_adaptor(), container_concat(array{0}, to_array(idx_off_ys))); - constexpr auto coord_ys_offset = make_tensor_coordinate( - typename Base::BottomTensorView{}.get_tensor_descriptor(), + // The ys index step is compile-time, but the descriptor's lengths and + // strides are not necessarily so. A default-constructed BottomTensorView + // has zero strides, which resolves every access past the first to offset + // 0 whenever the descriptor is built from runtime extents; take the + // offset from the real descriptor instead. + const auto coord_ys_offset = make_tensor_coordinate( + this->get_bottom_tensor_view().get_tensor_descriptor(), adapter_ys_offset.get_bottom_index()); return coord_ys_offset.get_offset(); }(); const vector_t vec_value = - this->get_bottom_tensor_view() - .template get_vectorized_elements( - bottom_tensor_thread_coord, linear_off, bool_constant{}); + this->get_bottom_tensor_view().template get_vectorized_elements( + bottom_tensor_thread_coord, + linear_off + lds_access_offset, + bool_constant{}); static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) { constexpr auto idx_ys = generate_tuple( [&](auto jj) { @@ -493,26 +499,30 @@ struct tile_window_with_static_distribution // data index [y0, y1, ...] constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - // Compute compile-time offset from access 0 to current access - constexpr auto lds_access_offset = [&]() { + // offset from access 0 to current access + const index_t lds_access_offset = [&]() { constexpr auto idx_off_ys = SFC_Ys::get_step_between(number<0>{}, iAccess); constexpr auto adapter_ys_offset = make_tensor_adaptor_coordinate( tile_dstr.get_ps_ys_to_xs_adaptor(), container_concat(array{0}, to_array(idx_off_ys))); - constexpr auto coord_ys_offset = make_tensor_coordinate( - typename Base::BottomTensorView{}.get_tensor_descriptor(), + // The ys index step is compile-time, but the descriptor's lengths and + // strides are not necessarily so. A default-constructed BottomTensorView + // has zero strides, which resolves every access past the first to offset + // 0 whenever the descriptor is built from runtime extents; take the + // offset from the real descriptor instead. + const auto coord_ys_offset = make_tensor_coordinate( + this->get_bottom_tensor_view().get_tensor_descriptor(), adapter_ys_offset.get_bottom_index()); return coord_ys_offset.get_offset(); }(); - // read from bottom tensor with compile-time offset + // read from bottom tensor const vector_t vec_value = - this->get_bottom_tensor_view() - .template get_vectorized_elements( - bottom_tensor_thread_coord, - linear_off, - bool_constant{}); + this->get_bottom_tensor_view().template get_vectorized_elements( + bottom_tensor_thread_coord, + linear_off + lds_access_offset, + bool_constant{}); // write into distributed tensor static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) { constexpr auto idx_ys = generate_tuple( diff --git a/projects/composablekernel/test/ck_tile/load_and_convert_tile/kernel.hpp b/projects/composablekernel/test/ck_tile/load_and_convert_tile/kernel.hpp index 515d396e4293..e113df4c2b3b 100644 --- a/projects/composablekernel/test/ck_tile/load_and_convert_tile/kernel.hpp +++ b/projects/composablekernel/test/ck_tile/load_and_convert_tile/kernel.hpp @@ -126,7 +126,11 @@ struct LoadAndConvertKernel static constexpr index_t kBlockSize = Problem::BlockShape::BlockSize; - CK_TILE_HOST static auto BlockSize() { return kBlockSize; } + // kBlockSize is derived from get_warp_size(), which on the host pass always reports the + // wave64 size. On a wave32 device the kernel only addresses half that many threads, so + // launching kBlockSize threads leaves the upper half writing results nobody asked for -- + // which masked the LDS offset bug this test is meant to catch. + CK_TILE_HOST static auto BlockSize() { return is_wave32() ? kBlockSize / 2 : kBlockSize; } private: CK_TILE_DEVICE static constexpr auto get_block_dims()