[mlir][python] Include pipeline parse errors in exception message

[rkayaith]

Summary

Currently any errors during pipeline parsing are reported to stderr.
This adds a new pipeline parsing function to the C api that reports
errors through a callback, and updates the python bindings to use it.

[rkayaith]

update description

[rkayaith]

mlir/lib/CAPI/IR/Pass.cpp

85–86

i'll fix this in D136403

[rkayaith]

add capi test for error capture, rebase

mlir/include/mlir-c/Pass.h

 
 /// Add a pass and transfer ownership to the provided mlirOpPassManager. If the
 /// pass is not a generic operation pass or matching the type of the provided
 /// PassManager, a new OpPassManager is implicitly nested under the provided
 /// PassManager.
 MLIR_CAPI_EXPORTED void
 mlirOpPassManagerAddOwnedPass(MlirOpPassManager passManager, MlirPass pass);
 
+/// Parse a sequence of textual MLIR pass pipeline elements and add them to the
+/// provided OpPassManager. If parsing fails an error message is reported using
+/// the provided callback.
+MLIR_CAPI_EXPORTED MlirLogicalResult mlirOpPassManagerAddPipeline(
+    MlirOpPassManager passManager, MlirStringRef pipelineElements,
+    MlirStringCallback callback, void *userData);
+
 /// Print a textual MLIR pass pipeline by sending chunks of the string
 /// representation and forwarding `userData to `callback`. Note that the
 /// callback may be called several times with consecutive chunks of the string.
 MLIR_CAPI_EXPORTED void mlirPrintPassPipeline(MlirOpPassManager passManager,
                                               MlirStringCallback callback,
                                               void *userData);
 
 /// Parse a textual MLIR pass pipeline and add it to the provided OpPassManager.

mlir/lib/Bindings/Python/Pass.cpp

py::class_<PyPassManager>(m, "PassManager", py::module_local())
       .def(
           "enable_verifier",
           [](PyPassManager &passManager, bool enable) {
             mlirPassManagerEnableVerifier(passManager.get(), enable);
           },
           py::arg("enable"), "Enable / disable verify-each.")
       .def_static(
           "parse",
-          [](const std::string pipeline, DefaultingPyMlirContext context) {
+          [](const std::string &pipeline, DefaultingPyMlirContext context) {
             MlirPassManager passManager = mlirPassManagerCreate(context->get());
-            MlirLogicalResult status = mlirParsePassPipeline(
+            PyPrintAccumulator errorMsg;
+            MlirLogicalResult status = mlirOpPassManagerAddPipeline(
                 mlirPassManagerGetAsOpPassManager(passManager),
-                mlirStringRefCreate(pipeline.data(), pipeline.size()));
+                mlirStringRefCreate(pipeline.data(), pipeline.size()),
+                errorMsg.getCallback(), errorMsg.getUserData());
             if (mlirLogicalResultIsFailure(status))
-              throw SetPyError(PyExc_ValueError,
-                               llvm::Twine("invalid pass pipeline '") +
-                                   pipeline + "'.");
+              throw SetPyError(PyExc_ValueError, std::string(errorMsg.join()));
             return new PyPassManager(passManager);
           },
           py::arg("pipeline"), py::arg("context") = py::none(),
           "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",

mlir/lib/CAPI/IR/Pass.cpp

void mlirPassManagerAddOwnedPass(MlirPassManager passManager, MlirPass pass) {
   unwrap(passManager)->addPass(std::unique_ptr<Pass>(unwrap(pass)));
 }
 
 void mlirOpPassManagerAddOwnedPass(MlirOpPassManager passManager,
                                    MlirPass pass) {
   unwrap(passManager)->addPass(std::unique_ptr<Pass>(unwrap(pass)));
 }
 
+MlirLogicalResult mlirOpPassManagerAddPipeline(MlirOpPassManager passManager,
+                                               MlirStringRef pipelineElements,
+                                               MlirStringCallback callback,
+                                               void *userData) {
+  detail::CallbackOstream stream(callback, userData);
+  return wrap(parsePassPipeline(unwrap(pipelineElements), *unwrap(passManager),
+                                stream));
+}
+
 void mlirPrintPassPipeline(MlirOpPassManager passManager,
                            MlirStringCallback callback, void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(passManager)->printAsTextualPipeline(stream);
 }
 
 MlirLogicalResult mlirParsePassPipeline(MlirOpPassManager passManager,
                                         MlirStringRef pipeline) {
   // TODO: errors are sent to std::errs() at the moment, we should pass in a
   // stream and redirect to a diagnostic.
   return wrap(mlir::parsePassPipeline(unwrap(pipeline), *unwrap(passManager)));
 }
 
 //===----------------------------------------------------------------------===//
 // External Pass API.
 //===----------------------------------------------------------------------===//
 
 namespace mlir {

mlir/test/CAPI/pass.c

"}"));
   mlirContextDestroy(ctx);
 }
 
 static void printToStderr(MlirStringRef str, void *userData) {
   (void)userData;
   fwrite(str.data, 1, str.length, stderr);
 }
 
