[Tuner] Add support for flag files in the simple tuner (#786)

This generalizes the simple tuner such that it can be used for
real-world models without having to modify the python code.

Author: kuhar

tuner/examples/simple/README.md

@@ -29,17 +29,20 @@ For an initial trial to test the tuning loop, use:
 cd ../../
 python -m examples.simple examples/simple/double_mmt.mlir \
     examples/simple/tmp/mmt_benchmark.mlir \
+    --simple-compile-flags-file=examples/simple/compile_flags.txt \
+    --simple-model-benchmark-flags-file=examples/simple/model_benchmark_flags.txt \
     --devices=hip://0 --num-candidates=30 \
     --simple-num-dispatch-candidates=5 --simple-num-model-candidates=3 \
 ```
 
 ### Basic Usage
 ```shell
 python -m examples.simple <model_file_path> <benchmark_file_path> \
-    --devices=hip://0 --num-candidates=1024 \
-    --test-num-dispatch-candidates=<num_dispatch_candidates> \
-    --test-num-model-candidates=<num_model_candidates> \
-    --test-hip-target=<hip_target> \
+    --devices=hip://0,hip://1 --num-candidates=1024 \
+    --simple-compile-flags-file=<compile_flags_path> \
+    --simple-model-benchmark-flags-file=<model_benchmark_flags_path> \
+    --simple-num-dispatch-candidates=<num_dispatch_candidates> \
+    --simple-num-model-candidates=<num_model_candidates> \
     --num-candidates=<num_generated_candidates> \
     --codegen-pipeline=<codegen_pipeline>
 ```

tuner/examples/simple/compile_flags.txt

@@ -0,0 +1,2 @@
+--iree-hal-target-backends=rocm
+--iree-hip-target=gfx942

tuner/examples/simple/model_benchmark_flags.txt

@@ -0,0 +1,3 @@
+--device_allocator=caching
+--input=2048x2048xf16
+--input=2048x2048xf16

tuner/examples/simple/simple_tuner.py

@@ -13,8 +13,8 @@
 class SimpleTuner(libtuner.TuningClient):
     def __init__(self, tuner_context: libtuner.TunerContext):
         super().__init__(tuner_context)
-        self.compile_flags = ["--compile-from=executable-sources"]
-        self.benchmark_flags = ["--benchmark_repetitions=3", "--input=1"]
+        self.compile_flags: list[str] = []
+        self.benchmark_flags: list[str] = []
 
     def get_iree_compile_flags(self) -> list[str]:
         return self.compile_flags
@@ -26,6 +26,14 @@ def get_benchmark_timeout_s(self) -> int:
         return 10
 
 
+def read_flags_file(flags_file: str) -> list[str]:
+    if not flags_file:
+        return []
+
+    with open(flags_file) as file:
+        return file.read().splitlines()
+
+
 def main():
     # Custom arguments for the example tuner file.
     parser = argparse.ArgumentParser(description="Autotune sample script")
@@ -46,10 +54,16 @@ def main():
         help="Number of model candidates to produce after tuning.",
     )
     client_args.add_argument(
-        "--simple-hip-target",
+        "--simple-compile-flags-file",
+        type=str,
+        default="",
+        help="Path to the flags file for iree-compile.",
+    )
+    client_args.add_argument(
+        "--simple-model-benchmark-flags-file",
         type=str,
-        default="gfx942",
-        help="Hip target for tuning.",
+        default="",
+        help="Path to the flags file for iree-benchmark-module for model benchmarking.",
     )
     # Remaining arguments come from libtuner
     args = libtuner.parse_arguments(parser)
@@ -69,6 +83,11 @@ def main():
         libtuner.validate_devices(args.devices)
         print("Validation successful!\n")
 
+    compile_flags: list[str] = read_flags_file(args.simple_compile_flags_file)
+    model_benchmark_flags: list[str] = read_flags_file(
+        args.simple_model_benchmark_flags_file
+    )
+
     print("Generating candidate tuning specs...")
     with TunerContext() as tuner_context:
         simple_tuner = SimpleTuner(tuner_context)
@@ -80,13 +99,17 @@ def main():
             return
 
         print("Compiling dispatch candidates...")
+        simple_tuner.compile_flags = compile_flags + [
+            "--compile-from=executable-sources"
+        ]
         compiled_candidates = libtuner.compile(
             args, path_config, candidates, candidate_trackers, simple_tuner
         )
         if stop_after_phase == libtuner.ExecutionPhases.compile_dispatches:
             return
 
         print("Benchmarking compiled dispatch candidates...")
+        simple_tuner.benchmark_flags = ["--input=1", "--benchmark_repetitions=3"]
         top_candidates = libtuner.benchmark(
             args,
             path_config,
@@ -99,10 +122,7 @@ def main():
             return
 
         print("Compiling models with top candidates...")
-        simple_tuner.compile_flags = [
-            "--iree-hal-target-backends=rocm",
-            f"--iree-hip-target={args.simple_hip_target}",
-        ]
+        simple_tuner.compile_flags = compile_flags
         compiled_model_candidates = libtuner.compile(
             args,
             path_config,
@@ -115,11 +135,7 @@ def main():
             return
 
         print("Benchmarking compiled model candidates...")
-        simple_tuner.benchmark_flags = [
-            "--benchmark_repetitions=3",
-            "--input=2048x2048xf16",
-            "--input=2048x2048xf16",
-        ]
+        simple_tuner.benchmark_flags = model_benchmark_flags
         top_model_candidates = libtuner.benchmark(
             args,
             path_config,