Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "8557901d027482d483314701d23e879e68f01bfd"
Commit 53766479 authored by Qianfeng Zhang's avatar Qianfeng Zhang
Browse files

Merge branch 'ck_tile/improve_async_pipeline' of... 

Merge branch 'ck_tile/improve_async_pipeline' of https://github.com/ROCm/composable_kernel into ck_tile/improve_async_pipeline
parents 02b6c6c2 a151fd6d
Hide whitespace changes
Inline Side-by-side
Showing with 35 additions and 23 deletions
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp
View file @ 53766479
...@@ -168,6 +168,8 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -168,6 +168,8 @@ struct BlockFmhaPipelineQRKSVSAsync
static_assert(2 <= k0_loops); static_assert(2 <= k0_loops);
static_assert(2 <= k1_loops); static_assert(2 <= k1_loops);
constexpr bool kPreloadWholeNextIterationK = (kM0 <= 64);
constexpr auto NumKLdsBuffers = Policy::template GetNumKLdsBuffers<Problem>(); constexpr auto NumKLdsBuffers = Policy::template GetNumKLdsBuffers<Problem>();
constexpr auto NumVLdsBuffers = Policy::template GetNumVLdsBuffers<Problem>(); constexpr auto NumVLdsBuffers = Policy::template GetNumVLdsBuffers<Problem>();
...@@ -210,11 +212,10 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -210,11 +212,10 @@ struct BlockFmhaPipelineQRKSVSAsync
using k_tile_type = decltype(load_tile(k_dram_window)); using k_tile_type = decltype(load_tile(k_dram_window));
auto k_tiles = [&]() { auto k_tiles = [&]() {
// for hdim-96 and hdim-160, try to save vgprs if constexpr(kPreloadWholeNextIterationK)
if constexpr(kQKHeaddim < kSubQKHeaddim)
return statically_indexed_array<k_tile_type, 2>{};
else
return statically_indexed_array<k_tile_type, k0_loops>{}; return statically_indexed_array<k_tile_type, k0_loops>{};
else
return statically_indexed_array<k_tile_type, 1>{};
}(); }();
k_tiles[I0] = load_tile(k_dram_window); k_tiles[I0] = load_tile(k_dram_window);
...@@ -330,7 +331,7 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -330,7 +331,7 @@ struct BlockFmhaPipelineQRKSVSAsync
do do
{ {
if constexpr(kQKHeaddim == kSubQKHeaddim) if constexpr(kPreloadWholeNextIterationK)
{ {
if(i_total_loops == 0) // executed by fist iteration if(i_total_loops == 0) // executed by fist iteration
{ {
...@@ -341,7 +342,8 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -341,7 +342,8 @@ struct BlockFmhaPipelineQRKSVSAsync
k_tiles[number<i_k0>{}]); k_tiles[number<i_k0>{}]);
k_tiles[number<i_k0 + 1>{}] = load_tile(k_dram_window); k_tiles[number<i_k0 + 1>{}] = load_tile(k_dram_window);
move_tile_window(k_dram_window, {0, kK0}); if constexpr(i_k0 < k0_loops - 2)
move_tile_window(k_dram_window, {0, kK0});
if constexpr(i_k0 == 0) if constexpr(i_k0 == 0)
clear_tile(s_acc); clear_tile(s_acc);
...@@ -356,7 +358,7 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -356,7 +358,7 @@ struct BlockFmhaPipelineQRKSVSAsync
store_tile(k_lds_windows[number<(k0_loops - 1) % NumKLdsBuffers>{}], store_tile(k_lds_windows[number<(k0_loops - 1) % NumKLdsBuffers>{}],
k_tiles[number<k0_loops - 1>{}]); k_tiles[number<k0_loops - 1>{}]);
move_tile_window(k_dram_window, {kN0, -k0_loops * kK0}); move_tile_window(k_dram_window, {kN0, -(k0_loops - 1) * kK0});
static_for<0, k0_loops, 1>{}([&](auto i_k0) { static_for<0, k0_loops, 1>{}([&](auto i_k0) {
k_tiles[number<i_k0>{}] = load_tile(k_dram_window); k_tiles[number<i_k0>{}] = load_tile(k_dram_window);
...@@ -375,15 +377,13 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -375,15 +377,13 @@ struct BlockFmhaPipelineQRKSVSAsync
} }
else // there is only single iteration else // there is only single iteration
{ {
static_for<0, k0_loops, 1>{}([&](auto i_k0) { static_for<0, k0_loops - 1, 1>{}([&](auto i_k0) {
store_tile(k_lds_windows[number<i_k0 % NumKLdsBuffers>{}], store_tile(k_lds_windows[number<i_k0 % NumKLdsBuffers>{}],
k_tiles[number<i_k0>{}]); k_tiles[number<i_k0>{}]);
if constexpr(i_k0 < k0_loops - 1) k_tiles[number<i_k0 + 1>{}] = load_tile(k_dram_window);
{ if constexpr(i_k0 < k0_loops - 2)
k_tiles[number<i_k0 + 1>{}] = load_tile(k_dram_window);
move_tile_window(k_dram_window, {0, kK0}); move_tile_window(k_dram_window, {0, kK0});
};
if constexpr(i_k0 == 0) if constexpr(i_k0 == 0)
clear_tile(s_acc); clear_tile(s_acc);
...@@ -395,6 +395,14 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -395,6 +395,14 @@ struct BlockFmhaPipelineQRKSVSAsync
k_lds_windows[number<i_k0 % NumKLdsBuffers>{}]); k_lds_windows[number<i_k0 % NumKLdsBuffers>{}]);
}); });
store_tile(k_lds_windows[number<(k0_loops - 1) % NumKLdsBuffers>{}],
k_tiles[number<k0_loops - 1>{}]);
block_sync_lds();
gemm_0(s_acc,
q_tiles[number<k0_loops - 1>{}],
k_lds_windows[number<(k0_loops - 1) % NumKLdsBuffers>{}]);
// move_tile_window(k_dram_window, {0, -k0_loops * kK0}); // move_tile_window(k_dram_window, {0, -k0_loops * kK0});
} }
} }
...@@ -440,21 +448,18 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -440,21 +448,18 @@ struct BlockFmhaPipelineQRKSVSAsync
}; };
}; };
} }
else // for hdim-96, hdim-160 else // only preload one unroll of K for next iteration
{ {
static_for<0, k0_loops, 1>{}([&](auto i_k0) { static_for<0, k0_loops - 1, 1>{}([&](auto i_k0) {
store_tile(k_lds_windows[number<i_k0 % NumKLdsBuffers>{}], store_tile(k_lds_windows[number<i_k0 % NumKLdsBuffers>{}], k_tiles[I0]);
k_tiles[number<i_k0 % 2>{}]);
if constexpr(i_k0 < k0_loops - 1)
{
k_tiles[number<(i_k0 + 1) % 2>{}] = load_tile(k_dram_window);
move_tile_window(k_dram_window, {0, kK0});
};
if constexpr(i_k0 == 0) if constexpr(i_k0 == 0)
clear_tile(s_acc); clear_tile(s_acc);
k_tiles[I0] = load_tile(k_dram_window);
if constexpr(i_k0 < k0_loops - 2)
move_tile_window(k_dram_window, {0, kK0});
block_sync_lds(); block_sync_lds();
// execute current unroll of gemm_0 // execute current unroll of gemm_0
gemm_0(s_acc, gemm_0(s_acc,
...@@ -462,9 +467,16 @@ struct BlockFmhaPipelineQRKSVSAsync ...@@ -462,9 +467,16 @@ struct BlockFmhaPipelineQRKSVSAsync
k_lds_windows[number<i_k0 % NumKLdsBuffers>{}]); k_lds_windows[number<i_k0 % NumKLdsBuffers>{}]);
}); });
store_tile(k_lds_windows[number<(k0_loops - 1) % NumKLdsBuffers>{}], k_tiles[I0]);
block_sync_lds();
gemm_0(s_acc,
q_tiles[number<k0_loops - 1>{}],
k_lds_windows[number<(k0_loops - 1) % NumKLdsBuffers>{}]);
if(i_total_loops < num_total_loop - 1) if(i_total_loops < num_total_loop - 1)
{ {
move_tile_window(k_dram_window, {kN0, -k0_loops * kK0}); move_tile_window(k_dram_window, {kN0, -(k0_loops - 1) * kK0});
k_tiles[I0] = load_tile(k_dram_window); k_tiles[I0] = load_tile(k_dram_window);
move_tile_window(k_dram_window, {0, kK0}); move_tile_window(k_dram_window, {0, kK0});
}; };
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment