mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-17 00:58:44 +00:00
Add async_load for tile_scatter_gather
This commit is contained in:
Feng Shijie
@@ -475,6 +475,96 @@ struct tile_scatter_gather
|
||||
});
|
||||
}
|
||||
|
||||
template <typename LdsTileWindow_,
|
||||
index_t i_access_unsupport_ = -1,
|
||||
bool oob_conditional_check = true>
|
||||
CK_TILE_DEVICE auto async_load(LdsTileWindow_&& lds_tile,
|
||||
number<i_access_unsupport_> = {},
|
||||
bool_constant<oob_conditional_check> = {}) const
|
||||
{
|
||||
using LdsTileWindow = remove_cvref_t<LdsTileWindow_>;
|
||||
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<iCoord * NumAccessPerCoord + iCoordAccess>{};
|
||||
|
||||
// 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<YsGatherDim>{}];
|
||||
const auto page_offset = page_idx_[idx_gather];
|
||||
|
||||
// read from bottom tensor
|
||||
if constexpr(std::is_same_v<ValidArray, std::nullptr_t>)
|
||||
this->get_bottom_tensor_view().template async_get_vectorized_elements<vector_t>(
|
||||
smem,
|
||||
bottom_tensor_thread_coord,
|
||||
page_offset,
|
||||
bool_constant<oob_conditional_check>{});
|
||||
else
|
||||
this->get_bottom_tensor_view().template async_get_vectorized_elements<vector_t>(
|
||||
smem,
|
||||
bottom_tensor_thread_coord,
|
||||
page_offset,
|
||||
valids_[idx_gather],
|
||||
bool_constant<oob_conditional_check>{});
|
||||
|
||||
// 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<NDimY>{});
|
||||
|
||||
constexpr auto idx_diff_ps_ys = container_concat(
|
||||
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}),
|
||||
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<NDimP>{}),
|
||||
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 <typename LdsTileWindow_,
|
||||
index_t i_access_unsupport_ = -1,
|
||||
@@ -963,6 +1053,31 @@ make_tile_scatter_gather(const TensorView_& tensor_view,
|
||||
tensor_view, window_lengths, origin, tile_distribution, page_idx, valids};
|
||||
}
|
||||
|
||||
template <typename NewTensorView_,
|
||||
typename OldTensorView_,
|
||||
typename WindowLengths_,
|
||||
typename StaticTileDistribution_,
|
||||
typename StaticPageIndexArray_,
|
||||
typename StaticValidArray_,
|
||||
index_t HsGatherDim = 0,
|
||||
index_t NumCoord = 1>
|
||||
CK_TILE_DEVICE auto replace_bottom_tensor_view(const NewTensorView_& new_tensor_view,
|
||||
const tile_scatter_gather<OldTensorView_,
|
||||
WindowLengths_,
|
||||
StaticTileDistribution_,
|
||||
StaticPageIndexArray_,
|
||||
StaticValidArray_,
|
||||
HsGatherDim,
|
||||
NumCoord>& 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 <typename TensorView,
|
||||
typename WindowLengths,
|
||||
typename StaticTileDistribution,
|
||||
|
||||
@@ -837,6 +837,24 @@ make_tile_window_raw(const TensorView_& tensor_view,
|
||||
return w;
|
||||
}
|
||||
|
||||
template <typename NewTensorView_,
|
||||
typename OldTensorView_,
|
||||
typename WindowLengths_,
|
||||
typename StaticTileDistribution_,
|
||||
index_t NumCoord = 1>
|
||||
CK_TILE_DEVICE auto
|
||||
replace_bottom_tensor_view(const NewTensorView_& new_tensor_view,
|
||||
const tile_window_with_static_distribution<OldTensorView_,
|
||||
WindowLengths_,
|
||||
StaticTileDistribution_,
|
||||
NumCoord>& 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 <typename TensorView_,
|
||||
typename WindowLengths_,
|
||||
typename StaticTileDistribution_,
|
||||
@@ -946,6 +964,15 @@ make_tile_window_raw(const tile_window_with_static_lengths<TensorView, WindowLen
|
||||
return w;
|
||||
}
|
||||
|
||||
template <typename NewTensorView_, typename OldTensorView_, typename WindowLengths_>
|
||||
CK_TILE_DEVICE auto replace_bottom_tensor_view(
|
||||
const NewTensorView_& new_tensor_view,
|
||||
const tile_window_with_static_lengths<OldTensorView_, WindowLengths_>& tile_window)
|
||||
{
|
||||
return make_tile_window(
|
||||
new_tensor_view, tile_window.get_window_lengths(), tile_window.get_window_origin());
|
||||
}
|
||||
|
||||
template <typename TensorView_, typename WindowLengths_>
|
||||
CK_TILE_DEVICE void move_tile_window(
|
||||
tile_window_with_static_lengths<TensorView_, WindowLengths_>& window,
|
||||
|
||||
@@ -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<Problem>(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<Problem>(
|
||||
a_copy_dram_window_.get_bottom_tensor_view()),
|
||||
a_copy_dram_window_);
|
||||
|
||||
// A tile in LDS
|
||||
ADataType* p_a_lds_ping = static_cast<ADataType*>(p_smem_ping);
|
||||
|
||||
Reference in New Issue
Block a user