diff --git a/mlir/include/mlir/Bindings/Python/Diagnostics.h b/mlir/include/mlir/Bindings/Python/Diagnostics.h index ea80e14dde0f3..167002d561931 100644 --- a/mlir/include/mlir/Bindings/Python/Diagnostics.h +++ b/mlir/include/mlir/Bindings/Python/Diagnostics.h @@ -9,12 +9,13 @@ #ifndef MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H #define MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H -#include -#include - #include "mlir-c/Diagnostics.h" #include "mlir-c/IR.h" -#include "llvm/ADT/StringRef.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include namespace mlir { namespace python { @@ -24,33 +25,45 @@ namespace python { class CollectDiagnosticsToStringScope { public: explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) { - handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage, - /*deleteUserData=*/nullptr); + handlerID = + mlirContextAttachDiagnosticHandler(ctx, &handler, &messageStream, + /*deleteUserData=*/nullptr); } ~CollectDiagnosticsToStringScope() { - assert(errorMessage.empty() && "unchecked error message"); + assert(message.empty() && "unchecked error message"); mlirContextDetachDiagnosticHandler(context, handlerID); } - [[nodiscard]] std::string takeMessage() { return std::move(errorMessage); } + [[nodiscard]] std::string takeMessage() { + std::string newMessage; + std::swap(message, newMessage); + return newMessage; + } private: static MlirLogicalResult handler(MlirDiagnostic diag, void *data) { auto printer = +[](MlirStringRef message, void *data) { - *static_cast(data) += - llvm::StringRef(message.data, message.length); + *static_cast(data) + << std::string_view(message.data, message.length); }; MlirLocation loc = mlirDiagnosticGetLocation(diag); - *static_cast(data) += "at "; + *static_cast(data) << "at "; mlirLocationPrint(loc, printer, data); - *static_cast(data) += ": "; + *static_cast(data) << ": "; mlirDiagnosticPrint(diag, printer, data); + for (intptr_t i = 0; i < mlirDiagnosticGetNumNotes(diag); i++) { + *static_cast(data) << "\n"; + MlirDiagnostic note = mlirDiagnosticGetNote(diag, i); + handler(note, data); + } return mlirLogicalResultSuccess(); } MlirContext context; MlirDiagnosticHandlerID handlerID; - std::string errorMessage = ""; + + std::string message; + llvm::raw_string_ostream messageStream{message}; }; } // namespace python diff --git a/mlir/test/python/ir/diagnostic_handler.py b/mlir/test/python/ir/diagnostic_handler.py index d516cda819897..6d273e5092e42 100644 --- a/mlir/test/python/ir/diagnostic_handler.py +++ b/mlir/test/python/ir/diagnostic_handler.py @@ -2,6 +2,9 @@ import gc from mlir.ir import * +from mlir._mlir_libs._mlirPythonTestNanobind import ( + test_diagnostics_with_errors_and_notes, +) def run(f): @@ -222,3 +225,16 @@ def callback2(d): # CHECK: CALLBACK2: foobar # CHECK: CALLBACK1: foobar loc.emit_error("foobar") + + +# CHECK-LABEL: TEST: testBuiltInDiagnosticsHandler +@run +def testBuiltInDiagnosticsHandler(): + ctx = Context() + + try: + test_diagnostics_with_errors_and_notes(ctx) + except ValueError as e: + # CHECK: created error + # CHECK: attached note + print(e) diff --git a/mlir/test/python/lib/PythonTestCAPI.cpp b/mlir/test/python/lib/PythonTestCAPI.cpp index cb7d7677714fe..ee6a5ae2d9c3d 100644 --- a/mlir/test/python/lib/PythonTestCAPI.cpp +++ b/mlir/test/python/lib/PythonTestCAPI.cpp @@ -11,6 +11,8 @@ #include "mlir-c/BuiltinTypes.h" #include "mlir/CAPI/Registration.h" #include "mlir/CAPI/Wrap.h" +#include "mlir/IR/Diagnostics.h" +#include "mlir/IR/Location.h" MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(PythonTest, python_test, python_test::PythonTestDialect) @@ -42,3 +44,9 @@ MlirTypeID mlirPythonTestTestTypeGetTypeID(void) { bool mlirTypeIsAPythonTestTestTensorValue(MlirValue value) { return mlirTypeIsATensor(wrap(unwrap(value).getType())); } + +void mlirPythonTestEmitDiagnosticWithNote(MlirContext ctx) { + auto diag = + mlir::emitError(unwrap(mlirLocationUnknownGet(ctx)), "created error"); + diag.attachNote() << "attached note"; +} diff --git a/mlir/test/python/lib/PythonTestCAPI.h b/mlir/test/python/lib/PythonTestCAPI.h index 43f8fdcbfae12..5ffd0cbf6201e 100644 --- a/mlir/test/python/lib/PythonTestCAPI.h +++ b/mlir/test/python/lib/PythonTestCAPI.h @@ -10,6 +10,7 @@ #define MLIR_TEST_PYTHON_LIB_PYTHONTESTCAPI_H #include "mlir-c/IR.h" +#include "mlir-c/Support.h" #ifdef __cplusplus extern "C" { @@ -33,6 +34,8 @@ MLIR_CAPI_EXPORTED MlirTypeID mlirPythonTestTestTypeGetTypeID(void); MLIR_CAPI_EXPORTED bool mlirTypeIsAPythonTestTestTensorValue(MlirValue value); +MLIR_CAPI_EXPORTED void mlirPythonTestEmitDiagnosticWithNote(MlirContext ctx); + #ifdef __cplusplus } #endif diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp index 99c81eae97a0c..dd0a9f2b66ea6 100644 --- a/mlir/test/python/lib/PythonTestModuleNanobind.cpp +++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp @@ -11,9 +11,12 @@ #include "PythonTestCAPI.h" #include "mlir-c/BuiltinAttributes.h" #include "mlir-c/BuiltinTypes.h" +#include "mlir-c/Diagnostics.h" #include "mlir-c/IR.h" +#include "mlir/Bindings/Python/Diagnostics.h" #include "mlir/Bindings/Python/Nanobind.h" #include "mlir/Bindings/Python/NanobindAdaptors.h" +#include "nanobind/nanobind.h" namespace nb = nanobind; using namespace mlir::python::nanobind_adaptors; @@ -45,6 +48,13 @@ NB_MODULE(_mlirPythonTestNanobind, m) { }, nb::arg("registry")); + m.def("test_diagnostics_with_errors_and_notes", [](MlirContext ctx) { + mlir::python::CollectDiagnosticsToStringScope handler(ctx); + + mlirPythonTestEmitDiagnosticWithNote(ctx); + throw nb::value_error(handler.takeMessage().c_str()); + }); + mlir_attribute_subclass(m, "TestAttr", mlirAttributeIsAPythonTestTestAttribute, mlirPythonTestTestAttributeGetTypeID)