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

[CK_Tile] Fix gemm kernel for 4,64,16 and 64,4,16 warp tile sizes (#2262)

Browse Source 
* debugging issue

* debugging issue

* debugging

* debugging

* reverting debugging code

* clang formatted

* updating default_config.json

* fix ci failure

* clang formatted

[ROCm/composable_kernel commit: 59a85cb4bc]
This commit is contained in:
Khushbu Agarwal
2025-06-03 20:16:10 -07:00
committed by GitHub
parent 945a691ce5
commit 7f9758a186
6 changed files with 46 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
View File

@@ -92,7 +92,20 @@ struct CShuffleEpilogue
CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeC()
{
constexpr index_t MaxVectorStoreSize = 16;
return MaxVectorStoreSize / sizeof(ODataType);
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
{
return std::min(static_cast<int>(kNPerIteration),
static_cast<int>(MaxVectorStoreSize / sizeof(ODataType)));
}
else if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::ColumnMajor>)
{
return std::min(static_cast<int>(kMPerIteration),
static_cast<int>(MaxVectorStoreSize / sizeof(ODataType)));
}
else
{
static_assert(false, "Unsupported CLayout!");
}
}
template <typename Problem>

8
include/ck_tile/ops/epilogue/default_2d_epilogue.hpp
View File

@@ -149,7 +149,9 @@ struct DefaultGemm2DEpilogue : public Default2DEpilogue<Problem_, Policy_>
else
{
// In this case each thread has just a single item in Ndim
return WG::WarpGemmAttribute::Impl::kCNLane / WG::kN;
return (WG::WarpGemmAttribute::Impl::kCNLane *
WG::WarpGemmAttribute::Impl::kBNBlock) /
WG::kN;
}
}
// M is contiguous dimension
@@ -158,7 +160,9 @@ struct DefaultGemm2DEpilogue : public Default2DEpilogue<Problem_, Policy_>
if constexpr(isCTransposed)
{
// In this case each thread has just a single item in Mdim
return WG::WarpGemmAttribute::Impl::kCNLane / WG::kN;
return (WG::WarpGemmAttribute::Impl::kCNLane *
WG::WarpGemmAttribute::Impl::kAMBlock) /
WG::kN;
}
else
{