[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

3371–3372

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>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::StringRef;
 using llvm::Twine;
 
 //------------------------------------------------------------------------------
 // Docstrings (trivial, non-duplicated docstrings are included inline).
 //------------------------------------------------------------------------------
 
 static const char kContextParseTypeDocstring[] =
     R"(Parses the assembly form of a type.
 
-Returns a Type object or raises a ValueError if the type cannot be parsed.
+Returns a Type object or raises an MLIRError if the type cannot be parsed.
 
 See also: https://mlir.llvm.org/docs/LangRef/#type-system
 )";
 
 static const char kContextGetCallSiteLocationDocstring[] =
     R"(Gets a Location representing a caller and callsite)";
 
 static const char kContextGetFileLocationDocstring[] =
     R"(Gets a Location representing a file, line and column)";
 
 static const char kContextGetFusedLocationDocstring[] =
     R"(Gets a Location representing a fused location with optional metadata)";
 
 static const char kContextGetNameLocationDocString[] =
     R"(Gets a Location representing a named location with optional child location)";
 
 static const char kModuleParseDocstring[] =
     R"(Parses a module's assembly format from a string.
 
-Returns a new MlirModule or raises a ValueError if the parsing fails.
+Returns a new MlirModule or raises an MLIRError if the parsing fails.
 
 See also: https://mlir.llvm.org/docs/LangRef/
 )";
 
 static const char kOperationCreateDocstring[] =
     R"(Creates a new operation.
 
 Args:
auto deleteCallback = +[](void *userData) {
     pyHandlerObject.dec_ref();
   };
 
   pyHandler->registeredID = mlirContextAttachDiagnosticHandler(
       get(), handlerCallback, static_cast<void *>(pyHandler), deleteCallback);
   return pyHandlerObject;
 }
 
+MlirLogicalResult PyMlirContext::ErrorCapture::handler(MlirDiagnostic diag,
+                                                       void *userData) {
+  auto *self = static_cast<ErrorCapture *>(userData);
+  // Check if the context requested we emit errors instead of capturing them.
+  if (self->ctx->emitErrorDiagnostics)
+    return mlirLogicalResultFailure();
+
+  if (mlirDiagnosticGetSeverity(diag) != MlirDiagnosticError)
+    return mlirLogicalResultFailure();
+
+  PyDiagnostic pyDiag(diag);
+  std::vector<std::pair<PyLocation, py::str>> notes;
+  for (py::handle n : pyDiag.getNotes()) {
+    auto note = n.cast<PyDiagnostic>();
+    notes.emplace_back(note.getLocation(), note.getMessage());
+  }
+  self->errors.emplace_back(pyDiag.getLocation(), pyDiag.getMessage(),
+                            std::move(notes));
+
+  return mlirLogicalResultSuccess();
+}
+
 PyMlirContext &DefaultingPyMlirContext::resolve() {
   PyMlirContext *context = PyThreadContextEntry::getDefaultContext();
   if (!context) {
     throw SetPyError(
         PyExc_RuntimeError,
         "An MLIR function requires a Context but none was provided in the call "
         "or from the surrounding environment. Either pass to the function with "
         "a 'context=' argument or establish a default using 'with Context():'");
PyOperationRef created = createInstance(std::move(contextRef), operation,
                                           std::move(parentKeepAlive));
   created->attached = false;
   return created;
 }
 
 PyOperationRef PyOperation::parse(PyMlirContextRef contextRef,
                                   const std::string &sourceStr,
                                   const std::string &sourceName) {
+  PyMlirContext::ErrorCapture errors(contextRef);
   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", errors.take());
   return PyOperation::createDetached(std::move(contextRef), op);
 }
 
 void PyOperation::checkValid() const {
   if (!valid) {
     throw SetPyError(PyExc_RuntimeError, "the operation has been invalidated");
   }
 }
void PyOperationBase::moveBefore(PyOperationBase &other) {
   PyOperation &operation = getOperation();
   PyOperation &otherOp = other.getOperation();
   operation.checkValid();
   otherOp.checkValid();
   mlirOperationMoveBefore(operation, otherOp);
   operation.parentKeepAlive = otherOp.parentKeepAlive;
 }
 
+bool PyOperationBase::verify() {
+  PyOperation &op = getOperation();
+  PyMlirContext::ErrorCapture errors(op.getContext());
+  if (!mlirOperationVerify(op.get()))
+    throw MLIRError("Verification failed", errors.take());
+  return true;
+}
+
 std::optional<PyOperationRef> PyOperation::getParentOperation() {
   checkValid();
   if (!isAttached())
     throw SetPyError(PyExc_ValueError, "Detached operations have no parent");
   MlirOperation operation = mlirOperationGetParentOperation(get());
   if (mlirOperationIsNull(operation))
     return {};
   return PyOperation::forOperation(getContext(), operation);
py::class_<PyMlirContext>(m, "_BaseContext", py::module_local())
           },
           py::arg("operation_name"))
       .def(
           "append_dialect_registry",
           [](PyMlirContext &self, PyDialectRegistry &registry) {
             mlirContextAppendDialectRegistry(self.get(), registry);
           },
           py::arg("registry"))
