From feeb6f7cd831dbf9feb02df9d224eb73690dc54f Mon Sep 17 00:00:00 2001
From: Qun Lin <qlin@amd.com>
Date: Thu, 15 May 2025 17:48:21 +0800
Subject: [PATCH] refine universal gemm multi-config

---
 example/ck_tile/03_gemm/gemm_utils.hpp     | 79 +++++++++++++++++-----
 example/ck_tile/03_gemm/universal_gemm.cpp | 32 +++++----
 2 files changed, 81 insertions(+), 30 deletions(-)

diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp
index 78f42e50f3..632a0b75d9 100644
--- a/example/ck_tile/03_gemm/gemm_utils.hpp
+++ b/example/ck_tile/03_gemm/gemm_utils.hpp
@@ -55,13 +55,13 @@ struct GemmConfigBase
     static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave;
 };
 
-
+template<typename PrecType>
 struct GemmConfig_1 : public GemmConfigBase
 {
     // Memory friendly for Interwave scheduler
     static constexpr ck_tile::index_t M_Tile = 128;
     static constexpr ck_tile::index_t N_Tile = 32;
-    static constexpr ck_tile::index_t K_Tile = 64;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
 
     static constexpr ck_tile::index_t M_Warp = 4;
     static constexpr ck_tile::index_t N_Warp = 1;
@@ -69,18 +69,19 @@ struct GemmConfig_1 : public GemmConfigBase
 
     static constexpr ck_tile::index_t M_Warp_Tile = 32;
     static constexpr ck_tile::index_t N_Warp_Tile = 32;
-    static constexpr ck_tile::index_t K_Warp_Tile = 8;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 8 : 16;
 
     static constexpr bool DoubleSmemBuffer = false;
     static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY;
 };
 
+template<typename PrecType>
 struct GemmConfig_1_1 : public GemmConfigBase
 {
     // Memory friendly for Interwave scheduler
     static constexpr ck_tile::index_t M_Tile = 128;
     static constexpr ck_tile::index_t N_Tile = 32;
-    static constexpr ck_tile::index_t K_Tile = 64;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
 
     static constexpr ck_tile::index_t M_Warp = 4;
     static constexpr ck_tile::index_t N_Warp = 1;
@@ -88,19 +89,20 @@ struct GemmConfig_1_1 : public GemmConfigBase
 
     static constexpr ck_tile::index_t M_Warp_Tile = 32;
     static constexpr ck_tile::index_t N_Warp_Tile = 32;
-    static constexpr ck_tile::index_t K_Warp_Tile = 8;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 8 : 16;
 
     static constexpr bool DoubleSmemBuffer = false;
     static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY;
     static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave;
 };
 
+template<typename PrecType>
 struct GemmConfig_2 : public GemmConfigBase
 {
     // Compute friendly for Intrawave scheduler
     static constexpr ck_tile::index_t M_Tile = 128;
     static constexpr ck_tile::index_t N_Tile = 128;
-    static constexpr ck_tile::index_t K_Tile = 64;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
 
     static constexpr ck_tile::index_t M_Warp = 2;
     static constexpr ck_tile::index_t N_Warp = 2;
@@ -108,7 +110,7 @@ struct GemmConfig_2 : public GemmConfigBase
 
     static constexpr ck_tile::index_t M_Warp_Tile = 16;
     static constexpr ck_tile::index_t N_Warp_Tile = 16;
-    static constexpr ck_tile::index_t K_Warp_Tile = 32;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 32 : 128;
 
     static constexpr bool DoubleSmemBuffer = false;
     static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
@@ -116,12 +118,13 @@ struct GemmConfig_2 : public GemmConfigBase
     static constexpr int kBlockPerCu                         = 2;
 };
 
