From 39eaa5015dc575bbd11a91d2e867ef7e89086bef Mon Sep 17 00:00:00 2001 From: "assistant-librarian[bot]" <210906412+assistant-librarian[bot]@users.noreply.github.com> Date: Tue, 10 Feb 2026 22:52:00 -0700 Subject: [PATCH] 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 --- 🔁 Imported from [ROCm/composable_kernel#3671](https://github.com/ROCm/composable_kernel/pull/3671) 🧑‍💻 Originally authored by @ThomasNing --------- Co-authored-by: ThomasNing Co-authored-by: systems-assistant[bot] Co-authored-by: illsilin_amdeng --- .../ops/epilogue/cshuffle_epilogue.hpp | 111 +++++++++++++++++- .../ops/softmax/block/block_softmax_2d.hpp | 6 +- 2 files changed, 109 insertions(+), 8 deletions(-) diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 11c42bf4bb..e9a11909c7 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -302,19 +302,118 @@ struct CShuffleEpilogue template 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) { - return make_naive_tensor_descriptor( - make_tuple(number{}, number{}), - make_tuple(number{}, 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{}, + number{}, + number{}), + make_tuple(number{}, + number{}, + number<1>{}), + number{}, + number<1>{}); + + constexpr auto lds_block_desc_1 = transform_tensor_descriptor( + lds_block_desc_0, + make_tuple(make_pass_through_transform(number{}), + make_unmerge_transform(make_tuple( + number{}, number{})), + make_pass_through_transform(number{})), + 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{}, number{})), + make_merge_transform_v3_division_mod(make_tuple( + number{}, number{}))), + 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) { - return make_naive_tensor_descriptor( - make_tuple(number{}, number{}), - make_tuple(number<1>{}, number{})); + 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{}, + number{}, + number{}), + make_tuple(number{}, + number{}, + number<1>{}), + number{}, + number<1>{}); + + constexpr auto lds_block_desc_1 = transform_tensor_descriptor( + lds_block_desc_0, + make_tuple(make_pass_through_transform(number{}), + make_unmerge_transform(make_tuple( + number{}, number{})), + make_pass_through_transform(number{})), + 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{}, number{})), + make_merge_transform_v3_division_mod(make_tuple( + number{}, number{}))), + make_tuple(sequence<0, 1>{}, sequence<2, 3>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + + return lds_block_desc; } else { diff --git a/include/ck_tile/ops/softmax/block/block_softmax_2d.hpp b/include/ck_tile/ops/softmax/block/block_softmax_2d.hpp index 58e768b319..cb861935f2 100644 --- a/include/ck_tile/ops/softmax/block/block_softmax_2d.hpp +++ b/include/ck_tile/ops/softmax/block/block_softmax_2d.hpp @@ -40,7 +40,8 @@ struct BlockSoftmax2D #endif // compute row max - auto reduce_row_max = BlockReduce2D{x, -numeric::infinity()}; + using X = remove_cvref_t; + BlockReduce2D reduce_row_max{x, -numeric::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{y, DataType{0}}; + using Y = remove_cvref_t; + BlockReduce2D 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