From 4b5859f717e4eb7657aac3f2f98dfefef9586ec1 Mon Sep 17 00:00:00 2001 From: coderfeli Date: Tue, 8 Apr 2025 03:42:46 +0000 Subject: [PATCH] refine apis --- include/ck_tile/core/tensor/load_tile.hpp | 54 ++----------------- .../core/tensor/tile_scatter_gather.hpp | 39 -------------- include/ck_tile/core/tensor/tile_window.hpp | 26 --------- .../core/tensor/tile_window_linear.hpp | 15 ------ .../ck_tile/core/tensor/tile_window_utils.hpp | 8 +++ .../pipeline/block_fmha_pipeline_qr_ks_vs.hpp | 15 ++---- 6 files changed, 17 insertions(+), 140 deletions(-) diff --git a/include/ck_tile/core/tensor/load_tile.hpp b/include/ck_tile/core/tensor/load_tile.hpp index b280a1725d..173f23441b 100644 --- a/include/ck_tile/core/tensor/load_tile.hpp +++ b/include/ck_tile/core/tensor/load_tile.hpp @@ -18,32 +18,10 @@ namespace ck_tile { -template -CK_TILE_DEVICE auto load_tile(const tile_window_with_static_distribution& tile_window, - number = {}, - bool_constant = {}) -{ - return tile_window.load(number{}, bool_constant{}); -} - -template -CK_TILE_DEVICE auto load_tile(const tile_window_linear& tile_window, +CK_TILE_DEVICE auto load_tile(const TileWindow_& tile_window, number = {}, bool_constant = {}) { @@ -51,35 +29,11 @@ CK_TILE_DEVICE auto load_tile(const tile_window_linear CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile, - const tile_window_with_static_distribution& tile_window, - number = {}, - bool_constant = {}) -{ - return tile_window.load(dst_tile, number{}, bool_constant{}); -} - -template -CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile, - const tile_window_linear& tile_window, + const TileWindow_& tile_window, number = {}, bool_constant = {}) { diff --git a/include/ck_tile/core/tensor/tile_scatter_gather.hpp b/include/ck_tile/core/tensor/tile_scatter_gather.hpp index 3e3148e5bc..3233ceb1d0 100644 --- a/include/ck_tile/core/tensor/tile_scatter_gather.hpp +++ b/include/ck_tile/core/tensor/tile_scatter_gather.hpp @@ -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 -// CK_TILE_DEVICE auto -// make_tile_window_raw(const TensorView_& tensor_view, -// const WindowLengths_& window_lengths, -// const multi_index& origin, -// const StaticTileDistribution_& tile_distribution, -// number = {}) -// { -// auto w = tile_scatter_gather, -// remove_cvref_t, -// remove_cvref_t, -// NumCoord>{ -// tensor_view, window_lengths, origin, tile_distribution}; -// w.init_raw(); -// return w; -// } - -// template -// CK_TILE_DEVICE void move_tile_window( -// tile_scatter_gather& window, -// const typename tile_scatter_gather::BottomTensorIndex& step) -// { -// window.move(step); -// } - - template CK_TILE_DEVICE constexpr auto make_tile_scatter_gather(const tile_window_with_static_lengths& tile_window, diff --git a/include/ck_tile/core/tensor/tile_window.hpp b/include/ck_tile/core/tensor/tile_window.hpp index 59a5236ceb..2a3c9fdc64 100644 --- a/include/ck_tile/core/tensor/tile_window.hpp +++ b/include/ck_tile/core/tensor/tile_window.hpp @@ -1097,23 +1097,6 @@ make_tile_window_raw(const TensorView_& tensor_view, return w; } -template -CK_TILE_DEVICE void move_tile_window( - tile_window_with_static_distribution& window, - const typename tile_window_with_static_distribution::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 -CK_TILE_DEVICE void move_tile_window( - tile_window_with_static_lengths& window, - const typename tile_window_with_static_lengths::BottomTensorIndex& - step) -{ - window.move(step); -} - } // namespace ck_tile diff --git a/include/ck_tile/core/tensor/tile_window_linear.hpp b/include/ck_tile/core/tensor/tile_window_linear.hpp index 1e24e660f6..6af6813e0c 100644 --- a/include/ck_tile/core/tensor/tile_window_linear.hpp +++ b/include/ck_tile/core/tensor/tile_window_linear.hpp @@ -1200,19 +1200,4 @@ make_tile_window_linear_raw(const TileWindow_& tile_window, LinearBottomDims_{}); } -template -CK_TILE_DEVICE void move_tile_window( - tile_window_linear& - window, - const typename tile_window_linear::BottomTensorIndex& step) -{ - window.move(step); -} - } // namespace ck_tile diff --git a/include/ck_tile/core/tensor/tile_window_utils.hpp b/include/ck_tile/core/tensor/tile_window_utils.hpp index 71a72329f8..a6d4fcde36 100644 --- a/include/ck_tile/core/tensor/tile_window_utils.hpp +++ b/include/ck_tile/core/tensor/tile_window_utils.hpp @@ -18,6 +18,14 @@ #pragma once namespace ck_tile { +template +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 diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp index 632f2fa428..877e58d183 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp @@ -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{}), 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 }); }