ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-08 23:38:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix tail handler bug

Browse Source
This commit is contained in:
lalala-sh
2025-07-17 08:40:35 +00:00
parent fb76450e63
commit 9aa3396a79
4 changed files with 19 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
example/ck_tile/18_flatmm/flatmm_basic.cpp
View File

@@ -80,7 +80,7 @@ float flatmm_calc(const ck_tile::FlatmmHostArgs<>& args, const ck_tile::stream_c
constexpr auto scheduler = FlatmmConfig::Scheduler;
constexpr auto memory_operation = memory_operation_.value;
using CodegenPipelineProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
using CodegenPipelineProblem = ck_tile::FlatmmPipelineProblem<ADataType,
BDataType,
AccDataType,
CodegenFlatmmShape,

3
example/ck_tile/18_flatmm/flatmm_basic.hpp
View File

@@ -83,7 +83,10 @@ struct FlatmmConfig16
template <typename DataType>
struct FlatmmConfig16_950 : public FlatmmConfig16<DataType>
{
static constexpr ck_tile::index_t N_Tile = 256;
static constexpr ck_tile::index_t K_Tile = 256 / sizeof(DataType);
static constexpr ck_tile::index_t K_Warp_Tile = sizeof(DataType) == 2 ? 32 : 128;
static constexpr int kBlockPerCu = 1;
};
template <typename ADataType>

12
include/ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1.hpp
View File

@@ -25,9 +25,19 @@ struct BaseFlatmmPipelineAGmemBGmemCRegV1
return num_loop % 2 == 0 ? TailNumber::Even : TailNumber::Odd;
}
template <typename RunFunction>
CK_TILE_HOST_DEVICE static auto TailHandler(const RunFunction& run_func, bool, TailNumber)
CK_TILE_HOST_DEVICE static auto TailHandler(const RunFunction& run_func, bool, TailNumber tail_num)
{
if (TailNumber::Even == tail_num)
{
return run_func(bool_constant<true>{}, integral_constant<TailNumber, TailNumber::Even>{});
}
else if (TailNumber::Odd == tail_num)
{
return run_func(bool_constant<true>{}, integral_constant<TailNumber, TailNumber::Odd>{});
}
// assert(false);
return run_func(bool_constant<true>{}, integral_constant<TailNumber, TailNumber::Empty>{});
// return run_func(bool_constant<true>{}, integral_constant<TailNumber, TailNumber::Empty>{});
}
};

6
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp
View File

@@ -178,11 +178,13 @@ using GemmPipelineProblem = GemmPipelineProblemBase<ADataType_,
VectorSizeA_,
VectorSizeB_>;
template <typename ADataType_,
typename BDataType_,
typename CDataType_,
typename BlockGemmShape_,
typename Traits_,
GemmPipelineScheduler Scheduler_ = GemmPipelineScheduler::Intrawave,
bool HasHotLoop_ = true,
TailNumber TailNum_ = TailNumber::Full,
typename ComputeDataType_ = ADataType_>
@@ -202,7 +204,7 @@ struct FlatmmPipelineProblem
using CLayout = remove_cvref_t<typename Traits::CLayout>;
static constexpr bool TransposeC = Traits::TransposeC;
static constexpr index_t NumWaveGroups = Traits::NumWaveGroups;
static constexpr bool UseStructuredSparsity = Traits::UseStructuredSparsity;
static constexpr index_t kBlockSize = BlockGemmShape::NumWarps * get_warp_size();
@@ -318,7 +320,7 @@ struct FlatmmPipelineProblem
return kPadM ? 1 : GetAlignmentC();
}
}();
};
};
template <typename ADataType_,
typename BDataType_,