+      .def_property("emit_error_diagnostics", nullptr,
+                    &PyMlirContext::setEmitErrorDiagnostics,
+                    "Emit error diagnostics to diagnostic handlers. By default "
+                    "error diagnostics are captured and reported through "
+                    "MLIRError exceptions.")
       .def("load_all_available_dialects", [](PyMlirContext &self) {
         mlirContextLoadAllAvailableDialects(self.get());
       });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialectDescriptor
   //----------------------------------------------------------------------------
   py::class_<PyDialectDescriptor>(m, "DialectDescriptor", py::module_local())
void mlir::python::populateIRCore(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of Module
   //----------------------------------------------------------------------------
   py::class_<PyModule>(m, "Module", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyModule::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyModule::createFromCapsule)
       .def_static(
           "parse",
-          [](const std::string moduleAsm, DefaultingPyMlirContext context) {
+          [](const std::string &moduleAsm, DefaultingPyMlirContext context) {
+            PyMlirContext::ErrorCapture errors(context->getRef());
             MlirModule module = mlirModuleCreateParse(
                 context->get(), toMlirStringRef(moduleAsm));
-            // TODO: Rework error reporting once diagnostic engine is exposed
-            // in C API.
-            if (mlirModuleIsNull(module)) {
-              throw SetPyError(
-                  PyExc_ValueError,
-                  "Unable to parse module assembly (see diagnostics)");
-            }
+            if (mlirModuleIsNull(module))
+              throw MLIRError("Unable to parse module assembly", errors.take());
             return PyModule::forModule(module).releaseObject();
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kModuleParseDocstring)
       .def_static(
           "create",
           [](DefaultingPyLocation loc) {
             MlirModule module = mlirModuleCreateEmpty(loc);
py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
            // Careful: Lots of arguments must match up with get_asm method.
            py::arg("binary") = false,
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
            py::arg("pretty_debug_info") = false,
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false,
            py::arg("assume_verified") = false, kOperationGetAsmDocstring)
-      .def(
-          "verify",
-          [](PyOperationBase &self) {
-            return mlirOperationVerify(self.getOperation());
-          },
-          "Verify the operation and return true if it passes, false if it "
-          "fails.")
+      .def("verify", &PyOperationBase::verify,
+           "Verify the operation. Raises MLIRError if verification fails, and "
+           "returns true otherwise.")
       .def("move_after", &PyOperationBase::moveAfter, py::arg("other"),
            "Puts self immediately after the other operation in its parent "
            "block.")
       .def("move_before", &PyOperationBase::moveBefore, py::arg("other"),
            "Puts self immediately before the other operation in its parent "
            "block.")
       .def(
           "detach_from_parent",
opViewClass.attr("parse") = classmethod(
 
         // Check if the expected operation was parsed, and cast to to the
         // appropriate `OpView` subclass if successful.
         // 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 =
+        MlirStringRef identifier =
             mlirIdentifierStr(mlirOperationGetName(*parsed.get()));
-        if (!mlirStringRefEqual(parsedOpName, toMlirStringRef(clsOpName)))
-          throw py::value_error(
-              "Expected a '" + clsOpName + "' op, got: '" +
-              std::string(parsedOpName.data, parsedOpName.length) + "'");
+        std::string_view parsedOpName(identifier.data, identifier.length);
+        if (clsOpName != parsedOpName)
+          throw MLIRError(Twine("Expected a '") + clsOpName + "' op, got: '" +
+                          parsedOpName + "'");
         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.
py::class_<PyAttribute>(m, "Attribute", py::module_local())
       .def(py::init<PyAttribute &>(), py::arg("cast_from_type"),
            "Casts the passed attribute to the generic Attribute")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAttribute::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAttribute::createFromCapsule)
       .def_static(
           "parse",
           [](std::string attrSpec, DefaultingPyMlirContext context) {
+            PyMlirContext::ErrorCapture errors(context->getRef());
             MlirAttribute type = mlirAttributeParseGet(
                 context->get(), toMlirStringRef(attrSpec));
-            // TODO: Rework error reporting once diagnostic engine is exposed
-            // in C API.
-            if (mlirAttributeIsNull(type)) {
-              throw SetPyError(PyExc_ValueError,
-                               Twine("Unable to parse attribute: '") +
-                                   attrSpec + "'");
-            }
+            if (mlirAttributeIsNull(type))
+              throw MLIRError("Unable to parse attribute", errors.take());
             return PyAttribute(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
-          "Parses an attribute from an assembly form")
+          "Parses an attribute from an assembly form. Raises an MLIRError on "
+          "failure.")
       .def_property_readonly(
           "context",
           [](PyAttribute &self) { return self.getContext().getObject(); },
           "Context that owns the Attribute")
       .def_property_readonly("type",
                              [](PyAttribute &self) {
                                return PyType(self.getContext()->getRef(),
                                              mlirAttributeGetType(self));
py::class_<PyType>(m, "Type", py::module_local())
       // extend the backing context which owns the MlirType.
       .def(py::init<PyType &>(), py::arg("cast_from_type"),
            "Casts the passed type to the generic Type")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyType::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyType::createFromCapsule)
       .def_static(
           "parse",
           [](std::string typeSpec, DefaultingPyMlirContext context) {
+            PyMlirContext::ErrorCapture errors(context->getRef());
             MlirType type =
                 mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec));
-            // TODO: Rework error reporting once diagnostic engine is exposed
-            // in C API.
-            if (mlirTypeIsNull(type)) {
-              throw SetPyError(PyExc_ValueError,
-                               Twine("Unable to parse type: '") + typeSpec +
-                                   "'");
-            }
+            if (mlirTypeIsNull(type))
+              throw MLIRError("Unable to parse type", errors.take());
             return PyType(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kContextParseTypeDocstring)
       .def_property_readonly(
           "context", [](PyType &self) { return self.getContext().getObject(); },
           "Context that owns the Type")
       .def("__eq__", [](PyType &self, PyType &other) { return self == other; })

   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) {
+    // We can't define exceptions with custom fields through pybind, so instead
+    // the exception class is defined in python and imported here.
+    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_Exception, obj.ptr());
+    }
+  });
 }

mlir/lib/Bindings/Python/IRModule.h

public:
   void contextExit(const pybind11::object &excType,
                    const pybind11::object &excVal,
                    const pybind11::object &excTb);
 
   /// Attaches a Python callback as a diagnostic handler, returning a
   /// registration object (internally a PyDiagnosticHandler).
   pybind11::object attachDiagnosticHandler(pybind11::object callback);
 
+  /// Controls whether error diagnostics should be propagated to diagnostic
+  /// handlers, instead of being captured by `ErrorCapture`.
+  void setEmitErrorDiagnostics(bool value) { emitErrorDiagnostics = value; }
+
+  /// RAII object that captures any error diagnostics emitted to the provided
+  /// context.
+  struct ErrorCapture {
+    ErrorCapture(PyMlirContextRef ctx)
+        : ctx(ctx), handlerID(mlirContextAttachDiagnosticHandler(
+                        ctx->get(), handler, /*userData=*/this,
+                        /*deleteUserData=*/nullptr)) {}
+    ~ErrorCapture() {
+      mlirContextDetachDiagnosticHandler(ctx->get(), handlerID);
+      assert(errors.empty() && "unhandled captured errors");
+    }
+
+    // Tuple of {location, message, notes} for an error diagnostic.
+    using ErrorInfo =
+        std::tuple<PyLocation, pybind11::str,
+                   std::vector<std::pair<PyLocation, pybind11::str>>>;
+    std::vector<ErrorInfo> take() { return std::move(errors); };
+
+  private:
+    PyMlirContextRef ctx;
+    MlirDiagnosticHandlerID handlerID;
+    std::vector<ErrorInfo> errors;
+
+    static MlirLogicalResult handler(MlirDiagnostic diag, void *userData);
+  };
+
 private:
   PyMlirContext(MlirContext context);
   // Interns the mapping of live MlirContext::ptr to PyMlirContext instances,
   // preserving the relationship that an MlirContext maps to a single
   // PyMlirContext wrapper. This could be replaced in the future with an
   // extension mechanism on the MlirContext for stashing user pointers.
   // Note that this holds a handle, which does not imply ownership.
   // Mappings will be removed when the context is destructed.
private:
   // Interns all live operations associated with this context. Operations
   // tracked in this map are valid. When an operation is invalidated, it is
   // removed from this map, and while it still exists as an instance, any
   // attempt to access it will raise an error.
   using LiveOperationMap =
       llvm::DenseMap<void *, std::pair<pybind11::handle, PyOperation *>>;
   LiveOperationMap liveOperations;
 
+  bool emitErrorDiagnostics = false;
+
   MlirContext context;
   friend class PyModule;
   friend class PyOperation;
 };
 
 /// Used in function arguments when None should resolve to the current context
 /// manager set instance.
 class DefaultingPyMlirContext
public:
 
   // Implement the bound 'writeBytecode' method.
   void writeBytecode(const pybind11::object &fileObject);
 
   /// Moves the operation before or after the other operation.
   void moveAfter(PyOperationBase &other);
   void moveBefore(PyOperationBase &other);
 
+  /// Verify the operation. Throws `MLIRError` if verification fails, and
+  /// returns `true` otherwise.
+  bool verify();
+
   /// Each must provide access to the raw Operation.
   virtual PyOperation &getOperation() = 0;
 };
 
 /// Wrapper around PyOperation.
 /// Operations exist in either an attached (dependent) or detached (top-level)
 /// state. In the detached state (as on creation), an operation is owned by
 /// the creator and its lifetime extends either until its reference count
