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

[rkayaith]

Summary

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

mlir/lib/Bindings/Python/Pass.cpp

py::class_<PyPassManager>(m, "PassManager", py::module_local())
             if (mlirLogicalResultIsFailure(status))
               throw SetPyError(PyExc_ValueError, std::string(errorMsg.join()));
           },
           py::arg("pipeline"),
           "Add textual pipeline elements to the pass manager. Throws a "
           "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) {
             MlirPassManager passManager = self.get();
             PyPrintAccumulator printAccum;
             mlirPrintPassPipeline(
                 mlirPassManagerGetAsOpPassManager(passManager),
                 printAccum.getCallback(), printAccum.getUserData());
             return printAccum.join();
           },
           "Print the textual representation for this PassManager, suitable to "
           "be passed to `parse` for round-tripping.");
 }

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

def __init__(self, options: str, opt_level: int, shared_libs: Sequence[str]):
     self.shared_libs = shared_libs
 
   def __call__(self, module: ir.Module):
     """Convenience application method."""
     self.compile(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."""
     return execution_engine.ExecutionEngine(
         module, opt_level=self.opt_level, shared_libs=self.shared_libs)
 
   def compile_and_jit(self,
                       module: ir.Module) -> execution_engine.ExecutionEngine:
     """Compiles and jits the module."""
     self.compile(module)
     return self.jit(module)

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

def __init__(self, options: str, opt_level: int, shared_libs: Sequence[str]):
     self.shared_libs = shared_libs
 
   def __call__(self, module: ir.Module):
     """Convenience application method."""
     self.compile(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."""
     return execution_engine.ExecutionEngine(
         module, opt_level=self.opt_level, shared_libs=self.shared_libs)
 
   def compile_and_jit(self,
                       module: ir.Module) -> execution_engine.ExecutionEngine:
     """Compiles and jits the module."""
     self.compile(module)
     return self.jit(module)

mlir/test/python/execution_engine.py

 
 
 run(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
 
 
 # Test simple ExecutionEngine execution
 # CHECK-LABEL: TEST: testInvokeVoid
 def testInvokeVoid():
   with Context():
     module = Module.parse(r"""

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

def transform(module, boilerplate):
   pm.add("func.func(convert-math-to-llvm)")
   pm.add("func.func(convert-scf-to-cf)")
   pm.add("func.func(arith-expand)")
   pm.add("func.func(memref-expand)")
   pm.add("convert-vector-to-llvm")
   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
 
 
 def test_elemwise_builtin():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     i8 = IntegerType.get_signless(8)

mlir/test/python/pass_manager.py

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 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.
 def log(*args):
   print(*args, file=sys.stderr)
   sys.stderr.flush()
 
 def run(f):
else:
       log("Exception not produced")
 run(testInvalidNesting)
 
 
 # Verify that a pass manager can execute on IR
 # 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)