ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-20 21:09:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplify the codes in block_gemm_areg_bsmem_trload_creg

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-11-10 15:27:34 +00:00
parent bd0444f365
commit 303818a851
1 changed files with 15 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
example/ck_tile/18_hstu_attention/block_gemm_areg_bsmem_trload_creg_v2_hack_1.hpp
View File

@@ -146,8 +146,15 @@ struct BlockGemmARegBSmemTrLoadCRegV2Hack_1
__builtin_amdgcn_sched_barrier(0);
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
// read C warp tensor from C block tensor
CWarpTensor c_warp_tensor;
c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
// read A warp tensor from A block tensor
AWarpTensor a_warp_tensor;
@@ -155,22 +162,15 @@ struct BlockGemmARegBSmemTrLoadCRegV2Hack_1
merge_sequences(sequence<mIter, kIter>{}, a_warp_y_index_zeros),
merge_sequences(sequence<1, 1>{}, a_warp_y_lengths));
// read C warp tensor from C block tensor
CWarpTensor c_warp_tensor;
c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
// warp GEMM
WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensors[nIter][kIter]);
// write C warp tensor into C block tensor
c_block_tensor.set_y_sliced_thread_data(
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
c_warp_tensor.get_thread_buffer());
});
// write C warp tensor into C block tensor
c_block_tensor.set_y_sliced_thread_data(
merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
c_warp_tensor.get_thread_buffer());
});
});
}