public:
   /// Casts the bindings class into the C API structure.
   operator MlirSymbolTable() { return symbolTable; }
 
 private:
   PyOperationRef operation;
   MlirSymbolTable symbolTable;
 };
 
+/// Custom exception that allows access to error diagnostic information. This is
+/// converted to the `ir.MLIRError` python exception when thrown.
+struct MLIRError {
+  MLIRError(llvm::Twine message,
+            std::vector<PyMlirContext::ErrorCapture::ErrorInfo>
+                &&errorDiagnostics = {})
+      : message(message.str()), errorDiagnostics(std::move(errorDiagnostics)) {}
+  std::string message;
+  std::vector<PyMlirContext::ErrorCapture::ErrorInfo> errorDiagnostics;
+};
+
 void populateIRAffine(pybind11::module &m);
 void populateIRAttributes(pybind11::module &m);
 void populateIRCore(pybind11::module &m);
 void populateIRInterfaces(pybind11::module &m);
 void populateIRTypes(pybind11::module &m);
 
 } // namespace python
 } // namespace mlir

mlir/lib/Bindings/Python/Pass.cpp

py::class_<PyPassManager>(m, "PassManager", py::module_local())
               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, PyOperationBase &op) {
+            PyMlirContext::ErrorCapture errors(op.getOperation().getContext());
             MlirLogicalResult status =
                 mlirPassManagerRun(passManager.get(), op.getOperation().get());
             if (mlirLogicalResultIsFailure(status))
-              throw SetPyError(PyExc_RuntimeError,
-                               "Failure while executing pass pipeline.");
+              throw MLIRError("Failure while executing pass pipeline",
+                              errors.take());
           },
           py::arg("operation"),
