[Mlir-commits] [mlir] 6cb1c0c - Add Python binding to run a PassManager on a MLIR Module

[Mehdi Amini]

Author: Mehdi Amini
Date: 2020-11-10T20:06:23Z
New Revision: 6cb1c0cae0dfcdeabb7d764b1cb3be637a227717

URL: https://github.com/llvm/llvm-project/commit/6cb1c0cae0dfcdeabb7d764b1cb3be637a227717
DIFF: https://github.com/llvm/llvm-project/commit/6cb1c0cae0dfcdeabb7d764b1cb3be637a227717.diff

LOG: Add Python binding to run a PassManager on a MLIR Module

Reviewed By: ftynse, stellaraccident

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

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/Pass.cpp
    mlir/test/Bindings/Python/pass_manager.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index 228f69b6deba..cfd9c281e734 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -60,6 +60,17 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
           "Parse a textual pass-pipeline and return a top-level PassManager "
           "that can be applied on a Module. Throw a ValueError if the pipeline "
           "can't be parsed")
+      .def(
+          "run",
+          [](PyPassManager &passManager, PyModule &module) {
+            MlirLogicalResult status =
+                mlirPassManagerRun(passManager.get(), module.get());
+            if (mlirLogicalResultIsFailure(status))
+              throw SetPyError(PyExc_RuntimeError,
+                               "Failure while executing pass pipeline.");
+          },
+          "Run the pass manager on the provided module, throw a RuntimeError "
+          "on failure.")
       .def(
           "__str__",
           [](PyPassManager &self) {

diff  --git a/mlir/test/Bindings/Python/pass_manager.py b/mlir/test/Bindings/Python/pass_manager.py
index c8e041c2eae0..62a92484f0f4 100644
--- a/mlir/test/Bindings/Python/pass_manager.py
+++ b/mlir/test/Bindings/Python/pass_manager.py
@@ -5,8 +5,7 @@
 from mlir.passmanager import *
 
 # Log everything to stderr and flush so that we have a unified stream to match
-# errors emitted by MLIR to stderr. TODO: this shouldn't be needed when
-# everything is plumbed.
+# errors/info emitted by MLIR to stderr.
 def log(*args):
   print(*args, file=sys.stderr)
   sys.stderr.flush()
@@ -52,3 +51,18 @@ def testParseFail():
     else:
       log("Exception not produced")
 run(testParseFail)
+
+
+# Verify that a pass manager can execute on IR
+# CHECK-LABEL: TEST: testRun
+def testRunPipeline():
+  with Context():
+    pm = PassManager.parse("print-op-stats")
+    module = Module.parse(r"""func @successfulParse() { return }""")
+    pm.run(module)
+# CHECK: Operations encountered:
+# CHECK: func              , 1
+# CHECK: module            , 1
+# CHECK: module_terminator , 1
+# CHECK: std.return        , 1
+run(testRunPipeline)