[mlir][python] Allow running pass manager on any operation

`PassManager.run` is currently restricted to running on `builtin.module`
ops, but this restriction doesn't exist on the C++ side. This updates it
to take `ir.Operation/OpView` instead of `ir.Module`.

Depends on D143354

Reviewed By: mehdi_amini

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

Author: rkayaith

mlir/lib/Bindings/Python/Pass.cpp

@@ -116,16 +116,16 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
           "ValueError if the pipeline can't be parsed.")
       .def(
           "run",
-          [](PyPassManager &passManager, PyModule &module) {
+          [](PyPassManager &passManager, PyOperationBase &op) {
             MlirLogicalResult status = mlirPassManagerRunOnOp(
-                passManager.get(), mlirModuleGetOperation(module.get()));
+                passManager.get(), op.getOperation().get());
             if (mlirLogicalResultIsFailure(status))
               throw SetPyError(PyExc_RuntimeError,
                                "Failure while executing pass pipeline.");
           },
-          py::arg("module"),
-          "Run the pass manager on the provided module, throw a RuntimeError "
-          "on failure.")
+          py::arg("operation"),
+          "Run the pass manager on the provided operation, throw a "
+          "RuntimeError on failure.")
       .def(
           "__str__",
           [](PyPassManager &self) {

mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py

@@ -24,7 +24,7 @@ def __call__(self, module: ir.Module):
 
   def compile(self, module: ir.Module):
     """Compiles the module by invoking the sparse copmiler pipeline."""
-    passmanager.PassManager.parse(self.pipeline).run(module)
+    passmanager.PassManager.parse(self.pipeline).run(module.operation)
 
   def jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:
     """Wraps the module in a JIT execution engine."""

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_sparse_compiler.py

@@ -27,7 +27,7 @@ def __call__(self, module: ir.Module):
 
   def compile(self, module: ir.Module):
     """Compiles the module by invoking the sparse copmiler pipeline."""
-    passmanager.PassManager.parse(self.pipeline).run(module)
+    passmanager.PassManager.parse(self.pipeline).run(module.operation)
 
   def jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:
     """Wraps the module in a JIT execution engine."""

mlir/test/python/execution_engine.py

@@ -64,7 +64,7 @@ def testInvalidModule():
 def lowerToLLVM(module):
   pm = PassManager.parse(
       "builtin.module(convert-complex-to-llvm,finalize-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts)")
-  pm.run(module)
+  pm.run(module.operation)
   return module
 
 

mlir/test/python/integration/dialects/linalg/opsrun.py

@@ -202,7 +202,7 @@ def transform(module, boilerplate):
   pm.add("finalize-memref-to-llvm")
   pm.add("convert-func-to-llvm")
   pm.add("reconcile-unrealized-casts")
-  pm.run(mod)
+  pm.run(mod.operation)
   return mod
 
 

mlir/test/python/pass_manager.py

@@ -3,6 +3,7 @@
 import gc, sys
 from mlir.ir import *
 from mlir.passmanager import *
+from mlir.dialects.func import FuncOp
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
@@ -120,11 +121,10 @@ def testInvalidNesting():
 # CHECK-LABEL: TEST: testRun
 def testRunPipeline():
   with Context():
-    pm = PassManager.parse("builtin.module(print-op-stats{json=false})")
-    module = Module.parse(r"""func.func @successfulParse() { return }""")
-    pm.run(module)
+    pm = PassManager.parse("any(print-op-stats{json=false})")
+    func = FuncOp.parse(r"""func.func @successfulParse() { return }""")
+    pm.run(func)
 # CHECK: Operations encountered:
-# CHECK: builtin.module    , 1
 # CHECK: func.func      , 1
 # CHECK: func.return        , 1
 run(testRunPipeline)