@@ -13,10 +13,9 @@ def setup_passes(mlir_module):
13
13
"""Setup pass pipeline parameters for benchmark functions."""
14
14
opt = (
15
15
"parallelization-strategy=none"
16
- " vectorization-strategy=none vl=1 enable-simd-index32=False"
17
16
)
18
- pipeline = f"sparsifier{{{ opt } }}"
19
- PassManager .parse (pipeline ).run (mlir_module )
17
+ pipeline = f"builtin.module( sparsifier{{{ opt } }}) "
18
+ PassManager .parse (pipeline ).run (mlir_module . operation )
20
19
21
20
22
21
def create_sparse_np_tensor (dimensions , number_of_elements ):
@@ -73,7 +72,7 @@ def emit_benchmark_wrapped_main_func(kernel_func, timer_func):
73
72
create a "time measuring" variant of a function.
74
73
"""
75
74
i64_type = ir .IntegerType .get_signless (64 )
76
- memref_of_i64_type = ir .MemRefType .get ([- 1 ], i64_type )
75
+ memref_of_i64_type = ir .MemRefType .get ([ir . ShapedType . get_dynamic_size () ], i64_type )
77
76
wrapped_func = func .FuncOp (
78
77
# Same signature and an extra buffer of indices to save timings.
79
78
"main" ,
@@ -86,7 +85,7 @@ def emit_benchmark_wrapped_main_func(kernel_func, timer_func):
86
85
with ir .InsertionPoint (wrapped_func .add_entry_block ()):
87
86
timer_buffer = wrapped_func .arguments [- 1 ]
88
87
zero = arith .ConstantOp .create_index (0 )
89
- n_iterations = memref .DimOp (ir . IndexType . get (), timer_buffer , zero )
88
+ n_iterations = memref .DimOp (timer_buffer , zero )
90
89
one = arith .ConstantOp .create_index (1 )
91
90
iter_args = list (wrapped_func .arguments [- num_results - 1 : - 1 ])
92
91
loop = scf .ForOp (zero , n_iterations , one , iter_args )
0 commit comments