+static void dontPrint(MlirStringRef str, void *userData) {
+  (void)str;
+  (void)userData;
+}
+
 void testPrintPassPipeline() {
   MlirContext ctx = mlirContextCreate();
   MlirPassManager pm = mlirPassManagerCreate(ctx);
   // Populate the pass-manager
   MlirOpPassManager nestedModulePm = mlirPassManagerGetNestedUnder(
       pm, mlirStringRefCreateFromCString("builtin.module"));
   MlirOpPassManager nestedFuncPm = mlirOpPassManagerGetNestedUnder(
       nestedModulePm, mlirStringRefCreateFromCString("func.func"));

 
 void testParsePassPipeline() {
   MlirContext ctx = mlirContextCreate();
   MlirPassManager pm = mlirPassManagerCreate(ctx);
   // Try parse a pipeline.
   MlirLogicalResult status = mlirParsePassPipeline(
       mlirPassManagerGetAsOpPassManager(pm),
       mlirStringRefCreateFromCString(
-          "builtin.module(func.func(print-op-stats{json=false}),"
-          " func.func(print-op-stats{json=false}))"));
+          "builtin.module(func.func(print-op-stats{json=false}))"));
   // Expect a failure, we haven't registered the print-op-stats pass yet.
   if (mlirLogicalResultIsSuccess(status)) {
     fprintf(
         stderr,
         "Unexpected success parsing pipeline without registering the pass\n");
     exit(EXIT_FAILURE);
   }
   // Try again after registrating the pass.
   mlirRegisterTransformsPrintOpStats();
   status = mlirParsePassPipeline(
       mlirPassManagerGetAsOpPassManager(pm),
       mlirStringRefCreateFromCString(
-          "builtin.module(func.func(print-op-stats{json=false}),"
-          " func.func(print-op-stats{json=false}))"));
+          "builtin.module(func.func(print-op-stats{json=false}))"));
   // Expect a failure, we haven't registered the print-op-stats pass yet.
   if (mlirLogicalResultIsFailure(status)) {
     fprintf(stderr,
             "Unexpected failure parsing pipeline after registering the pass\n");
     exit(EXIT_FAILURE);
   }
 
-  //      CHECK: Round-trip: builtin.module(builtin.module(
-  // CHECK-SAME:   func.func(print-op-stats{json=false}),
-  // CHECK-SAME:   func.func(print-op-stats{json=false})
-  // CHECK-SAME: ))
+  //      CHECK: Round-trip: builtin.module(
+  // CHECK-SAME:   builtin.module(func.func(print-op-stats{json=false}))
+  // CHECK-SAME: )
   fprintf(stderr, "Round-trip: ");
   mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
                         NULL);
   fprintf(stderr, "\n");
