From 5abb073bffff7571577da7616748468bbaf60e51 Mon Sep 17 00:00:00 2001
From: Amily Wu <amilywu2@amd.com>
Date: Mon, 12 Aug 2024 18:37:37 -0500
Subject: [PATCH] Add class for args.mode

---
 tuning/autotune.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tuning/autotune.py b/tuning/autotune.py
index a6e4ac8..5139117 100755
--- 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)"