-          "Run the pass manager on the provided operation, throw a "
-          "RuntimeError on failure.")
+          "Run the pass manager on the provided operation, raising an "
+          "MLIRError 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/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(Exception):
+    """
+    An exception with diagnostic information. Has the following fields:
+      message: str
+      error_diagnostics: List[Tuple[ir.Location, str, List[Tuple[ir.Location, str]]]]
+    """
+    def __init__(self, message, error_diagnostics):
+      self.message = message
+      self.error_diagnostics = error_diagnostics
+      super().__init__(message, error_diagnostics)
+
+    def __str__(self):
+      s = self.message
+      if self.error_diagnostics:
+        s += ':'
+      for loc, msg, notes in self.error_diagnostics:
+        s += "\nerror: "  + str(loc)[4:-1] + ": " + msg.replace('\n', '\n  ')
+        for loc, msg, in notes:
+          s += "\n note: "  + str(loc)[4:-1] + ": " + msg.replace('\n', '\n  ')
+      return s
+  ir.MLIRError = MLIRError
+
 
 _site_initialize()

mlir/test/python/ir/attributes.py

def testParsePrint():
   gc.collect()
   # CHECK: "hello"
   print(str(t))
   # CHECK: Attribute("hello")
   print(repr(t))
 
 
 # CHECK-LABEL: TEST: testParseError
