ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-07 08:15:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Only expose necessary methods (not attributes)

Browse Source
This commit is contained in:
PoYen, Chen
2024-08-13 07:26:26 +00:00
parent c54de6416a
commit 19c19d8bd3
1 changed files with 12 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
include/ck_tile/ops/fmha/block/page_block_navigator.hpp
View File

@@ -46,13 +46,6 @@ struct TrivialPageBlockNavigator
return /*block_index=*/0;
}
/// TODO: remove this method after finish debuging
CK_TILE_DEVICE static constexpr int32_t
get_block_index(const WindowOrigin& /*global_window_origin*/)
{
return /*block_index=*/0;
}
CK_TILE_DEVICE static constexpr WindowOrigin
to_local_window_origin(const WindowOrigin& global_window_origin)
{
@@ -169,17 +162,6 @@ struct PageBlockNavigator
tile_window.set_bottom_tensor_view_data_ptr(get_block_ptr(new_block_index));
}
CK_TILE_HOST_DEVICE
DataType* get_block_ptr(index_t block_index) const
{
return physical_blocks + physical_block_indices[block_index] * block_stride + fixed_offset;
}
CK_TILE_HOST_DEVICE int32_t get_block_index(const WindowOrigin& global_window_origin) const
{
return integer_divide_floor(global_window_origin.at(number<VirtualDim>{}), page_block_size);
}
CK_TILE_HOST_DEVICE WindowOrigin
to_local_window_origin(const WindowOrigin& global_window_origin) const
{
@@ -216,6 +198,18 @@ struct PageBlockNavigator
}
}
private:
CK_TILE_HOST_DEVICE
DataType* get_block_ptr(index_t block_index) const
{
return physical_blocks + physical_block_indices[block_index] * block_stride + fixed_offset;
}
CK_TILE_HOST_DEVICE int32_t get_block_index(const WindowOrigin& global_window_origin) const
{
return integer_divide_floor(global_window_origin.at(number<VirtualDim>{}), page_block_size);
}
DataType* physical_blocks;
long_index_t block_stride;
long_index_t fixed_offset;