Add persistent-dp option to streamk example

Author: CongMa13

example/ck_tile/40_streamk_gemm/README.md

@@ -23,6 +23,7 @@ args:
           -b_layout    tensor B data layout (default: C)
           -c_layout    tensor C data layout (default: R)
 -reduction_strategy    strategy for storing results in C tensor. atomic/reduction (default:atomic)
+     -persistent_dp    persistent strategy for data-parallel section. 0. Non-persistent, 1 persistent.")
           -stride_a    tensor A stride (default:0)
           -stride_b    tensor B stride (default:0)
           -stride_c    tensor C stride (default:0)

example/ck_tile/40_streamk_gemm/gemm_utils.hpp

@@ -18,7 +18,6 @@ struct GemmConfigBase
 
     static constexpr bool TransposeC            = false;
     static constexpr bool UseStructuredSparsity = false;
-    static constexpr bool Persistent            = false;
 
     static constexpr int kBlockPerCu                = 1;
     static constexpr auto Scheduler                 = ck_tile::GemmPipelineScheduler::Intrawave;
@@ -27,7 +26,7 @@ struct GemmConfigBase
     static constexpr bool DoubleSmemBuffer          = false;
 };
 
-template <typename PrecType>
+template <typename PrecType, bool Persistent_>
 struct GemmConfigMemoryInterwave : public GemmConfigBase
 {
     static constexpr ck_tile::index_t M_Tile = 256;
@@ -42,7 +41,8 @@ struct GemmConfigMemoryInterwave : public GemmConfigBase
     static constexpr ck_tile::index_t N_Warp_Tile = 32;
     static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 8 : 16;
 
-    static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave;
+    static constexpr bool Persistent = Persistent_;
+    static constexpr auto Scheduler  = ck_tile::GemmPipelineScheduler::Intrawave;
 };
 
 template <typename ADataType_, typename BDataType_ = ADataType_, typename CDataType_ = ADataType_>
@@ -87,6 +87,9 @@ auto create_args(int argc, char* argv[])
         .insert("reduction_strategy",
                 "atomic",
                 "strategy for storing results in C tensor - atomic/reduction")
+        .insert("persistent_dp",
+                "0",
+                "0. Non-persistent data-parallel section, 1 Fully persistent kernel.")
         .insert("stride_a", "0", "Tensor A stride")
         .insert("stride_b", "0", "Tensor B stride")
         .insert("stride_c", "0", "Tensor C stride")

example/ck_tile/40_streamk_gemm/run_gemm_example.inc

@@ -275,8 +275,8 @@ int run_gemm_example_with_layouts(int argc,
               << " B_Type=" << DataTypeTraits<BDataType>::name
               << " C_Type=" << DataTypeTraits<CDataType>::name
               << " reduction_strategy=" << arg_parser.get_str("reduction_strategy") << " "
-              << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
-              << std::endl;
+              << " persistent_dp=" << arg_parser.get_str("persistent_dp") << " " << ave_time
+              << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl;
 
     bool pass = true;
 

example/ck_tile/40_streamk_gemm/streamk_gemm_basic.cpp

@@ -169,7 +169,7 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a
     return 0;
 }
 
-template <template <typename PreType> typename GemmConfig>
+template <template <typename PreType, bool Persistent_> typename GemmConfig>
 int run_gemm_example(int argc, char* argv[])
 {
     auto [result, arg_parser] = create_args(argc, argv);
@@ -179,18 +179,35 @@ int run_gemm_example(int argc, char* argv[])
     std::string data_type = arg_parser.get_str("prec");
     std::string a_layout  = arg_parser.get_str("a_layout");
     std::string b_layout  = arg_parser.get_str("b_layout");
+    auto persistent_dp    = arg_parser.get_bool("persistent_dp");
 
     if(data_type == "bf16")
     {
         using TypeConfig = StreamKGemmTypeConfig<ck_tile::bf16_t>;
-        return run_gemm_example_prec_type<GemmConfig<ck_tile::bf16_t>, TypeConfig>(
-            a_layout, b_layout, argc, argv);
+        if(persistent_dp)
+        {
+            return run_gemm_example_prec_type<GemmConfig<ck_tile::bf16_t, true>, TypeConfig>(
+                a_layout, b_layout, argc, argv);
+        }
+        else
+        {
+            return run_gemm_example_prec_type<GemmConfig<ck_tile::bf16_t, false>, TypeConfig>(
+                a_layout, b_layout, argc, argv);
+        }
     }
     else if(data_type == "fp16")
     {
         using TypeConfig = StreamKGemmTypeConfig<ck_tile::half_t>;
-        return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, TypeConfig>(
-            a_layout, b_layout, argc, argv);
+        if(persistent_dp)
+        {
+            return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t, true>, TypeConfig>(
+                a_layout, b_layout, argc, argv);
+        }
+        else
+        {
+            return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t, false>, TypeConfig>(
+                a_layout, b_layout, argc, argv);
+        }
     }
     else
     {