Add class for args.mode

nod-ai · Aug 12, 2024 · 5abb073 · 5abb073
1 parent e9e5d0f
commit 5abb073
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tuning/autotune.py b/tuning/autotune.py
@@ -144,6 +144,11 @@ def get_exe_format(self, path: Path) -> str:
         return f"./{path.as_posix()}"
 
 
+class CompilationMode(Enum):
+    DEFAULT = "default"
+    WINOGRAD = "winograd"
+
+
 @dataclass
 class TaskTuple:
     args: argparse.Namespace
@@ -353,7 +358,7 @@ def parse_arguments() -> argparse.Namespace:
 
     # Required arguments
     parser.add_argument(
-        "mode", choices=["default", "winograd"], help="Compilation mode"
+        "mode", choices=[m.value for m in CompilationMode], help="Compilation mode"
     )
     parser.add_argument(
         "input_file", type=Path, help="Path to the input benchmark file (.mlir)"