From b920696ed393fc887245bef3d710f2dcb20b65eb Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub@nod-labs.com>
Date: Thu, 9 Jan 2025 20:32:02 -0500
Subject: [PATCH] [Tuner] Fix large model benchmarking (#808)

* Add model-specific benchmark timeout
* Fix benchmark argument parsing to allow for `=` in command line
argument values
* Don't print candidate trackers at the very end (too much noise)
* Always promote operands
---
 tuner/examples/simple/simple_tuner.py | 11 +++++------
 tuner/tuner/dispatch_constraints.py   |  3 ++-
 tuner/tuner/libtuner.py               |  4 ++--
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tuner/examples/simple/simple_tuner.py b/tuner/examples/simple/simple_tuner.py
index d4ec089fe..3b8af79f9 100644
--- a/tuner/examples/simple/simple_tuner.py
+++ b/tuner/examples/simple/simple_tuner.py
@@ -15,7 +15,8 @@ def __init__(self, tuner_context: libtuner.TunerContext):
         super().__init__(tuner_context)
         self.compile_flags: list[str] = []
         self.benchmark_flags: list[str] = []
-        self.compile_timeout: int = 10
+        self.compile_timeout: int = 16
+        self.benchmark_timeout: int = 16
 
     def get_iree_compile_flags(self) -> list[str]:
         return self.compile_flags
@@ -27,7 +28,7 @@ def get_iree_benchmark_module_flags(self) -> list[str]:
         return self.benchmark_flags
 
     def get_benchmark_timeout_s(self) -> int:
-        return 10
+        return self.benchmark_timeout
 
 
 def read_flags_file(flags_file: str) -> list[str]:
@@ -127,7 +128,7 @@ def main():
 
         print("Compiling models with top candidates...")
         simple_tuner.compile_flags = compile_flags
-        simple_tuner.compile_timeout = 60
+        simple_tuner.compile_timeout = 120
         compiled_model_candidates = libtuner.compile(
             args,
             path_config,
@@ -141,6 +142,7 @@ def main():
 
         print("Benchmarking compiled model candidates...")
         simple_tuner.benchmark_flags = model_benchmark_flags
+        simple_tuner.benchmark_timeout = 60
         top_model_candidates = libtuner.benchmark(
             args,
             path_config,
@@ -154,6 +156,3 @@ def main():
 
         print("Check the detailed execution logs in:")
         print(path_config.run_log.resolve())
-
-    for candidate in candidate_trackers:
-        libtuner.logging.debug(candidate)
diff --git a/tuner/tuner/dispatch_constraints.py b/tuner/tuner/dispatch_constraints.py
index b41d808b9..df7151002 100644
--- a/tuner/tuner/dispatch_constraints.py
+++ b/tuner/tuner/dispatch_constraints.py
@@ -309,10 +309,11 @@ def generate_compilation_infos(
         "reduction": reduction_tile_sizes,
         "subgroup_m_count": subgroup_m_count,
         "subgroup_n_count": subgroup_n_count,
+        "promote_operands": [0, 1],
     }
     if codegen_pipeline == iree_codegen.DispatchLoweringPassPipeline.LLVMGPUTileAndFuse:
         lowering_config_args["subgroup"] = subgroup_tile_sizes
-        lowering_config_args["promote_operands"] = [0, 1]
+
     lowering_config = get_lowering_config(**lowering_config_args)
 
     # Create the TranslationInfoAttr
diff --git a/tuner/tuner/libtuner.py b/tuner/tuner/libtuner.py
index 8c5b15761..9942187ec 100644
--- a/tuner/tuner/libtuner.py
+++ b/tuner/tuner/libtuner.py
@@ -497,10 +497,10 @@ def run_iree_benchmark_module_command(benchmark_pack: BenchmarkPack):
         assert flag[:2] == "--", "iree_benchmark_module_flags should begin with '--'"
         split_key_value = flag[2:].split("=")
         assert (
-            len(split_key_value) == 2
+            len(split_key_value) >= 1
         ), "iree_benchmark_module_flags should have the format --<key>=<value>"
         key = split_key_value[0]
-        value = split_key_value[1]
+        value = "=".join(split_key_value[1:])
         # Allow the tuning client to pass `--function=@func_name`.
         if key == "function":
             func_name = value