Add support for waves_per_eu flag

The waves_per_eu flag is a optimization hint that specifies the minimum number of waves per execution unit. Signed-off-by: Harsh Menon <[email protected]>
iree-org · Dec 4, 2024 · 9bdc283 · 9bdc283
1 parent d8f0df3
commit 9bdc283
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/iree/turbine/kernel/wave/utils.py b/iree/turbine/kernel/wave/utils.py
@@ -557,6 +557,9 @@ def compile_and_invoke(
         "--iree-vm-target-truncate-unsupported-floats",
     ]
 
+    if config.get("waves_per_eu", None) is not None:
+        flags.append(f"--iree-hip-waves-per-eu={config['waves_per_eu']}")
+
     # TODO: More targets/backends support.
     if backend == "rocm":
         target = config["target"]

diff --git a/tests/kernel/wave/wave_evoformer_test.py b/tests/kernel/wave/wave_evoformer_test.py
@@ -97,6 +97,7 @@ def testEvoformerAttentionForward(
     symbols.update(get_default_scheduling_params())
 
     config = get_default_run_config()
+    config["waves_per_eu"] = 2
     if run_bench:
         config["benchmark_batch_size"] = 1000
         config["benchmark_repetitions"] = 3