-# TODO: Hook the diagnostic manager to capture a more meaningful error
-# message.
 @run
 def testParseError():
   with Context():
     try:
       t = Attribute.parse("BAD_ATTR_DOES_NOT_EXIST")
-    except ValueError as e:
-      # CHECK: Unable to parse attribute: 'BAD_ATTR_DOES_NOT_EXIST'
-      print("testParseError:", e)
+    except MLIRError as e:
+      # CHECK: testParseError: <
+      # CHECK:   Unable to parse attribute:
+      # CHECK:   error: "BAD_ATTR_DOES_NOT_EXIST":1:1: expected attribute value
+      # CHECK: >
+      print(f"testParseError: <{e}>")
     else:
       print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testAttrEq
 @run
 def testAttrEq():
   with Context():

mlir/test/python/ir/builtin_types.py

def testParsePrint():
   gc.collect()
   # CHECK: i32
   print(str(t))
   # CHECK: Type(i32)
   print(repr(t))
 
 
 # CHECK-LABEL: TEST: testParseError
-# TODO: Hook the diagnostic manager to capture a more meaningful error
-# message.
 @run
 def testParseError():
   ctx = Context()
   try:
     t = Type.parse("BAD_TYPE_DOES_NOT_EXIST", ctx)
-  except ValueError as e:
-    # CHECK: Unable to parse type: 'BAD_TYPE_DOES_NOT_EXIST'
-    print("testParseError:", e)
+  except MLIRError as e:
+    # CHECK: testParseError: <
+    # CHECK:   Unable to parse type:
+    # CHECK:   error: "BAD_TYPE_DOES_NOT_EXIST":1:1: expected non-function type
+    # CHECK: >
+    print(f"testParseError: <{e}>")
   else:
     print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testTypeEq
 @run
 def testTypeEq():
   ctx = Context()

