From 5aeee759c9c8377e82ceaa0e9e91b570f91d2078 Mon Sep 17 00:00:00 2001
From: Amily Wu <amilywu2@amd.com>
Date: Thu, 22 Aug 2024 11:58:08 -0500
Subject: [PATCH] Polish args.mode

---
 tuning/libtuner.py       | 14 +++++++++++---
 tuning/punet_autotune.py |  6 +++---
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/tuning/libtuner.py b/tuning/libtuner.py
index 7fb86e3..95ce4e8 100755
--- a/tuning/libtuner.py
+++ b/tuning/libtuner.py
@@ -307,18 +307,26 @@ class ExecutionPhases(str, Enum):
     benchmark_models = "benchmark-models"
 
 
+class CompilationMode(str, Enum):
+    default = "default"
+    winograd = "winograd"
+
+
 def parse_arguments() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Autotune script")
 
     # Required arguments
-    parser.add_argument(
-        "mode", choices=["default", "winograd"], help="Compilation mode"
-    )
     parser.add_argument(
         "input_file", type=Path, help="Path to the input benchmark file (.mlir)"
     )
 
     # General options
+    parser.add_argument(
+        "--mode",
+        choices=[x.value for x in CompilationMode],
+        default=CompilationMode.default,
+        help="Compilation mode",
+    )
     parser.add_argument(
         "--verbose", "-v", action="store_true", help="Enable verbose output to stdout"
     )
diff --git a/tuning/punet_autotune.py b/tuning/punet_autotune.py
index c281bbf..0c94139 100644
--- a/tuning/punet_autotune.py
+++ b/tuning/punet_autotune.py
@@ -11,17 +11,17 @@
 """
 Sample Usage:
 
-python punet_autotune.py winograd 1286.mlir --lhs-dims=bmk --rhs-dims=bkn --tile-dims=*mnk --devices=hip://0,hip://1 --num-candidates=64
+python punet_autotune.py 1286.mlir --lhs-dims=bmk --rhs-dims=bkn --tile-dims=*mnk --devices=hip://0,hip://1 --num-candidates=64
 
 
 Recommended Trial Run:
 
-python punet_autotune.py winograd 1286.mlir --num-candidates=1
+python punet_autotune.py 1286.mlir --num-candidates=1
 
 
 Dry Run Test (no gpu requried):
 
-python punet_autotune.py winograd 1286.mlir --num-candidates=64 --num-model-candidates=10 --dry-run
+python punet_autotune.py 1286.mlir --num-candidates=64 --num-model-candidates=10 --dry-run
 
 """