diff --git a/include/ck_tile/core/tensor/tile_scatter_gather.hpp b/include/ck_tile/core/tensor/tile_scatter_gather.hpp index 082207d1df..9d16b7cde2 100644 --- a/include/ck_tile/core/tensor/tile_scatter_gather.hpp +++ b/include/ck_tile/core/tensor/tile_scatter_gather.hpp @@ -475,6 +475,96 @@ struct tile_scatter_gather }); } + template + CK_TILE_DEVICE auto async_load(LdsTileWindow_&& lds_tile, + number = {}, + bool_constant = {}) const + { + using LdsTileWindow = remove_cvref_t; + using LdsDataType = typename LdsTileWindow::DataType; + using Traits = load_store_traits; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = TileDstr{}; + + // Precompute invariant values outside loops + const auto window_origin = lds_tile.get_window_origin(); + const auto& bottom_tensor_view = lds_tile.get_bottom_tensor_view(); + const auto& tensor_descriptor = bottom_tensor_view.get_tensor_descriptor(); + auto smem_base_ptr = bottom_tensor_view.get_buffer_view().p_data_; + + // loop over thread tensor space [y0, y1, ...] + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + /// TODO: use structure binding (to be captured later) if compiled in C++20 + auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + auto lds_window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0]; + auto lds_bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + + // Use precomputed window origin + auto lds_bottom_tensor_thread_idx = + window_origin + lds_window_adaptor_thread_coord.get_bottom_index(); + // Use precomputed tensor descriptor + const auto lds_coord = + make_tensor_coordinate(tensor_descriptor, lds_bottom_tensor_thread_idx); + // Calculate SMEM address using base pointer + CK_TILE_LDS_ADDR LdsDataType* smem = smem_base_ptr + lds_coord.get_offset(); + + // data index [y0, y1, ...] + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + constexpr auto idx_gather = idx_ys_start[number{}]; + const auto page_offset = page_idx_[idx_gather]; + + // read from bottom tensor + if constexpr(std::is_same_v) + this->get_bottom_tensor_view().template async_get_vectorized_elements( + smem, + bottom_tensor_thread_coord, + page_offset, + bool_constant{}); + else + this->get_bottom_tensor_view().template async_get_vectorized_elements( + smem, + bottom_tensor_thread_coord, + page_offset, + valids_[idx_gather], + bool_constant{}); + + // move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto forward_step_scatter = generate_tuple( + [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + number{}); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + forward_step_scatter); + // lds_diff doesn't need to mask the difference of the gather-dim. + constexpr auto lds_idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + idx_diff_ys); + + move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + move_window_adaptor_and_bottom_tensor_thread_coordinate( + lds_window_adaptor_thread_coord, + lds_bottom_tensor_thread_coord, + lds_idx_diff_ps_ys); + } + }); + }); + } + // TODO: currently async load only implemented in inline asm template +CK_TILE_DEVICE auto replace_bottom_tensor_view(const NewTensorView_& new_tensor_view, + const tile_scatter_gather& tile_window) +{ + return make_tile_scatter_gather(new_tensor_view, + tile_window.window_lengths_, + tile_window.window_origin_, + tile_window.tile_dstr_, + tile_window.page_idx_, + tile_window.valids_); +} + template +CK_TILE_DEVICE auto +replace_bottom_tensor_view(const NewTensorView_& new_tensor_view, + const tile_window_with_static_distribution& tile_window) +{ + return make_tile_window(new_tensor_view, + tile_window.get_window_lengths(), + tile_window.get_window_origin(), + tile_window.get_tile_distribution()); +} + template +CK_TILE_DEVICE auto replace_bottom_tensor_view( + const NewTensorView_& new_tensor_view, + const tile_window_with_static_lengths& tile_window) +{ + return make_tile_window( + new_tensor_view, tile_window.get_window_lengths(), tile_window.get_window_origin()); +} + template CK_TILE_DEVICE void move_tile_window( tile_window_with_static_lengths& window, diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index 2908845d3c..69e5917949 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -529,12 +529,10 @@ struct F16xMXF4FlatmmPipelineAGmemBGmemCRegV1 __builtin_amdgcn_sched_barrier(0); - auto a_dram_view = a_copy_dram_window_.get_bottom_tensor_view(); - auto a_copy_dram_window = make_tile_window( - PipelinePolicy::template TransformF16xF4_ATensorView(a_dram_view), - a_copy_dram_window_.get_window_lengths(), - a_copy_dram_window_.get_window_origin(), - a_copy_dram_window_.get_tile_distribution()); + auto a_copy_dram_window = replace_bottom_tensor_view( + PipelinePolicy::template TransformF16xF4_ATensorView( + a_copy_dram_window_.get_bottom_tensor_view()), + a_copy_dram_window_); // A tile in LDS ADataType* p_a_lds_ping = static_cast(p_smem_ping);