ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-01 12:11:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Ck tile gemm padding dim (#1516)

Browse Source 
* Support the N dimension padding

* Finished the padding feature for different dimension of K
This commit is contained in:
Thomas Ning
2024-09-18 11:32:29 -07:00
committed by GitHub
parent e84adec3ba
commit 694c300145
4 changed files with 33 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp
View File

@@ -123,14 +123,26 @@ struct GemmKernel
}
}();
auto ABlockWindow = make_tile_window(
auto a_pad_view = pad_tensor_view(
a_tensor_view,
make_tuple(number<TilePartitioner::kM>{}, number<TilePartitioner::kK>{}),
sequence < 0,
GemmPipeline::kPadA ? 1 : 0 > {});
auto ABlockWindow = make_tile_window(
a_pad_view,
make_tuple(number<TilePartitioner::kM>{}, number<TilePartitioner::kK>{}),
{i_m, 0});
auto BBlockWindow = make_tile_window(
auto b_pad_view = pad_tensor_view(
b_tensor_view,
make_tuple(number<TilePartitioner::kN>{}, number<TilePartitioner::kK>{}),
sequence < 0,
GemmPipeline::kPadB ? 1 : 0 > {});
auto BBlockWindow = make_tile_window(
b_pad_view,
make_tuple(number<TilePartitioner::kN>{}, number<TilePartitioner::kK>{}),
{i_n, 0});
// allocate LDS
@@ -163,12 +175,16 @@ struct GemmKernel
}
}();
auto CBlockWindow = make_tile_window(
auto c_pad_view = pad_tensor_view(
c_tensor_view,
make_tuple(number<TilePartitioner::kM>{}, number<TilePartitioner::kN>{}),
sequence < 0,
GemmPipeline::kPadC ? 1 : 0 > {});
auto CBlockWindow_pad = make_tile_window(
c_pad_view,
make_tuple(number<TilePartitioner::kM>{}, number<TilePartitioner::kN>{}),
{i_m, i_n});
// epilogue.
EpiloguePipeline{}(CBlockWindow, acc);
EpiloguePipeline{}(CBlockWindow_pad, acc);
}
};