ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 00:58:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add async_load for tile_scatter_gather

Browse Source
This commit is contained in:
Feng Shijie
2025-09-12 05:32:45 +00:00
parent a44d35172d
commit fdbb22d2b6
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
include/ck_tile/core/tensor/tile_scatter_gather.hpp
View File

@@ -522,18 +522,22 @@ struct tile_scatter_gather
constexpr auto idx_gather = idx_ys_start[number<YsGatherDim>{}];
const auto page_offset = page_idx_[idx_gather];
// merge page_offset into bottom_coord
auto mixed_bottom_thread_coord = bottom_tensor_thread_coord;
mixed_bottom_thread_coord.get_hidden_index()[number<0>{}] += page_offset;
// 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,
mixed_bottom_thread_coord,
number<0>{},
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,
mixed_bottom_thread_coord,
number<0>{},
valids_[idx_gather],
bool_constant<oob_conditional_check>{});

2
include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp
View File

@@ -7,7 +7,7 @@
namespace ck_tile {
#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE 1
#define CKTILE_FLATMM_USE_BUFFER_LOAD_LDS_AS_POSSIBLE 0
#if defined(__gfx950__)
#define CKTILE_FLATMM_ARCH_SUPPORT_BUFFER_LOAD_LDS_DWORDx4 1