[mlir] Use sparse-compiler pass in sparse benchmark

[SaurabhJha]

Summary

This patch uses sparse-compiler option in mlir sparse tensor benchmark

[SaurabhJha]

Rename variable pipeline_new -> pipeline

[SaurabhJha]

mlir/benchmark/python/common.py

21–22

@wrengr @aartbik I need some help here. I haven't been able to get this to work.

If I do pipeline = f"sparse-compiler{{{opt}}}", I get Can't add pass 'LinalgBufferize' restricted to 'builtin.func' on a PassManager intended to run on 'builtin.module', did you intend to nest?.

If I wrap this in builtin.func like pipeline = f"builtin.func(sparse-compiler{{{opt}}})", I get Can't add pass 'Sparsification' restricted to 'builtin.module' on a PassManager intended to run on 'builtin.func', did you intend to nest?

If I do pipeline = f"builtin.module(sparse-compiler{{{opt}}})", I get Can't add pass 'LinalgBufferize' restricted to 'builtin.func' on a PassManager intended to run on 'builtin.module', did you intend to nest?

And finally, if I do pipeline = f"builtin.module(builtin.func(sparse-compiler{{{opt}}}))", I get Can't add pass 'Sparsification' restricted to 'builtin.module' on a PassManager intended to run on 'builtin.func', did you intend to nest?.

I tried looking at examples in the common pipeline patch but couldn't get it to work. Could you please help me set this correctly?

[rriddle]

mlir/benchmark/python/common.py

21–22

I added a comment to D117919, but it looks like the pipeline is relying on an implicitly nesting pass manager.

[wrengr]

mlir/benchmark/python/common.py

21–22

I tried a few different approaches re nesting (prior to submitting D117919), but couldn't get any of them to work for the cpu-integration tests except for the version that landed. Afaict the python bindings for pipeline parsing differ somehow from what mlir-opt uses; but it's hard to say since both the python bindings and the pipeline stuff are poorly documented. I'll take another look this week and see if I can't figure out how to fix it

[SaurabhJha]

mlir/benchmark/python/common.py

21–22

Great, thanks Wren.

[wrengr]

mlir/benchmark/python/common.py

21–22

I just uploaded D118658 which should fix this

[aartbik]

mlir/benchmark/python/common.py

19–20

we disable all, which basically means you could do without opt at all if that makes things simpler for the first version

[SaurabhJha]

mlir/benchmark/python/common.py

19–20

Unfortunately, even if I remove opt and run it like

pipeline = f"builtin.module(sparse-compiler))"
PassManager.parse(pipeline).run(mlir_module)

I still face the same issues.

[wrengr]

mlir/benchmark/python/common.py

19–20

With D118658 you should be able to use pipeline = f'sparse-compiler' (i.e., no builtin.module nesting), which is what all the python integration tests use after that differential

[SaurabhJha]

mlir/benchmark/python/common.py

19–20

Great, I will rebase my changes once your changes are in.

[wrengr]

mlir/benchmark/python/common.py

19–20

Just landed it :)

[aartbik]

mlir/benchmark/python/common.py

19–20

@SaurabhJha what is the status of this revision?

[SaurabhJha]

mlir/benchmark/python/common.py

19–20

Sorry for the delay, the last few days have been quite busy for me. I will update this today pacific time morning.

[SaurabhJha]

Integrating fixes in sparse-compiler pass. It works now.

[SaurabhJha]

Mistakenly pushed the previous version before. Fixed it now.

mlir/benchmark/python/common.py

 from mlir.dialects import std
 from mlir.passmanager import PassManager
 
 
 def setup_passes(mlir_module):
     """Setup pass pipeline parameters for benchmark functions.
     """
     opt = (
         "parallelization-strategy=0"
         " vectorization-strategy=0 vl=1 enable-simd-index32=False"
     )
-    pipeline = (
-        f"builtin.func"
-        f"(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),"
-        f"sparsification{{{opt}}},"
-        f"sparse-tensor-conversion,"
-        f"builtin.func"
-        f"(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),"
-        f"convert-scf-to-std,"
-        f"func-bufferize,"
-        f"tensor-constant-bufferize,"
-        f"builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),"
-        f"convert-vector-to-llvm"
-        f"{{reassociate-fp-reductions=1 enable-index-optimizations=1}},"
-        f"lower-affine,"
-        f"convert-memref-to-llvm,"
-        f"convert-std-to-llvm,"
-        f"reconcile-unrealized-casts"
-    )
-    PassManager.parse(pipeline).run(mlir_module)
+    pipeline_new = f"builtin.module(sparse-compiler{{{opt}}})"
+    PassManager.parse(pipeline_new).run(mlir_module)
 
 
 def create_sparse_np_tensor(dimensions, number_of_elements):
     """Constructs a numpy tensor of dimensions `dimensions` that has only a
     specific number of nonzero elements, specified by the `number_of_elements`
     argument.
     """
     tensor = np.zeros(dimensions, np.float64)