[Mlir-commits] [mlir] [MLIR][Benchmark] Refactor setup_passes function (PR #85456)

[llvmlistbot at llvm.org]

https://github.com/morphide created https://github.com/llvm/llvm-project/pull/85456

None

>From d46d84d9a773cd771e5370bc12be8090d466d8e3 Mon Sep 17 00:00:00 2001
From: morphide <163585062+morphide at users.noreply.github.com>
Date: Fri, 15 Mar 2024 19:42:15 +0000
Subject: [PATCH] [MLIR][Benchmark] Refactor setup_passes function

---
 mlir/benchmark/python/common.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlir/benchmark/python/common.py b/mlir/benchmark/python/common.py
index 0b924172c8dad7..6ac19e50f93a4d 100644
--- a/mlir/benchmark/python/common.py
+++ b/mlir/benchmark/python/common.py
@@ -9,11 +9,11 @@
 from mlir.passmanager import PassManager
 
 
-def setup_passes(mlir_module):
-    """Setup pass pipeline parameters for benchmark functions."""
-    opt = (
-        "parallelization-strategy=none"
-    )
+def setup_passes(mlir_module, optimization_options=None):
+    if optimization_options is None:
+        optimization_options = {"parallelization-strategy": "none"}
+
+    opt = ", ".join([f"{key}={value}" for key, value in optimization_options.items()])
     pipeline = f"builtin.module(sparsifier{{{opt}}})"
     PassManager.parse(pipeline).run(mlir_module.operation)