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

Epilogue fixes.

Browse Source
This commit is contained in:
Ville Pietilä
2025-09-22 15:38:02 +00:00
parent d7da3d5089
commit 29e3112b9b
2 changed files with 116 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -257,7 +257,14 @@ struct CShuffleEpilogue
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
{
return MPerIterationShuffle * NPerIterationShuffle * sizeof(ODataType);
if constexpr(NumGroupsToMerge > 1)
{
return kMPerBlock * kNPerBlock * sizeof(ODataType);
}
else
{
return MPerIterationShuffle * NPerIterationShuffle * sizeof(ODataType);
}
}
template <typename ODramWindow, typename OAccTile, typename DsDramWindows>
@@ -267,17 +274,17 @@ struct CShuffleEpilogue
void* p_smem)
{
if constexpr (NumGroupsToMerge > 1)
{
//if constexpr (NumGroupsToMerge > 1)
//{
// When NumGroupsToMerge > 1, we want to write out only the diagonal blocks.
// Hence, we configure the shuffle such that it iterates one merge group block at a time.
return merged_op(out_dram_window, o_acc_tile, ds_dram_windows, p_smem);
}
else
{
//}
//else
//{
// When NumGroupsToMerge == 1, we want to write out all the blocks.
return unmerged_op(out_dram_window, o_acc_tile, ds_dram_windows, p_smem);
}
// return unmerged_op(out_dram_window, o_acc_tile, ds_dram_windows, p_smem);
//}
}
template <typename ODramWindow, typename OAccTile, typename DsDramWindows>
@@ -289,7 +296,11 @@ struct CShuffleEpilogue
constexpr auto LdsTileDistr = make_static_tile_distribution(MakeLdsDistributionEncode());
auto lds_tile = make_static_distributed_tensor<AccDataType>(LdsTileDistr);
constexpr auto lds_block_desc = MakeLdsBlockDescriptor<Problem>();
//constexpr auto lds_block_desc = MakeLdsBlockDescriptor<Problem>();
constexpr auto lds_block_desc = make_naive_tensor_descriptor(
make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}),
make_tuple(number<kNPerBlock>{}, number<1>{})); // Row-major layout
auto o_lds_block = make_tensor_view<address_space_enum::lds>(
static_cast<ODataType*>(p_smem), lds_block_desc);
@@ -314,21 +325,6 @@ struct CShuffleEpilogue
static_assert(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>,
"Currently, the CShuffle Epilogue only supports the Row Major Output layout");
using TileEncodingPattern = tile_distribution_encoding_pattern_2d<kBlockSize,
MPerIterationShuffle,
NPerIterationShuffle,
GetVectorSizeC(),
tile_distribution_pattern::sparse_row,
Problem::kNumWaveGroups>;
constexpr auto dram_tile_distribution =
TileEncodingPattern::make_2d_static_tile_distribution();
auto d_dram_windows = generate_tuple(
[&](auto idx) {
return make_tile_window(ds_dram_windows[idx], dram_tile_distribution);
},
number<NumDTensor>{});
constexpr auto c_warp_y_lengths =
to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t<CWarpDstr::NDimY, 0>{};
@@ -362,27 +358,46 @@ struct CShuffleEpilogue
});
block_sync_lds();
// Copy diagonal block from LDS to global memory.
// Set-up LDS to global memory copy.
// We copy only the diagonal blocks from LDS to global memory.
// Hence, we must configure the tile distrinbution and SFCs to work
// on the group size blocks.
constexpr index_t MPerGroup = kMPerBlock / NumGroupsToMerge;
constexpr index_t NPerGroup = kNPerBlock / NumGroupsToMerge;
auto out_lds_window = make_tile_window(
o_lds_block,
make_tuple(number<MPerGroup>{}, number<NPerGroup>{}),
{0, 0});
using SFC_lds = space_filling_curve<sequence<kMPerBlock, kNPerBlock>,
sequence<0, 1>,
sequence<MPerGroup, NPerGroup>>;
// TODO: Within the subtile, we have column-major data layout.
constexpr auto dram_tile_encoding = tile_distribution_encoding<sequence<>,
tuple<sequence<MPerGroup, 1>, sequence<NPerGroup,1>>,
tuple<sequence<1>, sequence<2>>,
tuple<sequence<0>, sequence<0>>,
sequence<1, 2>,
sequence<1, 1>>{};
using SFC_dram = space_filling_curve<sequence<kMPerBlock, kNPerBlock / NumGroupsToMerge>,
sequence<0, 1>,
sequence<MPerGroup, NPerGroup>>;
constexpr auto dram_tile_distribution = make_static_tile_distribution(dram_tile_encoding);
auto d_dram_windows = generate_tuple(
[&](auto idx) {
return make_tile_window(ds_dram_windows[idx], dram_tile_distribution);
},
number<NumDTensor>{});
static_for<0, NumGroupsToMerge, 1>{}
(
[&](auto group)
{
auto c_out_tensor = load_tile(make_tile_window(out_lds_window, dram_tile_distribution));
// Create LDS window at the correct diagonal position for this group
constexpr auto lds_start_m = group * number<MPerGroup>{};
constexpr auto lds_start_n = group * number<NPerGroup>{};
auto current_lds_window = make_tile_window(
o_lds_block,
make_tuple(number<MPerGroup>{}, number<NPerGroup>{}),
{lds_start_m, lds_start_n},
dram_tile_distribution);
block_sync_lds();
auto c_out_tensor = load_tile(current_lds_window);
//make_tile_window(current_lds_window, dram_tile_distribution));
const auto ds_tensor = generate_tuple(
[&](auto idx) { return load_tile(d_dram_windows[idx]); }, number<NumDTensor>{});
@@ -405,21 +420,13 @@ struct CShuffleEpilogue
if constexpr(group != NumGroupsToMerge - 1)
{
constexpr auto step_m = number<MPerGroup>{};
// Move the global memory window
constexpr auto step_dram = SFC_dram::get_forward_step(group);
move_tile_window(out_dram_window, {step_dram.at(number<0>{}), step_dram.at(number<1>{})});
move_tile_window(out_dram_window, {step_m, 0});
static_for<0, NumDTensor, 1>{}([&](auto idx) {
move_tile_window(d_dram_windows[idx],
{step_dram.at(number<0>{}), step_dram.at(number<1>{})});
move_tile_window(d_dram_windows[idx], {step_m, 0});
});
// Move the LDS window
constexpr auto iAccess = number<group * NumGroupsToMerge + group>{};
constexpr auto next_iAccess = number<(group+1) * NumGroupsToMerge + (group+1)>{};
constexpr auto step_lds = SFC_lds::get_step_between(iAccess, next_iAccess);
move_tile_window(out_lds_window, {step_lds.at(number<0>{}), step_lds.at(number<1>{})});
}
}
);

