Revert "Optimize fmha fwd decode & prefill for gfx950 (#2641)" (#2670)

This reverts commit b7322a521a.

include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp

@@ -42,8 +42,6 @@ struct BlockGemmARegBRegCRegV1
         static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN);
         static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK;
 
-        static constexpr auto BlockGemmLoopOrder = Problem::BlockGemmLoopOrder;
-
         static constexpr index_t KPack = WarpGemm::kKPerThread;
     };
 
@@ -54,9 +52,8 @@ struct BlockGemmARegBRegCRegV1
 
     using Traits = GemmTraits_<Problem, Policy>;
 
-    using WarpGemm                           = typename Traits::WarpGemm;
-    using BlockGemmShape                     = typename Traits::BlockGemmShape;
-    static constexpr auto BlockGemmLoopOrder = Traits::BlockGemmLoopOrder;
+    using WarpGemm       = typename Traits::WarpGemm;
+    using BlockGemmShape = typename Traits::BlockGemmShape;
 
     using ADataType = remove_cvref_t<typename Traits::ADataType>;
     using BDataType = remove_cvref_t<typename Traits::BDataType>;
@@ -89,36 +86,17 @@ struct BlockGemmARegBRegCRegV1
         }
         else
         {
-            if constexpr(BlockGemmLoopOrder == GemmLoopOrder::KMN)
-            {
-                constexpr auto a_block_outer_dstr_encoding = tile_distribution_encoding<
-                    sequence<NWarp>,
-                    tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
-                    tuple<sequence<1, 0>>,
-                    tuple<sequence<1, 0>>,
-                    sequence<2, 1>,
-                    sequence<0, 0>>{};
+            constexpr auto a_block_outer_dstr_encoding = tile_distribution_encoding<
+                sequence<NWarp>,
+                tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
+                tuple<sequence<1, 0>>,
+                tuple<sequence<1, 0>>,
+                sequence<1, 2>,
+                sequence<0, 0>>{};
+            constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
+                a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{});
 
-                constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
-                    a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{});
-
-                return a_block_dstr_encode;
-            }
-            else if constexpr(BlockGemmLoopOrder == GemmLoopOrder::MNK)
-            {
-                constexpr auto a_block_outer_dstr_encoding = tile_distribution_encoding<
-                    sequence<NWarp>,
-                    tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
-                    tuple<sequence<1, 0>>,
-                    tuple<sequence<1, 0>>,
-                    sequence<1, 2>,
-                    sequence<0, 0>>{};
-
-                constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
-                    a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{});
-
-                return a_block_dstr_encode;
-            }
+            return a_block_dstr_encode;
         }
     }
 
@@ -140,33 +118,17 @@ struct BlockGemmARegBRegCRegV1
         }
         else
         {
-            if constexpr(BlockGemmLoopOrder == GemmLoopOrder::KMN)
-            {
-                constexpr auto b_block_outer_dstr_encoding = tile_distribution_encoding<
-                    sequence<MWarp>,
-                    tuple<sequence<NIterPerWarp, NWarp>, sequence<KIterPerWarp>>,
-                    tuple<sequence<0, 1>>,
-                    tuple<sequence<0, 1>>,
-                    sequence<2, 1>,
-                    sequence<0, 0>>{};
-                constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
-                    b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{});
+            constexpr auto b_block_outer_dstr_encoding = tile_distribution_encoding<
+                sequence<MWarp>,
+                tuple<sequence<NIterPerWarp, NWarp>, sequence<KIterPerWarp>>,
+                tuple<sequence<0, 1>>,
+                tuple<sequence<0, 1>>,
+                sequence<1, 2>,
+                sequence<0, 0>>{};
+            constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
+                b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{});
 
-                return b_block_dstr_encode;
-            }
-            else if constexpr(BlockGemmLoopOrder == GemmLoopOrder::MNK)
-            {
-                constexpr auto b_block_outer_dstr_encoding = tile_distribution_encoding<
-                    sequence<MWarp>,
-                    tuple<sequence<NIterPerWarp, NWarp>, sequence<KIterPerWarp>>,
-                    tuple<sequence<0, 1>>,
-                    tuple<sequence<0, 1>>,
-                    sequence<1, 2>,
-                    sequence<0, 0>>{};
-                constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
-                    b_block_outer_dstr_encoding, typename WarpGemm::BWarpDstrEncoding{});
-                return b_block_dstr_encode;
-            }
+            return b_block_dstr_encode;
         }
     }
 
@@ -251,82 +213,40 @@ struct BlockGemmARegBRegCRegV1
         constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t<CWarpDstr::NDimY, 0>{};
 
         // hot loop:
