Polish args.mode

nod-ai · Aug 22, 2024 · 5aeee75 · 5aeee75
1 parent ff47eab
commit 5aeee75
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git 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
@@ -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
 
 """