116
include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp
View File

@@ -418,17 +418,18 @@ struct TransformConvBwdWeightToGemm
{
// NWGK
const index_t NDoHoWoStride = G_ * K_;
const index_t GStride = K_;
constexpr auto KStride = I1;
if constexpr (NumGroupsToMerge > 1)
{
const index_t KStride = G_;
constexpr auto GStride = I1;
return make_naive_tensor_descriptor(
make_tuple(K_, NumGroupsToMerge, N_ * Wo_),
make_tuple(KStride, GStride, NDoHoWoStride));
make_tuple(NumGroupsToMerge, K_, N_ * Wo_),
make_tuple(GStride, KStride, NDoHoWoStride));
}
else
{
constexpr auto KStride = I1;
return make_naive_tensor_descriptor(
make_tuple(K_, N_ * Wo_),
make_tuple(KStride, NDoHoWoStride));
@@ -441,17 +442,18 @@ struct TransformConvBwdWeightToGemm
// NWGC
const index_t NStride = Wi_ * G_ * C_;
const index_t WiStride = G_ * C_;
const index_t GStride = C_;
constexpr auto CStride = I1;
if constexpr (NumGroupsToMerge > 1)
{
const index_t CStride = G_;
constexpr auto GStride = I1;
return make_naive_tensor_descriptor(
make_tuple(N_, Wi_, NumGroupsToMerge, C_),
make_tuple(N_, Wi_, C_, NumGroupsToMerge),
make_tuple(NStride, WiStride, GStride, CStride));
}
else
{
constexpr auto CStride = I1;
return make_naive_tensor_descriptor(
make_tuple(N_, Wi_, C_),
make_tuple(NStride, WiStride, CStride));
@@ -488,17 +490,18 @@ struct TransformConvBwdWeightToGemm
{
// NHWGK
const index_t NDoHoWoStride = G_ * K_;
const index_t GStride = K_;
constexpr auto KStride = I1;
if constexpr (NumGroupsToMerge > 1)
{
const index_t KStride = G_;
constexpr auto GStride = I1;
return make_naive_tensor_descriptor(
make_tuple(K_, NumGroupsToMerge, N_ * Ho_ * Wo_),
make_tuple(KStride, GStride, NDoHoWoStride));
make_tuple(NumGroupsToMerge, K_, N_ * Ho_ * Wo_),
make_tuple(GStride, KStride, NDoHoWoStride));
}
else
{
constexpr auto KStride = I1;
return make_naive_tensor_descriptor(
make_tuple(K_, N_ * Ho_ * Wo_),
make_tuple(KStride, NDoHoWoStride));
@@ -512,17 +515,18 @@ struct TransformConvBwdWeightToGemm
const index_t NStride = Hi_ * Wi_ * G_ * C_;
const index_t HiStride = Wi_ * G_ * C_;
const index_t WiStride = G_ * C_;
const index_t GStride = C_;
constexpr auto CStride = I1;
if constexpr (NumGroupsToMerge > 1)
{
const index_t CStride = G_;
constexpr auto GStride = I1;
return make_naive_tensor_descriptor(
make_tuple(N_, Hi_, Wi_, NumGroupsToMerge, C_),
make_tuple(NStride, HiStride, WiStride, GStride, CStride));
make_tuple(N_, Hi_, Wi_, C_, NumGroupsToMerge),
make_tuple(NStride, HiStride, WiStride, CStride, GStride));
}
else
{
constexpr auto CStride = I1;
return make_naive_tensor_descriptor(
make_tuple(N_, Hi_, Wi_, C_),
make_tuple(NStride, HiStride, WiStride, CStride));
@@ -559,17 +563,18 @@ struct TransformConvBwdWeightToGemm
{
// NDHWGK
const index_t NDoHoWoStride = G_ * K_;
const index_t GStride = K_;
constexpr auto KStride = I1;
if constexpr (NumGroupsToMerge > 1)
{
const index_t KStride = G_;
constexpr auto GStride = I1;
return make_naive_tensor_descriptor(
make_tuple(K_, NumGroupsToMerge, N_ * Do_ * Ho_ * Wo_),
make_tuple(KStride, GStride, NDoHoWoStride));
make_tuple(NumGroupsToMerge, K_, N_ * Do_ * Ho_ * Wo_),
make_tuple(GStride, KStride, NDoHoWoStride));
}
else
{
constexpr auto KStride = I1;
return make_naive_tensor_descriptor(
make_tuple(K_, N_ * Do_ * Ho_ * Wo_),
make_tuple(KStride, NDoHoWoStride));
@@ -583,17 +588,18 @@ struct TransformConvBwdWeightToGemm
const index_t DiStride = Hi_ * Wi_ * G_ * C_;
const index_t HiStride = Wi_ * G_ * C_;
const index_t WiStride = G_ * C_;
const index_t GStride = C_;
constexpr auto CStride = I1;
if constexpr (NumGroupsToMerge > 1)
{
constexpr auto GStride = I1;
const index_t CStride = G_;
return make_naive_tensor_descriptor(
make_tuple(N_, Di_, Hi_, Wi_, NumGroupsToMerge, C_),
make_tuple(NStride, DiStride, HiStride, WiStride, GStride, CStride));
make_tuple(N_, Di_, Hi_, Wi_, C_, NumGroupsToMerge),
make_tuple(NStride, DiStride, HiStride, WiStride, CStride, GStride));
}
else
{
constexpr auto CStride = I1;
return make_naive_tensor_descriptor(
make_tuple(N_, Di_, Hi_, Wi_, C_),
make_tuple(NStride, DiStride, HiStride, WiStride, CStride));
@@ -637,12 +643,12 @@ struct TransformConvBwdWeightToGemm
{
// Output tensor transformation
// [0, 1, 2] -> [0, 1]
// [K, Gm, (N*Wo)] -> [(K*Gm), (N*Wo)]
// [Gm, K, (N*Wo)] -> [(Gm*K), (N*Wo)]
const auto out_gemm_m_gemm_k_grid_desc =
transform_tensor_descriptor(
out_grid_desc,
make_tuple(
make_merge_transform(make_tuple(K_, NumGroupsToMerge)),
make_merge_transform(make_tuple(NumGroupsToMerge, K_)),
make_pass_through_transform(N_ * Wo_)),
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -705,38 +711,38 @@ struct TransformConvBwdWeightToGemm
// Input tensor transformation, part 1.
// [N, Wi, Gm, C] -> [N, (Wi + InLeftPadW + InRightPadW), Gm, C] = [N, Wip, Gm, C]
const auto in_n_wip_gm_c_grid_desc = transform_tensor_descriptor(
const auto in_n_wip_c_gm_grid_desc = transform_tensor_descriptor(
in_grid_desc,
make_tuple(
make_pass_through_transform(N_),
make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
make_pass_through_transform(NumGroupsToMerge),
make_pass_through_transform(C_)),
make_pass_through_transform(C_),
make_pass_through_transform(NumGroupsToMerge)),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}));
// Input tensor transformation, part 2.
// [N, Wip, Gm, C] -> [N, X, Wo, Gm, C]
// [N, Wip, C, Gm] -> [N, X, Wo, C, Gm]
const auto in_n_x_wo_gm_c_grid_desc = transform_tensor_descriptor(
in_n_wip_gm_c_grid_desc,
in_n_wip_c_gm_grid_desc,
make_tuple(
make_pass_through_transform(N_),
make_embed_transform(
make_tuple(X_, Wo_),
make_tuple(ConvDilationW_, ConvStrideW_)),
make_pass_through_transform(NumGroupsToMerge),
make_pass_through_transform(C_)),
make_pass_through_transform(C_),
make_pass_through_transform(NumGroupsToMerge)),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}),
make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3>{}, sequence<4>{}));
// Input tensor transformation, part 3.
// [0, 1, 2, 3, 4] -> [0, 1]
// [N, X, Wo, Gm, C] -> [(X*Gm*C), (N*Wo)]
// [0, 1, 2, 3, 4] -> [0, 1]
// [N, X, Wo, C, Gm] -> [(Gm*X*C), (N*Wo)]
const auto in_gemm_n_gemm_k_grid_desc =
transform_tensor_descriptor(
in_n_x_wo_gm_c_grid_desc,
make_tuple(
make_merge_transform(make_tuple(X_, NumGroupsToMerge, C_)),
make_merge_transform(make_tuple(X_, C_, NumGroupsToMerge)),
make_merge_transform(make_tuple(N_, Wo_))),
make_tuple(sequence<1, 3, 4>{}, sequence<0, 2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -788,12 +794,12 @@ struct TransformConvBwdWeightToGemm
{
// Output tensor transformation
// [0, 1, 2] -> [0, 1]
// [K, Gm, (N*Ho*Wo)] -> [(K*Gm), (N*Ho*Wo)]
// [Gm, K, (N*Ho*Wo)] -> [(K*Gm), (N*Ho*Wo)]
const auto out_gemm_m_gemm_k_grid_desc =
transform_tensor_descriptor(
out_grid_desc,
make_tuple(
make_merge_transform(make_tuple(K_, NumGroupsToMerge)),
make_merge_transform(make_tuple(NumGroupsToMerge, K_)),
make_pass_through_transform(N_ * Ho_ * Wo_)),
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -855,22 +861,22 @@ struct TransformConvBwdWeightToGemm
make_tuple(sequence<0>{}, sequence<1>{}));
// Input tensor transformation, part 1.
// [N, Hi, Wi, Gm, C] -> [N, Hip, Wip, Gm, C]
const auto in_n_hip_wip_gm_c_grid_desc = transform_tensor_descriptor(
// [N, Hi, Wi, C, Gm] -> [N, Hip, Wip, C, Gm]
const auto in_n_hip_wip_c_gm_grid_desc = transform_tensor_descriptor(
in_grid_desc,
make_tuple(
make_pass_through_transform(N_),
make_pad_transform(Wi_, InLeftPadH_, InRightPadH_),
make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
make_pass_through_transform(NumGroupsToMerge),
make_pass_through_transform(C_)),
make_pass_through_transform(C_),
make_pass_through_transform(NumGroupsToMerge)),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}, sequence<4>{}),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}, sequence<4>{}));
// Input tensor transformation, part 2.
// [N, Hip, Wip, Gm, C] -> [N, (Y, Wo), (X, Wo), Gm, C]
const auto in_n_y_ho_x_wo_gm_c_grid_desc = transform_tensor_descriptor(
in_n_hip_wip_gm_c_grid_desc,
// [N, Hip, Wip, C, Gm] -> [N, (Y, Wo), (X, Wo), C, Gm]
const auto in_n_y_ho_x_wo_c_gm_grid_desc = transform_tensor_descriptor(
in_n_hip_wip_c_gm_grid_desc,
make_tuple(
make_pass_through_transform(N_),
make_embed_transform(
@@ -879,19 +885,19 @@ struct TransformConvBwdWeightToGemm
make_embed_transform(
make_tuple(X_, Wo_),
make_tuple(ConvDilationW_, ConvStrideW_)),
make_pass_through_transform(NumGroupsToMerge),
make_pass_through_transform(C_)),
make_pass_through_transform(C_),
make_pass_through_transform(NumGroupsToMerge)),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}, sequence<4>{}),
make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3, 4>{}, sequence<5>{}, sequence<6>{}));
// Input tensor transformation, part 3.
// [0, 1, 2, 3, 4 5 6] -> [0, 1]
// [N, Y, Ho, X, Wo, Gm, C] -> [(Y*X*Gm*C), (N*Ho*Wo)]
// [0, 1, 2, 3, 4 5 6] -> [0, 1]
// [N, Y, Ho, X, Wo, C, Gm] -> [(Gm*Y*X*C), (N*Ho*Wo)]
const auto in_gemm_n_gemm_k_grid_desc =
transform_tensor_descriptor(
in_n_y_ho_x_wo_gm_c_grid_desc,
in_n_y_ho_x_wo_c_gm_grid_desc,
make_tuple(
make_merge_transform(make_tuple(Y_, X_, NumGroupsToMerge, C_)),
make_merge_transform(make_tuple(Y_, X_, C_, NumGroupsToMerge)),
make_merge_transform(make_tuple(N_, Ho_, Wo_))),
make_tuple(sequence<1, 3, 5, 6>{}, sequence<0, 2, 4>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -942,12 +948,12 @@ struct TransformConvBwdWeightToGemm
{
// Output tensor transformation
// [0, 1, 2] -> [0, 1]
// [K, Gm, (N*Do*Ho*Wo)] -> [(K*Gm), (N*Do*Ho*Wo)]
// [Gm, K, (N*Do*Ho*Wo)] -> [(K*Gm), (N*Do*Ho*Wo)]
const auto out_gemm_m_gemm_k_grid_desc =
transform_tensor_descriptor(
out_grid_desc,
make_tuple(
make_merge_transform(make_tuple(K_, NumGroupsToMerge)),
make_merge_transform(make_tuple(NumGroupsToMerge, K_)),
make_pass_through_transform(N_ * Do_ * Ho_ * Wo_)),
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));