[mlir][python] Capture error diagnostics in exceptions

[rkayaith]

Summary

This updates most (all?) error-diagnostic-emitting python APIs to
capture error diagnostics and include them in the raised exception's
message:

>>> Operation.parse('"arith.addi"() : () -> ()'))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
mlir._mlir_libs.MLIRError: Unable to parse operation assembly:
error: "-":1:1: 'arith.addi' op requires one result
 note: "-":1:1: see current operation: "arith.addi"() : () -> ()

The diagnostic information is available on the exception for users who
may want to customize the error message:

>>> try:
...   Operation.parse('"arith.addi"() : () -> ()')
... except MLIRError as e:
...   print(e.message)
...   print(e.error_diagnostics)
...   print(e.error_diagnostics[0].message)
... 
Unable to parse operation assembly
[<mlir._mlir_libs._mlir.ir.DiagnosticInfo object at 0x7fed32bd6b70>]
'arith.addi' op requires one result

Error diagnostics captured in exceptions aren't propagated to diagnostic
handlers, to avoid double-reporting of errors. The context-level
emit_error_diagnostics option can be used to revert to the old
behaviour, causing error diagnostics to be reported to handlers instead
of as part of exceptions.

API changes:

• Operation.verify now raises an exception on verification failure, instead of returning false
• The exception raised by the following methods has been changed to MLIRError:
  • PassManager.run
  • {Module,Operation,Type,Attribute}.parse
  • {RankedTensorType,UnrankedTensorType}.get
  • {MemRefType,UnrankedMemRefType}.get
  • VectorType.get
  • FloatAttr.get

closes #60595

depends on D144804, D143830

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

3397–3398

I couldn't find a way to define an exception with custom __init__ through pybind, so it's defined in python + imported here. Another option would be to use the C API.

[rkayaith]

update

[rkayaith]

missed some methods

[rkayaith]