+
+  // Try appending a pass:
+  status = mlirOpPassManagerAddPipeline(
+      mlirPassManagerGetAsOpPassManager(pm),
+      mlirStringRefCreateFromCString("func.func(print-op-stats{json=false})"),
+      printToStderr, NULL);
+  if (mlirLogicalResultIsFailure(status)) {
+    fprintf(stderr, "Unexpected failure appending pipeline\n");
+    exit(EXIT_FAILURE);
+  }
+  //      CHECK: Appended: builtin.module(
+  // CHECK-SAME:   builtin.module(func.func(print-op-stats{json=false})),
+  // CHECK-SAME:   func.func(print-op-stats{json=false})
+  // CHECK-SAME: )
+  fprintf(stderr, "Appended: ");
+  mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
+                        NULL);
+  fprintf(stderr, "\n");
+
+  mlirPassManagerDestroy(pm);
+  mlirContextDestroy(ctx);
+}
+
+void testParseErrorCapture() {
+  // CHECK-LABEL: testParseErrorCapture:
+  fprintf(stderr, "\nTEST: testParseErrorCapture:\n");
+
+  MlirContext ctx = mlirContextCreate();
+  MlirPassManager pm = mlirPassManagerCreate(ctx);
+  MlirOpPassManager opm = mlirPassManagerGetAsOpPassManager(pm);
+  MlirStringRef invalidPipeline = mlirStringRefCreateFromCString("invalid");
+
+  // CHECK: mlirOpPassManagerAddPipeline:
+  // CHECK: 'invalid' does not refer to a registered pass or pass pipeline
+  fprintf(stderr, "mlirOpPassManagerAddPipeline:\n");
+  if (mlirLogicalResultIsSuccess(mlirOpPassManagerAddPipeline(
+          opm, invalidPipeline, printToStderr, NULL)))
+    exit(EXIT_FAILURE);
+  fprintf(stderr, "\n");
+
+  // Make sure all output is going through the callback.
+  // CHECK: dontPrint: <>
+  fprintf(stderr, "dontPrint: <");
+  if (mlirLogicalResultIsSuccess(
+          mlirOpPassManagerAddPipeline(opm, invalidPipeline, dontPrint, NULL)))
+    exit(EXIT_FAILURE);
+  fprintf(stderr, ">\n");
+
   mlirPassManagerDestroy(pm);
   mlirContextDestroy(ctx);
 }
 
 struct TestExternalPassUserData {
   int constructCallCount;
   int destructCallCount;
   int initializeCallCount;
"}"));
   mlirContextDestroy(ctx);
 }
 
 int main() {
   testRunPassOnModule();
   testRunPassOnNestedModule();
   testPrintPassPipeline();
   testParsePassPipeline();
+  testParseErrorCapture();
   testExternalPass();
   return 0;
 }

mlir/test/python/pass_manager.py

 
 # Verify successful round-trip.
 # CHECK-LABEL: TEST: testParseSuccess
 def testParseSuccess():
   with Context():
     # An unregistered pass should not parse.
     try:
       pm = PassManager.parse("builtin.module(func.func(not-existing-pass{json=false}))")
-      # TODO: this error should be propagate to Python but the C API does not help right now.
-      # CHECK: error: 'not-existing-pass' does not refer to a registered pass or pass pipeline
     except ValueError as e:
-      # CHECK: ValueError exception: invalid pass pipeline 'builtin.module(func.func(not-existing-pass{json=false}))'.
+      # CHECK: ValueError exception: {{.+}} 'not-existing-pass' does not refer to a registered pass
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 
     # A registered pass should parse successfully.
     pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")
     # CHECK: Roundtrip: builtin.module(builtin.module(func.func(print-op-stats{json=false})))
     log("Roundtrip: ", pm)
 run(testParseSuccess)
 
 # Verify failure on unregistered pass.
 # CHECK-LABEL: TEST: testParseFail
 def testParseFail():
   with Context():
     try:
       pm = PassManager.parse("unknown-pass")
     except ValueError as e:
-      # CHECK: ValueError exception: invalid pass pipeline 'unknown-pass'.
+      #      CHECK: ValueError exception: MLIR Textual PassPipeline Parser:1:1: error:
+      # CHECK-SAME: 'unknown-pass' does not refer to a registered pass or pass pipeline
+      #      CHECK: unknown-pass
+      #      CHECK: ^
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 run(testParseFail)
 
 
 # Verify failure on incorrect level of nesting.
 # CHECK-LABEL: TEST: testInvalidNesting
 def testInvalidNesting():
   with Context():
     try:
       pm = PassManager.parse("func.func(normalize-memrefs)")
     except ValueError as e:
-      # CHECK: Can't add pass 'NormalizeMemRefs' restricted to 'builtin.module' on a PassManager intended to run on 'func.func', did you intend to nest?
-      # CHECK: ValueError exception: invalid pass pipeline 'func.func(normalize-memrefs)'.
+      # CHECK: ValueError exception: Can't add pass 'NormalizeMemRefs' restricted to 'builtin.module' on a PassManager intended to run on 'func.func', did you intend to nest?
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 run(testInvalidNesting)
 
 
 # Verify that a pass manager can execute on IR
 # CHECK-LABEL: TEST: testRun