mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-04-20 06:49:15 +00:00
[rocm-libraries] ROCm/rocm-libraries#4274 (commit 7c380df)
Add padding to cshuffle epilogue to avoid bank conflict (#4274) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Proposed changes Added the padding to CShuffle Epilogue to avoid the bank conflicts of 64. Synced up and learned from the internal repo. ## Checklist Please put an `x` into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask. - [ ] I have added tests relevant to the introduced functionality, and the unit tests are passing locally - [ ] I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, **IF** the test takes more than 30 seconds to run. - [ ] I have added inline documentation which enables the maintainers with understanding the motivation - [ ] I have removed the stale documentation which is no longer relevant after this pull request - [ ] (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request - [ ] I have run `clang-format` on all changed files - [ ] Any dependent changes have been merged ## Discussion If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered
This commit is contained in:
Thomas Ning
committed by
assistant-librarian[bot]
assistant-librarian[bot]
parent
6d6ee8f023
commit
807efa703a
@@ -302,19 +302,118 @@ struct CShuffleEpilogue
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto MakeLdsBlockDescriptor()
|
||||
{
|
||||
constexpr auto DataTypeSize = sizeof(ODataType);
|
||||
constexpr index_t VectorLen = GetVectorSizeC();
|
||||
constexpr index_t banks = get_n_lds_banks();
|
||||
|
||||
constexpr index_t BytesPerBank = 4;
|
||||
|
||||
// N is contiguous dimension
|
||||
if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
|
||||
{
|
||||
return make_naive_tensor_descriptor(
|
||||
make_tuple(number<MPerIterationShuffle>{}, number<NPerIterationShuffle>{}),
|
||||
make_tuple(number<NPerIterationShuffle>{}, number<1>{}));
|
||||
constexpr index_t MLdsLayerRequired =
|
||||
banks * BytesPerBank / NPerIterationShuffle / DataTypeSize;
|
||||
constexpr auto MLdsLayer = max(1, MLdsLayerRequired);
|
||||
|
||||
constexpr index_t BaseStrideElems = NPerIterationShuffle * MLdsLayer;
|
||||
static_assert((BaseStrideElems * DataTypeSize) % BytesPerBank == 0,
|
||||
"LDS row stride must be 4B-aligned for bank-word padding logic");
|
||||
// calculate how many elements to pad to avoid bank conflict
|
||||
#if defined(__gfx950__)
|
||||
constexpr index_t ElemsPer4B = BytesPerBank / ck_tile::gcd(BytesPerBank, DataTypeSize);
|
||||
constexpr auto ToWords = [](index_t elems) constexpr {
|
||||
return (elems * DataTypeSize) / BytesPerBank;
|
||||
};
|
||||
constexpr index_t BaseWords = ToWords(BaseStrideElems);
|
||||
constexpr index_t PadWords = ((BaseWords % 2) == 0) ? 1 : 0;
|
||||
constexpr auto PaddingAmount = PadWords * ElemsPer4B;
|
||||
#else
|
||||
constexpr auto PaddingAmount = 0;
|
||||
#endif
|
||||
|
||||
constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(number<MPerIterationShuffle / MLdsLayer>{},
|
||||
number<NPerIterationShuffle / VectorLen * MLdsLayer>{},
|
||||
number<VectorLen>{}),
|
||||
make_tuple(number<NPerIterationShuffle * MLdsLayer + PaddingAmount>{},
|
||||
number<VectorLen>{},
|
||||
number<1>{}),
|
||||
number<VectorLen>{},
|
||||
number<1>{});
|
||||
|
||||
constexpr auto lds_block_desc_1 = transform_tensor_descriptor(
|
||||
lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(number<MPerIterationShuffle / MLdsLayer>{}),
|
||||
make_unmerge_transform(make_tuple(
|
||||
number<MLdsLayer>{}, number<NPerIterationShuffle / VectorLen>{})),
|
||||
make_pass_through_transform(number<VectorLen>{})),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3>{}));
|
||||
|
||||
constexpr auto lds_block_desc = transform_tensor_descriptor(
|
||||
lds_block_desc_1,
|
||||
make_tuple(make_merge_transform_v3_division_mod(make_tuple(
|
||||
number<MPerIterationShuffle / MLdsLayer>{}, number<MLdsLayer>{})),
|
||||
make_merge_transform_v3_division_mod(make_tuple(
|
||||
number<NPerIterationShuffle / VectorLen>{}, number<VectorLen>{}))),
|
||||
make_tuple(sequence<0, 1>{}, sequence<2, 3>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
|
||||
return lds_block_desc;
|
||||
}
|
||||
// M is contiguous dimension
|
||||
else if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::ColumnMajor>)
|
||||
{
|
||||
return make_naive_tensor_descriptor(
|
||||
make_tuple(number<MPerIterationShuffle>{}, number<NPerIterationShuffle>{}),
|
||||
make_tuple(number<1>{}, number<MPerIterationShuffle>{}));
|
||||
constexpr index_t NLdsLayerRequired =
|
||||
get_n_lds_banks() * BytesPerBank / MPerIterationShuffle / DataTypeSize;
|
||||
constexpr auto NLdsLayer = max(1, NLdsLayerRequired);
|
||||
|
||||
constexpr index_t BaseStrideElems = MPerIterationShuffle * NLdsLayer;
|
||||
|
||||
static_assert((BaseStrideElems * DataTypeSize) % BytesPerBank == 0,
|
||||
"LDS row stride must be 4B-aligned for bank-word padding logic");
|
||||
|
||||
#if defined(__gfx950__)
|
||||
constexpr index_t ElemsPer4B = BytesPerBank / ck_tile::gcd(BytesPerBank, DataTypeSize);
|
||||
constexpr auto ToWords = [](index_t elems) constexpr {
|
||||
return (elems * DataTypeSize) / BytesPerBank;
|
||||
};
|
||||
constexpr index_t BaseWords = ToWords(BaseStrideElems);
|
||||
constexpr index_t PadWords = ((BaseWords % 2) == 0) ? 1 : 0;
|
||||
constexpr auto PaddingAmount = PadWords * ElemsPer4B;
|
||||
#else
|
||||
constexpr auto PaddingAmount = 0;
|
||||
#endif
|
||||
|
||||
constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(number<NPerIterationShuffle / NLdsLayer>{},
|
||||
number<MPerIterationShuffle / VectorLen * NLdsLayer>{},
|
||||
number<VectorLen>{}),
|
||||
make_tuple(number<MPerIterationShuffle * NLdsLayer + PaddingAmount>{},
|
||||
number<VectorLen>{},
|
||||
number<1>{}),
|
||||
number<VectorLen>{},
|
||||
number<1>{});
|
||||
|
||||
constexpr auto lds_block_desc_1 = transform_tensor_descriptor(
|
||||
lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(number<NPerIterationShuffle / NLdsLayer>{}),
|
||||
make_unmerge_transform(make_tuple(
|
||||
number<NLdsLayer>{}, number<MPerIterationShuffle / VectorLen>{})),
|
||||
make_pass_through_transform(number<VectorLen>{})),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3>{}));
|
||||
|
||||
constexpr auto lds_block_desc = transform_tensor_descriptor(
|
||||
lds_block_desc_1,
|
||||
make_tuple(make_merge_transform_v3_division_mod(make_tuple(
|
||||
number<NPerIterationShuffle / NLdsLayer>{}, number<NLdsLayer>{})),
|
||||
make_merge_transform_v3_division_mod(make_tuple(
|
||||
number<MPerIterationShuffle / VectorLen>{}, number<VectorLen>{}))),
|
||||
make_tuple(sequence<0, 1>{}, sequence<2, 3>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
|
||||
return lds_block_desc;
|
||||
}
|
||||
else
|
||||
{
|
||||
|
||||
@@ -40,7 +40,8 @@ struct BlockSoftmax2D
|
||||
#endif
|
||||
|
||||
// compute row max
|
||||
auto reduce_row_max = BlockReduce2D<decltype(x)>{x, -numeric<DataType>::infinity()};
|
||||
using X = remove_cvref_t<decltype(x)>;
|
||||
BlockReduce2D<X> reduce_row_max{x, -numeric<DataType>::infinity()};
|
||||
#if _BLOCK_SOFTMAX_USE_UNPACK2
|
||||
auto row_max = reduce_row_max(f_max3, f_max, sequence<1, 2>{});
|
||||
#else
|
||||
@@ -52,7 +53,8 @@ struct BlockSoftmax2D
|
||||
});
|
||||
|
||||
// compute row sum
|
||||
auto reduce_row_sum = BlockReduce2D<decltype(y)>{y, DataType{0}};
|
||||
using Y = remove_cvref_t<decltype(y)>;
|
||||
BlockReduce2D<Y> reduce_row_sum{y, DataType{0}};
|
||||
#if _BLOCK_SOFTMAX_USE_UNPACK2
|
||||
auto row_sum = reduce_row_sum(f_sum3, f_sum, sequence<1, 2>{});
|
||||
#else
|
||||
|
||||
Reference in New Issue
Block a user