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

epilogue switched to cshuffle

Browse Source
This commit is contained in:
aska-0096
2025-07-23 03:40:15 +00:00
parent 14e0ab70c6
commit 769fbb62d5
7 changed files with 49 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
include/ck_tile/ops/elementwise/unary_element_wise_operation.hpp
View File

@@ -240,12 +240,13 @@ struct PassThrough
y = type_convert<float>(x);
}
template <>
CK_TILE_HOST_DEVICE void
operator()<ck_tile::bf16_t, ck_tile::fp16_t>(ck_tile::bf16_t& y, const ck_tile::fp16_t& x) const
{
y = type_convert<ck_tile::bf16_t>(x);
}
// template <>
// CK_TILE_HOST_DEVICE void
// operator()<ck_tile::bf16_t, ck_tile::fp16_t>(ck_tile::bf16_t& y, const ck_tile::fp16_t& x)
// const
// {
// y = type_convert<ck_tile::bf16_t>(x);
// }
template <>
CK_TILE_HOST_DEVICE void operator()<float, ck_tile::fp16_t>(float& y,

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

@@ -161,29 +161,23 @@ struct CShuffleEpilogue
* - NumNXdlPerWavePerShuffle: Number of XDL tiles in N dimension processed per wave
*/
static constexpr auto shuffle_tile_tuple = [] {
constexpr index_t elem_per_thread = MPerXdl * NPerXdl / get_warp_size();
if constexpr(elem_per_thread >= GetVectorSizeC())
constexpr index_t memory_friendly_bulk_length = 128 /sizeof(ODataType);
if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
{
return std::make_tuple(1, 1);
}
else
{
constexpr index_t num_xdl_shuffles = GetVectorSizeC() / elem_per_thread;
if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
{
static_assert((kMPerBlock % (MPerXdl * MWave) == 0) &&
(kMPerBlock % num_xdl_shuffles == 0),
"kMPerBlock must be divisible by MPerXdl*MWave and "
"num_xdl_shuffles for CShuffleEpilogue");
return std::make_tuple(min(num_xdl_shuffles, kMPerBlock / (MPerXdl * MWave)), 1);
if constexpr((kNPerBlock % memory_friendly_bulk_length) !=0){
return std::make_tuple(1, 1);
}
else
{
static_assert((kNPerBlock % (NPerXdl * NWave) == 0) &&
(kNPerBlock % num_xdl_shuffles == 0),
"kNPerBlock must be divisible by NPerXdl*NWave and "
"num_xdl_shuffles for CShuffleEpilogue");
return std::make_tuple(1, min(num_xdl_shuffles, kNPerBlock / (NPerXdl * NWave)));
else{
return std::make_tuple(1, max(1, memory_friendly_bulk_length / (NPerXdl * NWave)));
}
}
else{
if constexpr((kMPerBlock % memory_friendly_bulk_length) !=0){
return std::make_tuple(1, 1);
}
else{
return std::make_tuple(max(1, memory_friendly_bulk_length / (MPerXdl * MWave)), 1);
}
}
}();

2
include/ck_tile/ops/fmha/kernel/fmha_fwd_decode_kernel.hpp
View File

@@ -1116,7 +1116,7 @@ struct FmhaFwdDecodeKernel
make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kN1>{}),
{i_m0, i_n1});
EpiloguePipeline{}(o_acc_dram_window, o_acc_tile);
EpiloguePipeline{}(o_acc_dram_window, o_acc_tile, ck_tile::tuple<>{}, smem_ptr);
}
};