From 03c252551cb53e64051aa6bbbe6e420538dc815f Mon Sep 17 00:00:00 2001
From: Harisankar Sadasivan <hari.sadasivan@amd.com>
Date: Tue, 25 Jun 2024 18:32:17 +0000
Subject: [PATCH] corrected arg parsing for streamk and universal streamk

---
 example/01_gemm/common.hpp | 54 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 52 insertions(+), 2 deletions(-)

diff --git a/example/01_gemm/common.hpp b/example/01_gemm/common.hpp
index 7fd5772308..f6705829eb 100755
--- a/example/01_gemm/common.hpp
+++ b/example/01_gemm/common.hpp
@@ -136,6 +136,57 @@ bool parse_cmd_args<ProblemSize>(int argc,
     return true;
 }
 
+template <>
+bool parse_cmd_args<ProblemSizeStreamK_universal>(int argc,
+                                                  char* argv[],
+                                                  ProblemSizeStreamK& problem_size,
+                                                  ExecutionConfig& config)
+{
+    if(argc == 1)
+    {
+        // use default case
+    }
+    else if(argc == 4)
+    {
+        config.do_verification = std::stoi(argv[1]);
+        config.init_method     = std::stoi(argv[2]);
+        config.time_kernel     = std::stoi(argv[3]);
+    }
+    else if(argc >= 10)
+    {
+        config.do_verification = std::stoi(argv[1]);
+        config.init_method     = std::stoi(argv[2]);
+        config.time_kernel     = std::stoi(argv[3]);
+
+        problem_size.M = std::stoi(argv[4]);
+        problem_size.N = std::stoi(argv[5]);
+        problem_size.K = std::stoi(argv[6]);
+
+        problem_size.StrideA = std::stoi(argv[7]);
+        problem_size.StrideB = std::stoi(argv[8]);
+        problem_size.StrideC = std::stoi(argv[9]);
+
+        if(argc >= 11)
+        {
+            problem_size.Streamk_sel = std::stoi(argv[10]);
+            problem_size.Grid_size   = std::stoi(argv[11]);
+        }
+    }
+    else
+    {
+        std::cerr << "arg1: verification (0=no, 1=yes)" << std::endl
+                  << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)"
+                  << std::endl
+                  << "arg3: time kernel (0=no, 1=yes)" << std::endl
+                  << "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideC" << std::endl
+                  << "arg10: stream-k select (0: all DP, 1: 1-tile SK, 2: 2-tile SK)"
+                  << "\narg11: Grid_size(-1 for max occupancy)" << std::endl;
+        return false;
+    }
+
+    return true;
+}
+
 template <>
 bool parse_cmd_args<ProblemSizeStreamK>(int argc,
                                         char* argv[],
@@ -168,8 +219,7 @@ bool parse_cmd_args<ProblemSizeStreamK>(int argc,
 
         if(argc >= 11)
         {
-            problem_size.Streamk_sel = std::stoi(argv[10]);
-            problem_size.Grid_size   = std::stoi(argv[11]);
+            problem_size.NumSKBlocks = std::stoi(argv[10]);
         }
     }
     else