mlir/test/python/ir/diagnostic_handler.py

def testDiagnosticEmptyNotes():
   loc.emit_error("foobar")
   assert not handler.had_error
 
 
 # CHECK-LABEL: TEST: testDiagnosticNonEmptyNotes
 @run
 def testDiagnosticNonEmptyNotes():
   ctx = Context()
+  ctx.emit_error_diagnostics = True
   def callback(d):
     # CHECK: DIAGNOSTIC:
     # CHECK:   message='arith.addi' op requires one result
     # CHECK:   notes=['see current operation: "arith.addi"() : () -> ()']
     print(f"DIAGNOSTIC:")
     print(f"  message={d.message}")
     print(f"  notes={list(map(str, d.notes))}")
     return True
   handler = ctx.attach_diagnostic_handler(callback)
   loc = Location.unknown(ctx)
+  try:
     Operation.create('arith.addi', loc=loc).verify()
+  except MLIRError:
+    pass
   assert not handler.had_error
 
 # CHECK-LABEL: TEST: testDiagnosticCallbackException
 @run
 def testDiagnosticCallbackException():
   ctx = Context()
   def callback(d):
     raise ValueError("Error in handler")

mlir/test/python/ir/exception.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import gc
+from mlir.ir import *
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+  return f
+
+
+# CHECK-LABEL: TEST: test_exception
+@run
+def test_exception():
+  ctx =  Context()
+  ctx.allow_unregistered_dialects = True
+  try:
+    Operation.parse("""
+      func.func @foo() {
+          "test.use"(%0) : (i64) -> ()  loc("use")
+          %0 = "test.def"() : () -> i64 loc("def")
+          return
+      }
+    """, context=ctx)
+  except MLIRError as e:
+    # CHECK: Exception: <
+    # CHECK:   Unable to parse operation assembly:
+    # CHECK:   error: "use": operand #0 does not dominate this use
+    # CHECK:    note: "use": see current operation: "test.use"(%0) : (i64) -> ()
+    # CHECK:    note: "def": operand defined here (op in the same block)
+    # CHECK: >
+    print(f"Exception: <{e}>")
+
+    # CHECK: message: Unable to parse operation assembly
+    print(f"message: {e.message}")
+
+    #      CHECK: error_diagnostics: [(
+    # CHECK-SAME:   loc("use"), 'operand #0 does not dominate this use', [
+    # CHECK-SAME:     (loc("use"), 'see current operation: "test.use"(%0) : (i64) -> ()'),
+    # CHECK-SAME:     (loc("def"), 'operand defined here (op in the same block)')
+    # CHECK-SAME:   ]
+    # CHECK-SAME: )]
+    print(f"error_diagnostics: {e.error_diagnostics}")
+
+
+# CHECK-LABEL: test_emit_error_diagnostics
+@run
+def test_emit_error_diagnostics():
+  ctx = Context()
+  loc = Location.unknown(ctx)
+  handler_diags = []
+  def handler(d):
+    handler_diags.append(d.message)
+    return True
+  ctx.attach_diagnostic_handler(handler)
+
+  try:
+    Attribute.parse("not an attr", ctx)
+  except MLIRError as e:
+    # CHECK: emit_error_diagnostics=False:
+    # CHECK: e.error_diagnostics: [(loc("not an attr":1:1), 'expected attribute value', [])]
+    # CHECK: handler_diags: []
+    print(f"emit_error_diagnostics=False:")
+    print(f"e.error_diagnostics: {e.error_diagnostics}")
+    print(f"handler_diags: {handler_diags}")
+
+  ctx.emit_error_diagnostics = True
+  try:
+    Attribute.parse("not an attr", ctx)
+  except MLIRError as e:
+    # CHECK: emit_error_diagnostics=True:
+    # CHECK: e.error_diagnostics: []
+    # CHECK: handler_diags: ['expected attribute value']
+    print(f"emit_error_diagnostics=True:")
+    print(f"e.error_diagnostics: {e.error_diagnostics}")
+    print(f"handler_diags: {handler_diags}")

