refactor(ck): mx gemm kernel unification

projects/composablekernel/example/ck_tile/42_mx_gemm/mx_gemm.cpp

@@ -24,13 +24,13 @@ static constexpr inline auto is_row_major(Layout layout_)
 template <typename GemmConfig,
           typename ADataType,
           typename BDataType,
+          typename AScaleDataType,
+          typename BScaleDataType,
           typename AccDataType,
           typename CDataType,
           typename ALayout,
           typename BLayout,
           typename CLayout,
-          typename ScaleM,
-          typename ScaleN,
           bool UsePersistentKernel = false>
 float invoke_mx_gemm(ck_tile::DeviceMem& a_dev_buf,
                      ck_tile::DeviceMem& b_dev_buf,
@@ -42,36 +42,38 @@ float invoke_mx_gemm(ck_tile::DeviceMem& a_dev_buf,
                      ck_tile::index_t stride_B,
                      ck_tile::index_t stride_C,
                      ck_tile::index_t kbatch,
-                     ScaleM scale_m,
-                     ScaleN scale_n,
+                     ck_tile::DeviceMem& scale_m,
+                     ck_tile::DeviceMem& scale_n,
                      int n_warmup,
                      int n_repeat)
 {
-    MXGemmHostArgs<ScaleM, ScaleN> args(a_dev_buf.GetDeviceBuffer(),
-                                        b_dev_buf.GetDeviceBuffer(),
-                                        c_dev_buf.GetDeviceBuffer(),
-                                        kbatch,
-                                        M,
-                                        N,
-                                        K,
-                                        stride_A,
-                                        stride_B,
-                                        stride_C,
-                                        scale_m,
-                                        scale_n);
+    ck_tile::MxGemmHostArgs<1, 1, 0> args({static_cast<const void*>(a_dev_buf.GetDeviceBuffer())},
+                                          {static_cast<const void*>(scale_m.GetDeviceBuffer())},
+                                          {static_cast<const void*>(b_dev_buf.GetDeviceBuffer())},
+                                          {static_cast<const void*>(scale_n.GetDeviceBuffer())},
+                                          {},
+                                          c_dev_buf.GetDeviceBuffer(),
+                                          kbatch,
+                                          M,
+                                          N,
+                                          K,
+                                          {stride_A},
+                                          {stride_B},
+                                          {},
+                                          stride_C);
 
     // Simplified invocation - comp_async handles hot loop and tail internally
     auto invoke_splitk_path = [&](auto split_k_) {
         return mx_gemm_calc<GemmConfig,
                             ADataType,
                             BDataType,
+                            AScaleDataType,
+                            BScaleDataType,
                             AccDataType,
                             CDataType,
                             ALayout,
                             BLayout,
                             CLayout,
-                            ScaleM,
-                            ScaleN,
                             UsePersistentKernel,
                             split_k_.value>(
             args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat, true, true, 50});

projects/composablekernel/example/ck_tile/42_mx_gemm/mx_gemm.hpp

@@ -9,49 +9,8 @@
 #include "ck_tile/host/kernel_launch.hpp"
 #include "ck_tile/ops/epilogue.hpp"
 #include "ck_tile/ops/gemm.hpp"
-#include "ck_tile/ops/gemm_mx.hpp"
-#include "ck_tile/ops/gemm_mx/kernel/scale_pointer.hpp"
-
-template <typename ScaleM, typename ScaleN>
-struct MXGemmHostArgs : ck_tile::UniversalGemmHostArgs<1, 1, 0>
-{
-    using Base = ck_tile::UniversalGemmHostArgs<1, 1, 0>;
-
-    MXGemmHostArgs(const void* a_ptr,
-                   const void* b_ptr,
-                   void* c_ptr_,
-                   ck_tile::index_t k_batch_,
-                   ck_tile::index_t M_,
-                   ck_tile::index_t N_,
-                   ck_tile::index_t K_,
-                   ck_tile::index_t stride_A_,
-                   ck_tile::index_t stride_B_,
-                   ck_tile::index_t stride_C_,
-                   ScaleM scale_m_,
-                   ScaleN scale_n_)
-        : Base({a_ptr},
-               {b_ptr},
-               {},
-               c_ptr_,
-               k_batch_,
-               M_,
-               N_,
-               K_,
-               {stride_A_},
-               {stride_B_},
-               {},
-               stride_C_),
-          scale_m(scale_m_),
-          scale_n(scale_n_)
-    {
-    }
-
-    ScaleM scale_m;
-    ScaleN scale_n;
-};
 
 // GEMM config with 16x16 warp tile
-
 struct MxGemmConfig
 {
     static constexpr ck_tile::index_t M_Tile = 128;
@@ -78,12 +37,15 @@ struct MxGemmConfig
     static constexpr int TileParitionerM01          = 4;
     static constexpr auto Scheduler                 = ck_tile::GemmPipelineScheduler::Intrawave;
     static constexpr ck_tile::index_t NumWaveGroups = 1;
-    static constexpr bool DoubleSmemBuffer          = false; // comp_async uses double buffer
+    static constexpr bool DoubleSmemBuffer          = true; // comp_async uses double buffer
     static constexpr bool Preshuffle                = false;
     static constexpr ck_tile::index_t BContiguousItemsPerAccess = 16;
 
     static constexpr int N_Repeat          = N_Tile / N_Warp_Tile / N_Warp;
     static constexpr bool TiledMMAPermuteN = false;
+
+    using AScaleDataType = ck_tile::e8m0_t;
+    using BScaleDataType = ck_tile::e8m0_t;
 };
 
 struct MX_GemmConfigEightWaves : MxGemmConfig

projects/composablekernel/example/ck_tile/42_mx_gemm/mx_gemm_instance.hpp

@@ -5,10 +5,9 @@
 
 #include "ck_tile/host.hpp"
 #include "mx_gemm.hpp"
-#include "ck_tile/ops/gemm_mx/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp"
-#include "ck_tile/ops/gemm_mx/pipeline/gemm_pipeline_ag_bg_cr_comp_async_eight_waves.hpp"
-#include "ck_tile/ops/gemm_mx/kernel/gemm_mx_kernel.hpp"
-#include "ck_tile/ops/gemm_mx/pipeline/wp_pipeline_agmem_bgmem_creg_v1.hpp"
+#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp"
+#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_eight_waves.hpp"
+#include "ck_tile/ops/gemm/pipeline/wp_mx_pipeline_agmem_bgmem_creg_v1.hpp"
 
 template <typename Layout>
 using is_row_major_t = ck_tile::bool_constant<
@@ -17,16 +16,16 @@ using is_row_major_t = ck_tile::bool_constant<
 template <typename GemmConfig,
           typename ADataType,
           typename BDataType,
+          typename AScaleDataType,
+          typename BScaleDataType,
           typename AccDataType,
           typename CDataType,
           typename ALayout,
           typename BLayout,
           typename CLayout,
-          typename ScaleM,
-          typename ScaleN,
           bool persistent,
           bool Splitk>
-float mx_gemm_calc(const MXGemmHostArgs<ScaleM, ScaleN>& args, const ck_tile::stream_config& s)
+float mx_gemm_calc(const ck_tile::MxGemmHostArgs<1, 1, 0>& args, const ck_tile::stream_config& s)
 {
     using GemmShape = ck_tile::TileGemmShape<
         ck_tile::sequence<GemmConfig::M_Tile, GemmConfig::N_Tile, GemmConfig::K_Tile>,
@@ -51,29 +50,38 @@ float mx_gemm_calc(const MXGemmHostArgs<ScaleM, ScaleN>& args, const ck_tile::st
     static_assert(sizeof(ComputeDataType) >= sizeof(BDataType),
                   "mixed_prec_gemm requires ADataType is a wider type than BDataType");
 
-    using MXPipelineProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
-                                                                    BDataType,
-                                                                    AccDataType,
-                                                                    GemmShape,
-                                                                    MXGemmTraits,
-                                                                    GemmConfig::Scheduler>;
+    using AComputeDataType = ADataType;
+    using BComputeDataType = BDataType;
+
+    using MXPipelineProblem = ck_tile::MxGemmPipelineProblem<ADataType,
+                                                             BDataType,
+                                                             AccDataType,
+                                                             GemmShape,
+                                                             MXGemmTraits,
+                                                             GemmConfig::Scheduler,
+                                                             ck_tile::element_wise::PassThrough,
+                                                             ck_tile::element_wise::PassThrough,
+                                                             AComputeDataType,
+                                                             BComputeDataType,
+                                                             AScaleDataType,
+                                                             BScaleDataType>;
 
-    // Use the MX GEMM Preshuffle pipeline or
-    // the new MX comp_async pipeline with MX scaling support
     constexpr bool IsEightWave =
         (GemmConfig::M_Warp * GemmConfig::N_Warp * GemmConfig::K_Warp) == 8;
     using MXGemmPipeline = std::conditional_t<
         GemmConfig::Preshuffle,
         ck_tile::MXGemmPreshufflePipelineAGmemBGmemCRegV1<MXPipelineProblem>,
         std::conditional_t<IsEightWave,
-                           ck_tile::MXGemmPipelineAgBgCrCompAsyncEightWaves<MXPipelineProblem>,
-                           ck_tile::MXGemmPipelineAgBgCrCompAsync<MXPipelineProblem>>>;
+                           ck_tile::GemmPipelineAgBgCrCompAsyncEightWaves<MXPipelineProblem>,
+                           ck_tile::GemmPipelineAgBgCrCompAsync<MXPipelineProblem>>>;
 
     using TilePartitioner =
         ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
                                                    GemmConfig::TileParitionerGroupNum,
                                                    GemmConfig::TileParitionerM01>;
 
+    constexpr ck_tile::index_t BlockedXDLNPerWarp = GemmConfig::Preshuffle ? 2 : 1;
+
     using GemmEpilogue =
         std::conditional_t<GemmConfig::TiledMMAPermuteN,
                            ck_tile::PermuteNEpilogue<
@@ -115,38 +123,27 @@ float mx_gemm_calc(const MXGemmHostArgs<ScaleM, ScaleN>& args, const ck_tile::st
                                GemmConfig::NumWaveGroups,
                                false, // FixedVectorSize_ (Default)
                                1,     // VectorSizeC_ (Default)
-                               ck_tile::MXEpilogueTraits<GemmConfig>::BlockedXDLNPerWarp,
+                               BlockedXDLNPerWarp,
                                false,                      // DoubleSmemBuffer_ (Default)
                                ComputeDataType,            // AComputeDataType
                                ComputeDataType,            // BComputeDataType
                                !GemmConfig::Preshuffle>>>; // TilesPacked_ (because of
                                                            // packed scales)
 
-    using Kernel = ck_tile::MXGemmKernel<TilePartitioner, MXGemmPipeline, GemmEpilogue>;
-
-    auto kargs = Kernel::MakeKernelArgs(std::array<const void*, 1>{args.as_ptr},
-                                        std::array<const void*, 1>{args.bs_ptr},
-                                        std::array<const void*, 0>{},
-                                        args.e_ptr,
-                                        args.k_batch,
-                                        args.M,
-                                        args.N,
-                                        args.K,
-                                        std::array<ck_tile::index_t, 1>{args.stride_As},
-                                        std::array<ck_tile::index_t, 1>{args.stride_Bs},
-                                        std::array<ck_tile::index_t, 0>{},
-                                        args.stride_E,
-                                        args.scale_m,
-                                        args.scale_n);
+    using Kernel = ck_tile::MxGemmKernel<TilePartitioner, MXGemmPipeline, GemmEpilogue>;
+
+    auto kargs = Kernel::MakeKernelArgs(args);
 
     if(!Kernel::IsSupportedArgument(kargs))
     {
         throw std::runtime_error(
             "MX GEMM: unsupported shape/configuration (set CK_TILE_LOGGING=1 for details).");
     }
 
-    const auto kernel = ck_tile::make_kernel<Kernel::kBlockPerCu>(
-        Kernel{}, Kernel::GridSize(kargs), Kernel::BlockSize(), 0, kargs);
+    constexpr int kBlockPerCu = 1;
+
+    const auto kernel = ck_tile::make_kernel<kBlockPerCu>(
+        Kernel{}, Kernel::GridSize(args.M, args.N, args.k_batch), Kernel::BlockSize(), 0, kargs);
 
     // For split-K (k_batch > 1) the kernel's epilogue uses atomic_add into C, so C must be
     // zeroed before every kernel launch -- not just once before the first warmup iteration.