Epilogue cshuffle Improvement (#2312)

* add cshuffle's mxdlperwavepershuffle support, not finished * add epilogue functions * add cshuffle's mxdlperwavepershuffle support, not finished * add epilogue functions * update cshuffle logic * update cshuffle_logics * add some change within review * update some codes following the code review * update epilogue logic * remove from problem * update codes following review. * fix some issues * solve the previous PR error, refine the code * Update include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com> * Comment addressed * handling tile_engine failing case * handling tile_engine failing case --------- Co-authored-by: joyeamd <John.Ye@amd.com> Co-authored-by: joye <joye@amd.com> Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com> Co-authored-by: khushbu agarwal <khuagarw@amd.com> [ROCm/composable_kernel commit: 06e0b8436c]
2026-05-14 02:02:46 +00:00 · 2025-06-10 22:44:50 -07:00
parent a0af2eca3f
commit 2350191009
1 changed files with 133 additions and 61 deletions
--- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
+++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
@@ -17,11 +17,11 @@ template <typename ADataType_,
          index_t kBlockSize_,
          index_t kM_,
          index_t kN_,
-          index_t kMWave_,
-          index_t kNWave_,
-          index_t kMPerXdl_,
-          index_t kNPerXdl_,
-          index_t kKPerXdl_,
+          index_t MWave_,
+          index_t NWave_,
+          index_t MPerXdl_,
+          index_t NPerXdl_,
+          index_t KPerXdl_,
          bool isCTransposed_,
          memory_operation_enum MemoryOperation_>
 struct CShuffleEpilogueProblem
@@ -34,11 +34,11 @@ struct CShuffleEpilogueProblem
    static constexpr index_t kBlockSize                    = kBlockSize_;
    static constexpr index_t kMPerBlock                    = kM_;
    static constexpr index_t kNPerBlock                    = kN_;
-    static constexpr index_t kMWave                        = kMWave_;
-    static constexpr index_t kNWave                        = kNWave_;
-    static constexpr index_t kMPerXdl                      = kMPerXdl_;
-    static constexpr index_t kNPerXdl                      = kNPerXdl_;
-    static constexpr index_t kKPerXdl                      = kKPerXdl_;
+    static constexpr index_t MWave                         = MWave_;
+    static constexpr index_t NWave                         = NWave_;
+    static constexpr index_t MPerXdl                       = MPerXdl_;
+    static constexpr index_t NPerXdl                       = NPerXdl_;
+    static constexpr index_t KPerXdl                       = KPerXdl_;
    static constexpr index_t isCTransposed                 = isCTransposed_;
    static constexpr memory_operation_enum MemoryOperation = MemoryOperation_;
 };
@@ -59,25 +59,14 @@ struct CShuffleEpilogue
    static constexpr index_t kBlockSize                    = Problem::kBlockSize;
    static constexpr index_t kMPerBlock                    = Problem::kMPerBlock;
    static constexpr index_t kNPerBlock                    = Problem::kNPerBlock;
-    static constexpr index_t kMWave                        = Problem::kMWave;
-    static constexpr index_t kNWave                        = Problem::kNWave;
-    static constexpr index_t kMPerXdl                      = Problem::kMPerXdl;
-    static constexpr index_t kNPerXdl                      = Problem::kNPerXdl;
-    static constexpr index_t kKPerXdl                      = Problem::kKPerXdl;
+    static constexpr index_t MWave                         = Problem::MWave;
+    static constexpr index_t NWave                         = Problem::NWave;
+    static constexpr index_t MPerXdl                       = Problem::MPerXdl;
+    static constexpr index_t NPerXdl                       = Problem::NPerXdl;
+    static constexpr index_t KPerXdl                       = Problem::KPerXdl;
    static constexpr index_t isCTransposed                 = Problem::isCTransposed;
-    static constexpr index_t kMPerIteration                = kMPerXdl * kMWave;
-    static constexpr index_t kNPerIteration                = kNPerXdl * kNWave;
-
-    using WG = WarpGemmMfmaDispatcher<ADataType,
-                                      BTypeToUse,
-                                      AccDataType,
-                                      kMPerXdl,
-                                      kNPerXdl,
-                                      kKPerXdl,
-                                      isCTransposed>;
-
-    using CWarpDstr   = typename WG::CWarpDstr;
-    using CWarpTensor = typename WG::CWarpTensor;
+    static constexpr index_t MPerIteration                 = MPerXdl * MWave;
+    static constexpr index_t NPerIteration                 = NPerXdl * NWave;

    /**
     * @brief Get the vector store size for C tensor.
@@ -89,18 +78,18 @@ struct CShuffleEpilogue
     *
     * @return The vector store size for C tensor.
     */