-        if constexpr(BlockGemmLoopOrder == GemmLoopOrder::KMN)
-        {
-            static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
-                static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
-                    // read A warp tensor from A Block window
-                    AWarpTensor a_warp_tensor;
-                    a_warp_tensor.get_thread_buffer() = a_block_tensor.get_y_sliced_thread_data(
-                        merge_sequences(sequence<kIter, mIter>{}, a_warp_y_index_zeros),
-                        merge_sequences(sequence<1, 1>{}, a_warp_y_lengths));
-
-                    static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
-                        // read B warp tensor from B block tensor
-                        BWarpTensor b_warp_tensor;
-                        b_warp_tensor.get_thread_buffer() = b_block_tensor.get_y_sliced_thread_data(
-                            merge_sequences(sequence<kIter, nIter>{}, b_warp_y_index_zeros),
-                            merge_sequences(sequence<1, 1>{}, b_warp_y_lengths));
-
-                        // read C warp tensor from C block tensor
-                        using c_iter_idx = std::conditional_t<TransposeC,
-                                                              sequence<nIter, mIter>,
-                                                              sequence<mIter, nIter>>;
-                        CWarpTensor c_warp_tensor;
-                        c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
-                            merge_sequences(c_iter_idx{}, c_warp_y_index_zeros),
-                            merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
-
-                        // warp GEMM
-                        WarpGemm{}(c_warp_tensor, a_warp_tensor, b_warp_tensor);
-
-                        // write C warp tensor into C block tensor
-                        c_block_tensor.set_y_sliced_thread_data(
-                            merge_sequences(c_iter_idx{}, c_warp_y_index_zeros),
-                            merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
-                            c_warp_tensor.get_thread_buffer());
-                    });
-                });
-            });
-        }
-        else if constexpr(BlockGemmLoopOrder == GemmLoopOrder::MNK)
-        {
+        static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
             static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
+                // read A warp tensor from A Block window
+                AWarpTensor a_warp_tensor;
+                a_warp_tensor.get_thread_buffer() = a_block_tensor.get_y_sliced_thread_data(
+                    merge_sequences(sequence<mIter, kIter>{}, a_warp_y_index_zeros),
+                    merge_sequences(sequence<1, 1>{}, a_warp_y_lengths));
+
                 static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
-                    static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
-                        // read A warp tensor from A Block window
-                        AWarpTensor a_warp_tensor;
+                    // read B warp tensor from B block tensor
+                    BWarpTensor b_warp_tensor;
+                    b_warp_tensor.get_thread_buffer() = b_block_tensor.get_y_sliced_thread_data(
+                        merge_sequences(sequence<nIter, kIter>{}, b_warp_y_index_zeros),
+                        merge_sequences(sequence<1, 1>{}, b_warp_y_lengths));
 
-                        a_warp_tensor.get_thread_buffer() = a_block_tensor.get_y_sliced_thread_data(
-                            merge_sequences(sequence<mIter, kIter>{}, a_warp_y_index_zeros),
-                            merge_sequences(sequence<1, 1>{}, a_warp_y_lengths));
+                    // read C warp tensor from C block tensor
+                    using c_iter_idx = std::
+                        conditional_t<TransposeC, sequence<nIter, mIter>, sequence<mIter, nIter>>;
+                    CWarpTensor c_warp_tensor;
+                    c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
+                        merge_sequences(c_iter_idx{}, c_warp_y_index_zeros),
+                        merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
 
-                        // read B warp tensor from B block tensor
-                        BWarpTensor b_warp_tensor;
+                    // warp GEMM
+                    WarpGemm{}(c_warp_tensor, a_warp_tensor, b_warp_tensor);
 
-                        b_warp_tensor.get_thread_buffer() = b_block_tensor.get_y_sliced_thread_data(
-                            merge_sequences(sequence<nIter, kIter>{}, b_warp_y_index_zeros),
-                            merge_sequences(sequence<1, 1>{}, b_warp_y_lengths));
-
-                        // read C warp tensor from C block tensor
-                        CWarpTensor c_warp_tensor;
-
-                        c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
-                            merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
-                            merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
-
-                        // warp GEMM
-                        WarpGemm{}(c_warp_tensor, a_warp_tensor, b_warp_tensor);
-
-                        // write C warp tensor into C block tensor
-                        c_block_tensor.set_y_sliced_thread_data(
-                            merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
-                            merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
-                            c_warp_tensor.get_thread_buffer());
-                    });
+                    // write C warp tensor into C block tensor
+                    c_block_tensor.set_y_sliced_thread_data(
+                        merge_sequences(c_iter_idx{}, c_warp_y_index_zeros),
+                        merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
+                        c_warp_tensor.get_thread_buffer());
                 });
             });
-        }
+        });
     }
 
     CK_TILE_DEVICE static constexpr auto MakeCBlockTile()

include/ck_tile/ops/gemm/block/block_gemm_problem.hpp

@@ -4,7 +4,6 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
-#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp"
 
 namespace ck_tile {
 
@@ -14,8 +13,7 @@ template <typename ADataType_,
           typename CDataType_,
           index_t kBlockSize_,
           typename BlockGemmShape_,
-          GemmLoopOrder BlockGemmLoopOrder_ = GemmLoopOrder::KMN,
-          index_t NumWaveGroups_            = 1>
+          index_t NumWaveGroups_ = 1>
 struct BlockGemmProblem
 {
     using ADataType      = remove_cvref_t<ADataType_>;
@@ -23,9 +21,8 @@ struct BlockGemmProblem
     using CDataType      = remove_cvref_t<CDataType_>;
     using BlockGemmShape = remove_cvref_t<BlockGemmShape_>;
 
-    static constexpr index_t kBlockSize               = kBlockSize_;
-    static constexpr index_t NumWaveGroups            = NumWaveGroups_;
-    static constexpr GemmLoopOrder BlockGemmLoopOrder = BlockGemmLoopOrder_;
+    static constexpr index_t kBlockSize    = kBlockSize_;
+    static constexpr index_t NumWaveGroups = NumWaveGroups_;
 };
 
 } // namespace ck_tile