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

refine apis

Browse Source
This commit is contained in:
coderfeli
2025-04-08 03:42:46 +00:00
parent 8d62ff557a
commit 4b5859f717
6 changed files with 17 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
include/ck_tile/core/tensor/load_tile.hpp
View File

@@ -18,32 +18,10 @@
namespace ck_tile {
template <typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
index_t NumCoord,
template <typename TileWindow_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(const tile_window_with_static_distribution<BottomTensorView_,
WindowLengths_,
TileDistribution_,
NumCoord>& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{
return tile_window.load(number<i_access>{}, bool_constant<oob_conditional_check>{});
}
template <typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
typename LinearBottomDims_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(const tile_window_linear<BottomTensorView_,
WindowLengths_,
TileDistribution_,
LinearBottomDims_>& tile_window,
CK_TILE_DEVICE auto load_tile(const TileWindow_& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{
@@ -51,35 +29,11 @@ CK_TILE_DEVICE auto load_tile(const tile_window_linear<BottomTensorView_,
}
template <typename DistributedTensor_,
typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
index_t NumCoord,
typename TileWindow_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile,
const tile_window_with_static_distribution<BottomTensorView_,
WindowLengths_,
TileDistribution_,
NumCoord>& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{
return tile_window.load(dst_tile, number<i_access>{}, bool_constant<oob_conditional_check>{});
}
template <typename DistributedTensor_,
typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
typename LinearBottomDims_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile,
const tile_window_linear<BottomTensorView_,
WindowLengths_,
TileDistribution_,
LinearBottomDims_>& tile_window,
const TileWindow_& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{

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

@@ -602,45 +602,6 @@ make_tile_scatter_gather(const TensorView_& tensor_view,
tensor_view, window_lengths, origin, tile_distribution, page_idx};
}
// this version can't be called in a constexpr context
// template <typename TensorView_,
// typename WindowLengths_,
// typename StaticTileDistribution_,
// index_t NumCoord = 1>
// CK_TILE_DEVICE auto
// make_tile_window_raw(const TensorView_& tensor_view,
// const WindowLengths_& window_lengths,
// const multi_index<TensorView_::get_num_of_dimension()>& origin,
// const StaticTileDistribution_& tile_distribution,
// number<NumCoord> = {})
// {
// auto w = tile_scatter_gather<remove_cvref_t<TensorView_>,
// remove_cvref_t<WindowLengths_>,
// remove_cvref_t<StaticTileDistribution_>,
// NumCoord>{
// tensor_view, window_lengths, origin, tile_distribution};
// w.init_raw();
// return w;
// }
// template <typename TensorView_,
// typename WindowLengths_,
// typename StaticTileDistribution_,
// index_t NumCoord>
// CK_TILE_DEVICE void move_tile_window(
// tile_scatter_gather<TensorView_,
// WindowLengths_,
// StaticTileDistribution_,
// NumCoord>& window,
// const typename tile_scatter_gather<TensorView_,
// WindowLengths_,
// StaticTileDistribution_,
// NumCoord>::BottomTensorIndex& step)
// {
// window.move(step);
// }
template <typename TensorView, typename WindowLengths, typename StaticTileDistribution, typename StaticPageIndexArray, index_t PageIndexDim>
CK_TILE_DEVICE constexpr auto
make_tile_scatter_gather(const tile_window_with_static_lengths<TensorView, WindowLengths>& tile_window,

26
include/ck_tile/core/tensor/tile_window.hpp
View File

@@ -1097,23 +1097,6 @@ make_tile_window_raw(const TensorView_& tensor_view,
return w;
}
template <typename TensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
index_t NumCoord>
CK_TILE_DEVICE void move_tile_window(
tile_window_with_static_distribution<TensorView_,
WindowLengths_,
StaticTileDistribution_,
NumCoord>& window,
const typename tile_window_with_static_distribution<TensorView_,
WindowLengths_,
StaticTileDistribution_,
NumCoord>::BottomTensorIndex& step)
{
window.move(step);
}
/**
* @brief This class provides description of tile windowed view on the device memory.
*
@@ -1242,13 +1225,4 @@ make_tile_window_raw(const tile_window_with_static_lengths<TensorView, WindowLen
return w;
}
template <typename TensorView_, typename WindowLengths_>
CK_TILE_DEVICE void move_tile_window(
tile_window_with_static_lengths<TensorView_, WindowLengths_>& window,
const typename tile_window_with_static_lengths<TensorView_, WindowLengths_>::BottomTensorIndex&
step)
{
window.move(step);
}
} // namespace ck_tile

15
include/ck_tile/core/tensor/tile_window_linear.hpp
View File

@@ -1200,19 +1200,4 @@ make_tile_window_linear_raw(const TileWindow_& tile_window,
LinearBottomDims_{});
}
template <typename TensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
typename LinearBottomDims_>
CK_TILE_DEVICE void move_tile_window(
tile_window_linear<TensorView_, WindowLengths_, StaticTileDistribution_, LinearBottomDims_>&
window,
const typename tile_window_linear<TensorView_,
WindowLengths_,
StaticTileDistribution_,
LinearBottomDims_>::BottomTensorIndex& step)
{
window.move(step);
}
} // namespace ck_tile

8
include/ck_tile/core/tensor/tile_window_utils.hpp
View File

@@ -18,6 +18,14 @@
#pragma once
namespace ck_tile {
template <typename TileWindow_>
CK_TILE_DEVICE void move_tile_window(
TileWindow_& window,
const typename TileWindow_::BottomTensorIndex& step)
{
window.move(step);
}
// input a lds store tile, extract some information from it
// used to set m0 value for gfx9 serious
template <typename LdsTileWindow_>

15
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp
View File

@@ -296,15 +296,12 @@ struct BlockFmhaPipelineQRKSVS
k_dist,
k_offsets); // K DRAM tile window for
// auto k_block_tile = load_tile(k_dram_window);
auto k_block_tile = k_dram_window.load();
auto k_block_tile = load_tile(k_dram_window);
{
// move_tile_window(k_dram_window, {0, kK0});
k_dram_window.move({0, kK0});
move_tile_window(k_dram_window, {0, kK0});
clear_tile(s_acc); // initialize C
store_tile(k_lds_window, tile_elementwise_in(k_element_func, k_block_tile));
k_block_tile = k_dram_window.load();
k_block_tile = load_tile(k_dram_window);
}
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
@@ -329,14 +326,12 @@ struct BlockFmhaPipelineQRKSVS
sequence<kM0, (i_k0 + 1) * kK0>{}),
k_lds_window);
block_sync_lds();
k_dram_window.move({0, kK0});
// move_tile_window(k_dram_window, {0, kK0});
move_tile_window(k_dram_window, {0, kK0});
store_tile(
k_lds_window,
tile_elementwise_in(k_element_func, k_block_tile)); // LDS write i + 1
k_block_tile = k_dram_window.load(); // global read i + 2
// k_block_tile = load_tile(k_dram_window); // global read i + 2
k_block_tile = load_tile(k_dram_window); // global read i + 2
});
}