+template<typename PrecType>
 struct GemmConfig_2_1 : public GemmConfigBase
 {
     // Compute friendly for Intrawave scheduler
     static constexpr ck_tile::index_t M_Tile = 128;
     static constexpr ck_tile::index_t N_Tile = 128;
-    static constexpr ck_tile::index_t K_Tile = 64;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
 
     static constexpr ck_tile::index_t M_Warp = 2;
     static constexpr ck_tile::index_t N_Warp = 2;
@@ -129,7 +132,7 @@ struct GemmConfig_2_1 : public GemmConfigBase
 
     static constexpr ck_tile::index_t M_Warp_Tile = 16;
     static constexpr ck_tile::index_t N_Warp_Tile = 16;
-    static constexpr ck_tile::index_t K_Warp_Tile = 32;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 32 : 128;
 
     static constexpr bool DoubleSmemBuffer = false;
     static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
@@ -137,12 +140,13 @@ struct GemmConfig_2_1 : public GemmConfigBase
     static constexpr int kBlockPerCu = 2;
 };
 
+template<typename PrecType>
 struct GemmConfig_2_2 : public GemmConfigBase
 {
     // Compute friendly for Intrawave scheduler
     static constexpr ck_tile::index_t M_Tile = 256;
     static constexpr ck_tile::index_t N_Tile = 256;
-    static constexpr ck_tile::index_t K_Tile = 32;
+    static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType);
 
     static constexpr ck_tile::index_t M_Warp = 2;
     static constexpr ck_tile::index_t N_Warp = 2;
@@ -150,21 +154,40 @@ struct GemmConfig_2_2 : public GemmConfigBase
 
     static constexpr ck_tile::index_t M_Warp_Tile = 32;
     static constexpr ck_tile::index_t N_Warp_Tile = 32;
-    static constexpr ck_tile::index_t K_Warp_Tile = 16;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 16 : 64;
 
     static constexpr bool DoubleSmemBuffer = false;
     static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
-    //static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave;
-    //static constexpr int kBlockPerCu = 2;
 };
 
+template<typename PrecType>
+struct GemmConfig_2_3 : public GemmConfigBase
+{
+    // Compute friendly for Intrawave scheduler
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    static constexpr ck_tile::index_t K_Warp = 1;
+
+    static constexpr ck_tile::index_t M_Warp_Tile = 32;
+    static constexpr ck_tile::index_t N_Warp_Tile = 32;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 16 : 64;
+
+    static constexpr bool DoubleSmemBuffer = false;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
+};
+
+template<typename PrecType>
 struct GemmConfig_3 : public GemmConfigBase
 {
     // Compute friendly for Intrawave scheduler
     // Using the ping pong reader in the lds level
     static constexpr ck_tile::index_t M_Tile = 256;
     static constexpr ck_tile::index_t N_Tile = 256;
-    static constexpr ck_tile::index_t K_Tile = 32;
+    static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType);
 
     static constexpr ck_tile::index_t M_Warp = 2;
     static constexpr ck_tile::index_t N_Warp = 2;
@@ -172,19 +195,20 @@ struct GemmConfig_3 : public GemmConfigBase
 
     static constexpr ck_tile::index_t M_Warp_Tile = 32;
     static constexpr ck_tile::index_t N_Warp_Tile = 32;
-    static constexpr ck_tile::index_t K_Warp_Tile = 16;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 16 : 64;
 
     static constexpr bool DoubleSmemBuffer = true;
     static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4;
 };
 
+template<typename PrecType>
 struct GemmConfig_3_1 : public GemmConfigBase
 {
     // Compute friendly for Intrawave scheduler
     // Using the ping pong reader in the lds level
     static constexpr ck_tile::index_t M_Tile = 256;
     static constexpr ck_tile::index_t N_Tile = 256;
-    static constexpr ck_tile::index_t K_Tile = 32;
+    static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType);
 
     static constexpr ck_tile::index_t M_Warp = 2;
     static constexpr ck_tile::index_t N_Warp = 2;
@@ -192,13 +216,34 @@ struct GemmConfig_3_1 : public GemmConfigBase
 
     static constexpr ck_tile::index_t M_Warp_Tile = 32;
     static constexpr ck_tile::index_t N_Warp_Tile = 32;
-    static constexpr ck_tile::index_t K_Warp_Tile = 16;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 16 : 64;
 
     static constexpr bool DoubleSmemBuffer = true;
     static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4;
     static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave;
 };
 