mlir/test/python/ir/module.py

def testParseSuccess():
   ctx = None  # Ensure that module captures the context.
   gc.collect()
   module.dump()  # Just outputs to stderr. Verifies that it functions.
   print(str(module))
 
 
 # Verify parse error.
 # CHECK-LABEL: TEST: testParseError
-# CHECK: testParseError: Unable to parse module assembly (see diagnostics)
+# CHECK: testParseError: <
+# CHECK:   Unable to parse module assembly:
+# CHECK:   error: "-":1:1: expected operation name in quotes
+# CHECK: >
 @run
 def testParseError():
   ctx = Context()
   try:
     module = Module.parse(r"""}SYNTAX ERROR{""", ctx)
-  except ValueError as e:
-    print("testParseError:", e)
+  except MLIRError as e:
+    print(f"testParseError: <{e}>")
   else:
     print("Exception not produced")
 
 
 # Verify successful parse.
 # CHECK-LABEL: TEST: testCreateEmpty
 # CHECK: module {
 @run

mlir/test/python/ir/operation.py

 def testInvalidOperationStrSoftFails():
   ctx = Context()
   with Location.unknown(ctx):
     invalid_op = create_invalid_operation()
     # Verify that we fallback to the generic printer for safety.
     # CHECK: "builtin.module"() ({
     # CHECK: }) : () -> ()
     print(invalid_op)
-    # CHECK: .verify = False
-    print(f".verify = {invalid_op.operation.verify()}")
+    try:
+      invalid_op.verify()
+    except MLIRError as e:
+      # CHECK: Exception: <
+      # CHECK:   Verification failed:
+      # CHECK:   error: unknown: 'builtin.module' op requires one region
+      # CHECK:    note: unknown: see current operation:
+      # CHECK:     "builtin.module"() ({
+      # CHECK:     ^bb0:
+      # CHECK:     }, {
+      # CHECK:     }) : () -> ()
+      # CHECK: >
+      print(f"Exception: <{e}>")
 
 
 # CHECK-LABEL: TEST: testInvalidModuleStrSoftFails
 @run
 def testInvalidModuleStrSoftFails():
   ctx = Context()
   with Location.unknown(ctx):
     module = Module.create()
with Context() as ctx:
     assert isinstance(m, ModuleOp)
     assert type(o) is OpView
 
     # Parsing specific operation.
     m = ModuleOp.parse('module {}')
     assert isinstance(m, ModuleOp)
     try:
       ModuleOp.parse('"test.foo"() : () -> ()')
-    except ValueError as e:
+    except MLIRError 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)}")

mlir/test/python/pass_manager.py

except ValueError as e:
       # 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
+# CHECK-LABEL: TEST: testRunPipeline
 def testRunPipeline():
   with Context():
     pm = PassManager.parse("any(print-op-stats{json=false})")
     func = FuncOp.parse(r"""func.func @successfulParse() { return }""")
     pm.run(func)
 # CHECK: Operations encountered:
 # CHECK: func.func      , 1
 # CHECK: func.return        , 1
 run(testRunPipeline)
+
+# CHECK-LABEL: TEST: testRunPipelineError
+@run
+def testRunPipelineError():
+  with Context() as ctx:
+    ctx.allow_unregistered_dialects = True
+    op = Operation.parse('"test.op"() : () -> ()')
+    pm = PassManager.parse("any(cse)")
+    try:
+      pm.run(op)
+    except MLIRError as e:
+      # CHECK: Exception: <
+      # CHECK:   Failure while executing pass pipeline:
+      # CHECK:   error: "-":1:1: 'test.op' op trying to schedule a pass on an unregistered operation
+      # CHECK:    note: "-":1:1: see current operation: "test.op"() : () -> ()
+      # CHECK: >
+      print(f"Exception: <{e}>")