[Mlir-commits] [mlir] 5da5483 - [mlir][benchmark] Fix import in sparse benchmark.

[Ingo Müller]

Author: Ingo Müller
Date: 2022-07-15T07:15:51Z
New Revision: 5da5483ffb407d6ad1ffe3108b6459d0499c91ab

URL: https://github.com/llvm/llvm-project/commit/5da5483ffb407d6ad1ffe3108b6459d0499c91ab
DIFF: https://github.com/llvm/llvm-project/commit/5da5483ffb407d6ad1ffe3108b6459d0499c91ab.diff

LOG: [mlir][benchmark] Fix import in sparse benchmark.

The benchmark currently fails to run because it cannot find the `func`
symbol when using a `FuncOp`. I suppose that the breakage was introduced
by the extraction of the func dialect from the builtin dialect that
wasn't reflected in the benchmark yet.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D129738

Added: 
    

Modified: 
    mlir/benchmark/python/benchmark_sparse.py
    mlir/benchmark/python/common.py

Removed: 
    


################################################################################
diff  --git a/mlir/benchmark/python/benchmark_sparse.py b/mlir/benchmark/python/benchmark_sparse.py
index a96dbb2be8ae0..6d7a39690734c 100644
--- a/mlir/benchmark/python/benchmark_sparse.py
+++ b/mlir/benchmark/python/benchmark_sparse.py
@@ -10,7 +10,7 @@
 
 from mlir import ir
 from mlir import runtime as rt
-from mlir.dialects import builtin
+from mlir.dialects import func
 from mlir.dialects.linalg.opdsl import lang as dsl
 from mlir.execution_engine import ExecutionEngine
 

diff  --git a/mlir/benchmark/python/common.py b/mlir/benchmark/python/common.py
index 0c241c6fbb431..3d80892b0fb86 100644
--- a/mlir/benchmark/python/common.py
+++ b/mlir/benchmark/python/common.py
@@ -5,7 +5,6 @@
 
 from mlir import ir
 from mlir.dialects import arith
-from mlir.dialects import builtin
 from mlir.dialects import func
 from mlir.dialects import memref
 from mlir.dialects import scf
@@ -66,25 +65,26 @@ def emit_timer_func() -> func.FuncOp:
     return nanoTime
 
 
-def emit_benchmark_wrapped_main_func(func, timer_func):
+def emit_benchmark_wrapped_main_func(kernel_func, timer_func):
     """Takes a function and a timer function, both represented as FuncOp
     objects, and returns a new function. This new function wraps the call to
     the original function between calls to the timer_func and this wrapping
     in turn is executed inside a loop. The loop is executed
-    len(func.type.results) times. This function can be used to create a
-    "time measuring" variant of a function.
+    len(kernel_func.type.results) times. This function can be used to
+    create a "time measuring" variant of a function.
     """
     i64_type = ir.IntegerType.get_signless(64)
     memref_of_i64_type = ir.MemRefType.get([-1], i64_type)
     wrapped_func = func.FuncOp(
         # Same signature and an extra buffer of indices to save timings.
         "main",
-        (func.arguments.types + [memref_of_i64_type], func.type.results),
+        (kernel_func.arguments.types + [memref_of_i64_type],
+         kernel_func.type.results),
         visibility="public"
     )
     wrapped_func.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
 
-    num_results = len(func.type.results)
+    num_results = len(kernel_func.type.results)
     with ir.InsertionPoint(wrapped_func.add_entry_block()):
         timer_buffer = wrapped_func.arguments[-1]
         zero = arith.ConstantOp.create_index(0)
@@ -95,7 +95,7 @@ def emit_benchmark_wrapped_main_func(func, timer_func):
         with ir.InsertionPoint(loop.body):
             start = func.CallOp(timer_func, [])
             call = func.CallOp(
-                func,
+                kernel_func,
                 wrapped_func.arguments[:-num_results - 1] + loop.inner_iter_args
             )
             end = func.CallOp(timer_func, [])