+template<typename PrecType>
+struct GemmConfig_3_2 : public GemmConfigBase
+{
+    // Compute friendly for Intrawave scheduler
+    // Using the ping pong reader in the lds level
+    static constexpr ck_tile::index_t M_Tile = 256;
+    static constexpr ck_tile::index_t N_Tile = 256;
+    static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
+
+    static constexpr ck_tile::index_t M_Warp = 2;
+    static constexpr ck_tile::index_t N_Warp = 2;
+    static constexpr ck_tile::index_t K_Warp = 1;
+
+    static constexpr ck_tile::index_t M_Warp_Tile = 32;
+    static constexpr ck_tile::index_t N_Warp_Tile = 32;
+    static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 16 : 64;
+
+    static constexpr bool DoubleSmemBuffer = true;
+    static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4;
+};
+
 template <typename ADataType, typename BDataType = ADataType, typename CDataType = ADataType>
 struct GemmTypeConfig;
 
diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp
index ead1f3091e..23f34a648c 100644
--- a/example/ck_tile/03_gemm/universal_gemm.cpp
+++ b/example/ck_tile/03_gemm/universal_gemm.cpp
@@ -339,7 +339,7 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a
     }
 }
 
-template<typename GemmConfig>
+template<template<typename PreType> typename GemmConfig>
 int run_gemm_example(int argc, char* argv[])
 {
     auto [result, arg_parser] = create_args(argc, argv);
@@ -352,31 +352,35 @@ int run_gemm_example(int argc, char* argv[])
 
     if(data_type == "fp16")
     {
-        return run_gemm_example_prec_type<GemmConfig, ck_tile::half_t>(a_layout, b_layout, argc, argv);
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, ck_tile::half_t>(a_layout, b_layout, argc, argv);
     }
     else if(data_type == "bf16")
     {
-        return run_gemm_example_prec_type<GemmConfig, ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
     }
     else if(data_type == "fp8")
     {
-        return run_gemm_example_prec_type<GemmConfig, ck_tile::fp8_t, ck_tile::fp8_t, ck_tile::half_t>(
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::fp8_t>, ck_tile::fp8_t, ck_tile::fp8_t, ck_tile::half_t>(
             a_layout, b_layout, argc, argv);
     }
     else if(data_type == "bf8")
     {
-        return run_gemm_example_prec_type<GemmConfig, ck_tile::bf8_t, ck_tile::bf8_t, ck_tile::half_t>(
+        return run_gemm_example_prec_type<GemmConfig<ck_tile::bf8_t>, ck_tile::bf8_t, ck_tile::bf8_t, ck_tile::half_t>(
             a_layout, b_layout, argc, argv);
     }
-
-#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V3)
     else if(data_type == "pk_int4_t")
     {
-        // TODO: Add support for bhalf_t ADataType
-        return run_gemm_example_prec_type<GemmConfig, ck_tile::half_t, ck_tile::pk_int4_t, ck_tile::half_t>(
-            a_layout, b_layout, argc, argv);
+        // TODO: Add support for bhalf_t ADataType 
+        if constexpr(GemmConfig<ck_tile::half_t>::Pipeline == CK_TILE_PIPELINE_COMPUTE_V3)
+        {
+            return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, ck_tile::half_t, ck_tile::pk_int4_t, ck_tile::half_t>(
+                a_layout, b_layout, argc, argv);
+        }
+        else
+        {
+             throw std::runtime_error("Unsupported pipeline for this operation !!!");
+        }
     }
-#endif
     else
     {
         throw std::runtime_error("Unsupported data type for this operation !!!");
@@ -390,10 +394,12 @@ int main(int argc, char* argv[])
         run_gemm_example<GemmConfig_1>(argc, argv);
         run_gemm_example<GemmConfig_1_1>(argc, argv);
         run_gemm_example<GemmConfig_2>(argc, argv);
-        //run_gemm_example<GemmConfig_2_1>(argc, argv);
+        //run_gemm_example<GemmConfig_2_1>(argc, argv); xx
         run_gemm_example<GemmConfig_2_2>(argc, argv);
+        run_gemm_example<GemmConfig_2_3>(argc, argv);
         run_gemm_example<GemmConfig_3>(argc, argv);
-        //run_gemm_example<GemmConfig_3_1>(argc, argv);
+        //run_gemm_example<GemmConfig_3_1>(argc, argv); xx
+        run_gemm_example<GemmConfig_3_2>(argc, argv);
     }
     catch(const std::runtime_error& e)
     {