[Mlir-commits] [mlir] b15ccd4 - [mlir] Better Python diagnostics (#128581)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Mon Mar 10 15:59:50 PDT 2025


Author: Nikhil Kalra
Date: 2025-03-10T15:59:47-07:00
New Revision: b15ccd436aef320802635984f336fcfe5b69ff0b

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

LOG: [mlir] Better Python diagnostics (#128581)

Updated the Python diagnostics handler to emit notes (in addition to
errors) into the output stream so that users have more context as to
where in the IR the error is occurring.

Added: 
    

Modified: 
    mlir/include/mlir/Bindings/Python/Diagnostics.h
    mlir/test/python/ir/diagnostic_handler.py
    mlir/test/python/lib/PythonTestCAPI.cpp
    mlir/test/python/lib/PythonTestCAPI.h
    mlir/test/python/lib/PythonTestModuleNanobind.cpp

Removed: 
    


################################################################################
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 <cassert>
-#include <string>
-
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
-#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cassert>
+#include <cstdint>
+#include <string>
 
 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<std::string *>(data) +=
-          llvm::StringRef(message.data, message.length);
+      *static_cast<llvm::raw_string_ostream *>(data)
+          << std::string_view(message.data, message.length);
     };
     MlirLocation loc = mlirDiagnosticGetLocation(diag);
-    *static_cast<std::string *>(data) += "at ";
+    *static_cast<llvm::raw_string_ostream *>(data) << "at ";
     mlirLocationPrint(loc, printer, data);
-    *static_cast<std::string *>(data) += ": ";
+    *static_cast<llvm::raw_string_ostream *>(data) << ": ";
     mlirDiagnosticPrint(diag, printer, data);
+    for (intptr_t i = 0; i < mlirDiagnosticGetNumNotes(diag); i++) {
+      *static_cast<llvm::raw_string_ostream *>(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
+ at 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)


        


More information about the Mlir-commits mailing list