-    CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeC()
+    CK_TILE_HOST_DEVICE static constexpr index_t GetVectorSizeC()
    {
-        constexpr index_t MaxVectorStoreSize = 16;
+        constexpr index_t max_vector_size = 16;
        if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
        {
-            return std::min(static_cast<int>(kNPerIteration),
-                            static_cast<int>(MaxVectorStoreSize / sizeof(ODataType)));
+            return std::min(static_cast<int>(NPerIteration),
+                            static_cast<int>(max_vector_size / sizeof(ODataType)));
        }
        else if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::ColumnMajor>)
        {
-            return std::min(static_cast<int>(kMPerIteration),
-                            static_cast<int>(MaxVectorStoreSize / sizeof(ODataType)));
+            return std::min(static_cast<int>(MPerIteration),
+                            static_cast<int>(max_vector_size / sizeof(ODataType)));
        }
        else
        {
@@ -108,6 +97,65 @@ struct CShuffleEpilogue
        }
    }

+    /**
+     * @brief Shuffle tile configuration parameters
+     *
+     * @details These parameters control the number of XDL tiles processed per wave in each shuffle
+     * iteration:
+     * - NumMXdlPerWavePerShuffle: Number of XDL tiles in M dimension processed per wave
+     * - 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())
+        {
+            return std::make_tuple(1, 1);
+        }
+        else
+        {
+            constexpr index_t num_xdl_shuffles = GetVectorSizeC() / elem_per_thread;
+            if constexpr(std::is_same_v<CLayout, 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);
+            }
+            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)));
+            }
+        }
+    }();
+    static constexpr index_t NumMXdlPerWavePerShuffle = std::get<0>(shuffle_tile_tuple);
+    static constexpr index_t NumNXdlPerWavePerShuffle = std::get<1>(shuffle_tile_tuple);
+
+    static constexpr auto MNPerIterationShuffle = [] {
+        constexpr index_t m_val = MPerXdl * MWave * NumMXdlPerWavePerShuffle;
+        constexpr index_t n_val = NPerXdl * NWave * NumNXdlPerWavePerShuffle;
+        if constexpr(kMPerBlock % m_val != 0 || kNPerBlock % n_val != 0)
+            return std::make_tuple(MPerXdl * MWave, NPerXdl * NWave);
+        else
+            return std::make_tuple(m_val, n_val);
+    }();
+    static constexpr index_t MPerIterationShuffle = std::get<0>(MNPerIterationShuffle);
+    static constexpr index_t NPerIterationShuffle = std::get<1>(MNPerIterationShuffle);
+    using WG                                      = WarpGemmMfmaDispatcher<ADataType,
+                                      BTypeToUse,
+                                      AccDataType,
+                                      MPerXdl,
+                                      NPerXdl,
+                                      KPerXdl,
+                                      isCTransposed>;
+
+    using CWarpDstr   = typename WG::CWarpDstr;
+    using CWarpTensor = typename WG::CWarpTensor;
+
    template <typename Problem>
    CK_TILE_HOST_DEVICE static constexpr auto MakeLdsBlockDescriptor()
    {
@@ -115,15 +163,15 @@ struct CShuffleEpilogue
        if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
        {
            return make_naive_tensor_descriptor(
-                make_tuple(number<kMWave * kMPerXdl>{}, number<kNWave * kNPerXdl>{}),
-                make_tuple(number<kNWave * kNPerXdl>{}, number<1>{}));
+                make_tuple(number<MPerIterationShuffle>{}, number<NPerIterationShuffle>{}),
+                make_tuple(number<NPerIterationShuffle>{}, number<1>{}));
        }
        // M is contiguous dimension
        else if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::ColumnMajor>)
        {
            return make_naive_tensor_descriptor(
-                make_tuple(number<kMWave * kMPerXdl>{}, number<kNWave * kNPerXdl>{}),
-                make_tuple(number<1>{}, number<kMWave * kMPerXdl>{}));
+                make_tuple(number<MPerIterationShuffle>{}, number<NPerIterationShuffle>{}),
+                make_tuple(number<1>{}, number<MPerIterationShuffle>{}));
        }
        else
        {
@@ -131,40 +179,62 @@ struct CShuffleEpilogue
        }
    }

+    CK_TILE_DEVICE static constexpr auto MakeLdsDistributionEncode()
+    {
+        constexpr auto block_outer_dstr_encoding =
+            tile_distribution_encoding<sequence<>,
+                                       tuple<sequence<NumMXdlPerWavePerShuffle, MWave>,
+                                             sequence<NumNXdlPerWavePerShuffle, NWave>>,
+                                       tuple<sequence<1, 2>>,
+                                       tuple<sequence<1, 1>>,
+                                       sequence<1, 2>,
+                                       sequence<0, 0>>{};
+        constexpr auto block_dstr_encoding = detail::make_embed_tile_distribution_encoding(
+            block_outer_dstr_encoding, typename CWarpDstr::DstrEncode{});
+
+        return block_dstr_encoding;
+    }
+
    CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
    {
-        return kMWave * kNWave * kMPerXdl * kNPerXdl * sizeof(ODataType);
+        return MPerIterationShuffle * NPerIterationShuffle * sizeof(ODataType);
    }

    template <typename ODramWindow, typename OAccTile>
    CK_TILE_DEVICE auto
    operator()(ODramWindow& out_dram_window, const OAccTile& o_acc_tile, void* p_smem)
    {
+        constexpr auto LdsTileDistr = make_static_tile_distribution(MakeLdsDistributionEncode());

-        const index_t iMWarp = get_warp_id() / kNWave;
-        const index_t iNWarp = get_warp_id() - iMWarp * kNWave;
+        auto lds_tile = make_static_distributed_tensor<AccDataType>(LdsTileDistr);

        constexpr auto lds_block_desc = MakeLdsBlockDescriptor<Problem>();
        auto o_lds_block              = make_tensor_view<address_space_enum::lds>(
            static_cast<ODataType*>(p_smem), lds_block_desc);
-        auto in_lds_window =
-            make_tile_window(o_lds_block,
-                             make_tuple(number<kMPerXdl>{}, number<kNPerXdl>{}),
-                             {number<kMPerXdl>{} * iMWarp, number<kNPerXdl>{} * iNWarp});
-        auto out_lds_window =
-            make_tile_window(o_lds_block,
-                             make_tuple(number<kMWave * kMPerXdl>{}, number<kNWave * kNPerXdl>{}),
-                             {0, 0});
+
+        auto in_lds_window = make_tile_window(
+            o_lds_block,
+            make_tuple(number<MPerIterationShuffle>{}, number<NPerIterationShuffle>{}),
+            {0, 0},
+            LdsTileDistr);
+
+        auto out_lds_window = make_tile_window(
+            o_lds_block,
+            make_tuple(number<MPerIterationShuffle>{}, number<NPerIterationShuffle>{}),
+            {0, 0});

        using SFC                    = space_filling_curve<sequence<kMPerBlock, kNPerBlock>,
                                        sequence<0, 1>,
-                                        sequence<kMPerXdl * kMWave, kNPerXdl * kNWave>>;
+                                        sequence<MPerIterationShuffle, NPerIterationShuffle>>;
        constexpr index_t num_access = SFC::get_num_of_access();

+        static_assert(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>,
+                      "Currently, the CShuffle Epilogue only supports the Row Major Output layout");
+
        using TileEncodingPattern =
            TileDistributionEncodingPattern2D<kBlockSize,
-                                              kMPerIteration,
-                                              kNPerIteration,
+                                              MPerIterationShuffle,
+                                              NPerIterationShuffle,
                                              GetVectorSizeC(),
                                              tile_distribution_pattern::thread_raked>;
        constexpr auto dram_tile_distribution = TileEncodingPattern::Make2DStaticTileDistribution();
@@ -173,21 +243,23 @@ struct CShuffleEpilogue
            to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
        constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t<CWarpDstr::NDimY, 0>{};

-        CWarpTensor c_warp_in_tensor;
        static_for<0, num_access, 1>{}([&](auto iAccess) {
+            block_sync_lds();
            constexpr auto idx_y_start = SFC::get_index(iAccess);

-            constexpr auto mIter = number<idx_y_start.at(number<0>{}) / (kMPerXdl * kMWave)>{};
-            constexpr auto nIter = number<idx_y_start.at(number<1>{}) / (kNPerXdl * kNWave)>{};
+            constexpr auto mIter = number<idx_y_start.at(number<0>{}) / (MPerIterationShuffle)>{};
+            constexpr auto nIter = number<idx_y_start.at(number<1>{}) / (NPerIterationShuffle)>{};

-            c_warp_in_tensor.get_thread_buffer() = o_acc_tile.get_y_sliced_thread_data(
-                merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
-                merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
+            lds_tile.get_thread_buffer() = o_acc_tile.get_y_sliced_thread_data(
+                merge_sequences(
+                    sequence<mIter * NumMXdlPerWavePerShuffle, nIter * NumNXdlPerWavePerShuffle>{},
+                    c_warp_y_index_zeros),
+                merge_sequences(sequence<NumMXdlPerWavePerShuffle, NumNXdlPerWavePerShuffle>{},
+                                c_warp_y_lengths));

-            const auto c_warp_in_tensor_casted = cast_tile<ODataType>(c_warp_in_tensor);
+            const auto c_warptile_in_tensor_casted = cast_tile<ODataType>(lds_tile);

-            block_sync_lds();
-            store_tile(in_lds_window, c_warp_in_tensor_casted);
+            store_tile(in_lds_window, c_warptile_in_tensor_casted);
            block_sync_lds();

            const auto c_out_tensor =