Since I only changed error-emitting APIs to raise MLIRError, the exception type used is somewhat inconsistent, e.g. ComplexType.get still raises ValueError. Not sure what to do about those (raise MLIRError everywhere even if there's no diagnostics?)

[stellaraccident]

Nice, thanks! Fine to land as a stretch improvement but I think we need to further sugar the new exception class to avoid exposing implementation details.

mlir/python/mlir/_mlir_libs/__init__.py

109

I think we may want to further typify this vs having deeply nested tuples. Can be fine in a follow-up to sugar. Would make string printing easier and would leak fewer details to callers who want to extract the diagnostics.

[rkayaith]

Replace diagnostic tuples with a DiagnosticInfo class, which just mirrors the
main Diagnostic class but is safe to access outside handlers.

mlir/python/mlir/_mlir_libs/__init__.py

109

Took a stab at this, let me know if this is what you had in mind.

[rkayaith]

rebase

mlir/lib/Bindings/Python/IRCore.cpp

 
 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Debug.h"
+#include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 //#include "mlir-c/Registration.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <optional>
 #include <utility>
 
createCustomDialectWrapper(const std::string &dialectNamespace,
   // Create the custom implementation.
   return (*dialectClass)(std::move(dialectDescriptor));
 }
 
 static MlirStringRef toMlirStringRef(const std::string &s) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
+struct MLIRError : public std::exception {
+  MLIRError(Twine message, std::vector<py::tuple> &&errorDiagnostics)
+      : message(message.str()), errorDiagnostics(errorDiagnostics) {}
+  std::string message;
+  std::vector<py::tuple> errorDiagnostics;
+};
+
 /// Wrapper for the global LLVM debugging flag.
 struct PyGlobalDebugFlag {
   static void set(py::object &o, bool enable) { mlirEnableGlobalDebug(enable); }
 
   static bool get(const py::object &) { return mlirIsGlobalDebugEnabled(); }
 
   static void bind(py::module &m) {
     // Debug flags.
py::tuple PyDiagnostic::getNotes() {
   materializedNotes = py::tuple(numNotes);
   for (intptr_t i = 0; i < numNotes; ++i) {
     MlirDiagnostic noteDiag = mlirDiagnosticGetNote(diagnostic, i);
     (*materializedNotes)[i] = PyDiagnostic(noteDiag);
   }
   return *materializedNotes;
 }
 
+py::tuple PyDiagnostic::asTuple() {
+  py::list notes;
+  for (py::handle x : getNotes())
+    notes.append(x.cast<PyDiagnostic>().asTuple());
+  return py::make_tuple(getSeverity(), getLocation(), getMessage(), notes);
+}
+
 //------------------------------------------------------------------------------
 // PyDialect, PyDialectDescriptor, PyDialects, PyDialectRegistry
 //------------------------------------------------------------------------------
 
 MlirDialect PyDialects::getDialectForKey(const std::string &key,
                                          bool attrError) {
   MlirDialect dialect = mlirContextGetOrLoadDialect(getContext()->get(),
                                                     {key.data(), key.size()});
PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef,
   (void)liveOperations;
 
   PyOperationRef created = createInstance(std::move(contextRef), operation,
                                           std::move(parentKeepAlive));
   created->attached = false;
   return created;
 }
 
+struct PyScopedErrorDiagCapture {
+  MlirContext ctx;
+  MlirDiagnosticHandlerID handlerID;
+  std::vector<py::tuple> errorDiagnostics;
+
+  PyScopedErrorDiagCapture(MlirContext ctx)
+      : ctx(ctx), handlerID(mlirContextAttachDiagnosticHandler(
+                      ctx, handler, this, /*deleteUserData=*/nullptr)) {}
+  ~PyScopedErrorDiagCapture() {
+    mlirContextDetachDiagnosticHandler(ctx, handlerID);
+  }
+
+  static MlirLogicalResult handler(MlirDiagnostic diag, void *userData) {
+    if (mlirDiagnosticGetSeverity(diag) != MlirDiagnosticError)
+      return mlirLogicalResultFailure();
+    auto *diagCapture = static_cast<PyScopedErrorDiagCapture *>(userData);
+    diagCapture->errorDiagnostics.push_back(PyDiagnostic(diag).asTuple());
+    return mlirLogicalResultSuccess();
+  }
+};
+
 PyOperationRef PyOperation::parse(PyMlirContextRef contextRef,
                                   const std::string &sourceStr,
                                   const std::string &sourceName) {
+  PyScopedErrorDiagCapture diagCapture(contextRef->get());
   MlirOperation op =
       mlirOperationCreateParse(contextRef->get(), toMlirStringRef(sourceStr),
                                toMlirStringRef(sourceName));
-  // TODO: Include error diagnostic messages in the exception message
   if (mlirOperationIsNull(op))
-    throw py::value_error(
-        "Unable to parse operation assembly (see diagnostics)");
+    throw MLIRError("Unable to parse operation assembly",
+                    std::move(diagCapture.errorDiagnostics));
   return PyOperation::createDetached(std::move(contextRef), op);
 }
 
 void PyOperation::checkValid() const {
   if (!valid) {
     throw SetPyError(PyExc_RuntimeError, "the operation has been invalidated");
   }
 }
opViewClass.attr("parse") = classmethod(
         // NOTE: This accesses attributes that have been automatically added to
         // `OpView` subclasses, and is not intended to be used on `OpView`
         // directly.
         std::string clsOpName =
             py::cast<std::string>(cls.attr("OPERATION_NAME"));
         MlirStringRef parsedOpName =
             mlirIdentifierStr(mlirOperationGetName(*parsed.get()));
         if (!mlirStringRefEqual(parsedOpName, toMlirStringRef(clsOpName)))
-          throw py::value_error(
+          throw MLIRError(
               "Expected a '" + clsOpName + "' op, got: '" +
-              std::string(parsedOpName.data, parsedOpName.length) + "'");
+                  std::string(parsedOpName.data, parsedOpName.length) + "'",
+              {});
         return PyOpView::constructDerived(cls, *parsed.get());
       },
       py::arg("cls"), py::arg("source"), py::kw_only(),
       py::arg("source_name") = "", py::arg("context") = py::none(),
       "Parses a specific, generated OpView based on class level attributes");
 
   //----------------------------------------------------------------------------
   // Mapping of PyRegion.

   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
   PyGlobalDebugFlag::bind(m);
 
   // Attribute builder getter.
   PyAttrBuilderMap::bind(m);
+
+  py::register_local_exception_translator([](std::exception_ptr p) {
+    try {
+      if (p)
+        std::rethrow_exception(p);
+    } catch (const MLIRError &e) {
+      py::object obj = py::module_::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+                           .attr("MLIRError")(e.message, e.errorDiagnostics);
+      PyErr_SetObject(PyExc_ValueError, obj.ptr());
+    }
+  });
 }

mlir/lib/Bindings/Python/IRModule.h

 public:
   PyDiagnostic(MlirDiagnostic diagnostic) : diagnostic(diagnostic) {}
   void invalidate();
   bool isValid() { return valid; }
   MlirDiagnosticSeverity getSeverity();
   PyLocation getLocation();
   pybind11::str getMessage();
   pybind11::tuple getNotes();
+  pybind11::tuple asTuple();
 
 private:
   MlirDiagnostic diagnostic;
 
   void checkValid();
   /// If notes have been materialized from the diagnostic, then this will
   /// be populated with the corresponding objects (all castable to
   /// PyDiagnostic).

mlir/python/mlir/_mlir_libs/__init__.py

def __init__(self, *args, **kwargs):
       super().__init__(*args, **kwargs)
       self.append_dialect_registry(registry)
       for hook in post_init_hooks:
         hook(self)
       # TODO: There is some debate about whether we should eagerly load
       # all dialects. It is being done here in order to preserve existing
       # behavior. See: https://github.com/llvm/llvm-project/issues/56037
       self.load_all_available_dialects()
-
   ir.Context = Context
 
+  class MLIRError(ValueError):
+    def __init__(self, message, error_diagnostics):
+      self.message = message
+      self.error_diagnostics = error_diagnostics
+      super().__init__(str(self))
+
+    def __str__(self):
+      s = self.message
+      if self.error_diagnostics:
+        s += ':'
+      for severity, loc, msg, notes in self.error_diagnostics:
+        assert severity == ir.DiagnosticSeverity.ERROR
+        s += f"\nerror: {loc}: {msg}"
+        for severity, loc, msg, notes in notes:
+          assert severity == ir.DiagnosticSeverity.NOTE
+          assert not notes
+          s += f"\n       {loc}: {msg}"
+      return s
+  ir.MLIRError = MLIRError
+
 
 _site_initialize()

mlir/test/python/ir/operation.py

except ValueError as e:
       # CHECK: error: Expected a 'builtin.module' op, got: 'test.foo'
       print(f"error: {e}")
     else:
       assert False, "expected error"
 
     o = Operation.parse('"test.foo"() : () -> ()', source_name="my-source-string")
     # CHECK: op_with_source_name: "test.foo"() : () -> () loc("my-source-string":1:1)
     print(f"op_with_source_name: {o.get_asm(enable_debug_info=True, use_local_scope=True)}")
+
+
+# CHECK-LABEL: TEST: test_exception
+@run
+def test_exception():
+  from pprint import pprint
+  with Context() as ctx, Location.unknown():
+    ctx.allow_unregistered_dialects = True
+    try:
+      Operation.parse("""
+        func.func @foo() {
+          "test.use"(%0) : (i64) -> ()  loc("use")
+          %0 = "test.def"() : () -> i64 loc("def")
+          return
+        }
+      """)
+    except MLIRError as e:
+      # CHECK: exception:
+      # CHECK: Unable to parse operation assembly:
+      # CHECK: error: loc("use"): operand #0 does not dominate this use
+      # CHECK:        loc("use"): see current operation: "test.use"(%0) : (i64) -> ()
+      # CHECK:        loc("def"): operand defined here (op in the same block)
+      print(f"exception:\n{e}")
+
+      # CHECK: message: Unable to parse operation assembly
+      print(f"message: {e.message}")
+
+      # CHECK: error_diagnostics:
+      # CHECK: [(<DiagnosticSeverity.ERROR: 0>,
+      # CHECK:   loc("use"),
+      # CHECK:   'operand #0 does not dominate this use',
+      # CHECK:   [(<DiagnosticSeverity.NOTE: 2>,
+      # CHECK:     loc("use"),
+      # CHECK:     'see current operation: "test.use"(%0) : (i64) -> ()',
+      # CHECK:     []),
+      # CHECK:    (<DiagnosticSeverity.NOTE: 2>,
+      # CHECK:     loc("def"),
+      # CHECK:     'operand defined here (op in the same block)',
+      # CHECK:     [])])]
+      print("error_diagnostics: ")
+      pprint(e.error_diagnostics)