ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 18:17:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#5789 (commit 6654ca6)

Browse Source 
[CK][CK_TILE] Revert addional oob check in gemm IsSupported
 function (#5789)

## Motivation

fix ck_tile's oob check.

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
This commit is contained in:
joyeamd
2026-03-26 01:41:35 +00:00
committed by assistant-librarian[bot]
parent 0004a37de5
commit 046d3ac274
2 changed files with 39 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp
View File

@@ -448,23 +448,11 @@ struct UniversalGemmKernel
} }
if(kargs.K % vectorSizeA != 0) if(kargs.K % vectorSizeA != 0)
{ {
const auto remainder = kargs.K % vectorSizeA; if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
constexpr ck_tile::index_t APackedSize =
ck_tile::numeric_traits<ADataType>::PackedSize;
const auto remainder_in_bytes = remainder * sizeof(ADataType) / APackedSize;
// oob can support to dword level
if(remainder_in_bytes % 4 == 0)
{ {
AsTensorIsValid = true; CK_TILE_ERROR("K is not a multiple of vector load size for A tensor!");
}
else
{
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
{
CK_TILE_ERROR("K is not a multiple of vector load size for A tensor!");
}
AsTensorIsValid = false;
} }
AsTensorIsValid = false;
} }
} }
else else
@@ -480,24 +468,11 @@ struct UniversalGemmKernel
} }
if(kargs.M % vectorSizeA != 0) if(kargs.M % vectorSizeA != 0)
{ {
const auto remainder = kargs.M % vectorSizeA; if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
constexpr ck_tile::index_t APackedSize =
ck_tile::numeric_traits<ADataType>::PackedSize;
const auto remainder_in_bytes = remainder * sizeof(ADataType) / APackedSize;
// oob can support to dword level
if(remainder_in_bytes % 4 == 0)
{ {
CK_TILE_ERROR("M is not a multiple of vector load size for A tensor!");
AsTensorIsValid = true;
}
else
{
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
{
CK_TILE_ERROR("M is not a multiple of vector load size for A tensor!");
}
AsTensorIsValid = false;
} }
AsTensorIsValid = false;
} }
} }
}); });
@@ -520,58 +495,33 @@ struct UniversalGemmKernel
} }
if(kargs.N % vectorSizeB != 0) if(kargs.N % vectorSizeB != 0)
{ {
const auto remainder = kargs.N % vectorSizeB; if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
constexpr ck_tile::index_t BPackedSize =
ck_tile::numeric_traits<BDataType>::PackedSize;
const auto remainder_in_bytes = remainder * sizeof(BDataType) / BPackedSize;
// oob can support to dword level
if(remainder_in_bytes % 4 == 0)
{ {
BsTensorIsValid = true; CK_TILE_ERROR("N is not a multiple of vector load size for B tensor!");
}
else
{
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
{
CK_TILE_ERROR("N is not a multiple of vector load size for B tensor!");
}
BsTensorIsValid = false;
} }
BsTensorIsValid = false;
} }
else }
else
{
if(kargs.K % (TilePartitioner::KPerBlock * kargs.k_batch) != 0 &&
GemmPipeline::kPadK == false)
{ {
if(kargs.K % (TilePartitioner::KPerBlock * kargs.k_batch) != 0 && if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
GemmPipeline::kPadK == false)
{ {
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) CK_TILE_ERROR(
{ "Can't support K that is not a multiple of k_batch * KPerBlock "
CK_TILE_ERROR( "without padding!");
"Can't support K that is not a multiple of k_batch * KPerBlock "
"without padding!");
}
BsTensorIsValid = false;
} }
if(kargs.K % vectorSizeB != 0) BsTensorIsValid = false;
}
if(kargs.K % vectorSizeB != 0)
{
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
{ {
const auto remainder = kargs.K % vectorSizeB; CK_TILE_ERROR("K is not a multiple of vector load size for B tensor!");
constexpr ck_tile::index_t BPackedSize =
ck_tile::numeric_traits<BDataType>::PackedSize;
const auto remainder_in_bytes = remainder * sizeof(BDataType) / BPackedSize;
// oob can support to dword level
if(remainder_in_bytes % 4 == 0)
{
BsTensorIsValid = true;
}
else
{
if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
{
CK_TILE_ERROR(
"K is not a multiple of vector load size for B tensor!");
}
BsTensorIsValid = false;
}
} }
BsTensorIsValid = false;
} }
} }
}); });

31
test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc
View File

@@ -31,14 +31,7 @@ TYPED_TEST(TEST_SUITE_NAME, SmallM)
if constexpr(std::is_same_v<typename TestFixture::ALayout, if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>) ck_tile::tensor_layout::gemm::ColumnMajor>)
{ {
if(M * sizeof(typename TestFixture::ADataType) % 4 == 0) // oob fit dword EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
{
this->Run(M, N, K);
}
else
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
} }
else else
{ {
@@ -91,14 +84,7 @@ TYPED_TEST(TEST_SUITE_NAME, MidLargeM)
} }
else else
{ {
if(M * sizeof(typename TestFixture::ADataType) % 4 == 0) // oob fit dword EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
{
this->Run(M, N, K);
}
else
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
} }
} }
else else
@@ -120,7 +106,18 @@ TYPED_TEST(TEST_SUITE_NAME, PaddK)
for(int M : Ms) for(int M : Ms)
{ {
this->Run(M, N, K); if constexpr(std::is_same_v<typename TestFixture::BDataType, ck_tile::pk_int4_t>)
{
#if defined(ARCH_GFX12) || defined(ARCH_GFX11)
this->Run(M, N, K);
#else
EXPECT_THROW(this->Run(M, N, K), std::runtime_error);
#endif
}
else
{
this->Run(M, N, K);
}
} }
} }