[Mlir-commits] [mlir] a3a3784 - [mlir] Own the registrations of SourceMgrDiagnosticVerifierHandler (#226258)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Tue Sep 29 01:36:16 PDT 2026


Author: Henrich Lauko
Date: 2026-09-29T10:36:10+02:00
New Revision: a3a37848d8cac394d067a6c9c5d545753556e4e7

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

LOG: [mlir] Own the registrations of SourceMgrDiagnosticVerifierHandler (#226258)

`SourceMgrDiagnosticVerifierHandler::registerInContext` registered a
callback capturing `this` and dropped the handler id, so nothing could
erase it. The constructor did the same in its own context. After the
verifier was destroyed, the next diagnostic called into a dead object.
mlir-opt's per-buffer contexts die first, which hid the bug.

Make both registrations owned:

- The constructor installs its callback through `setHandler`, replacing
the base class's printing handler. That handler was unreachable anyway:
the verifier consumes every diagnostic and reports unexpected ones
itself. The base destructor now erases it.
- `registerInContext` returns a scoped registration. Callers destroy it
before the verifier and the context.

Erasing from every registered context in the destructor would not work:
mlir-opt's per-buffer contexts are already gone by then.

mlir-opt's `performActions` finalized the remark policy only after
textual output; on a failing pipeline the remarks deferred by
`--remark-policy=final` were emitted from `~MLIRContext`, after the
handlers were gone. It now finalizes on every exit, so those remarks
reach the verifier, or the printing handler in the non-verify path where
they were dropped before. A lit test pins this.

Assisted-by: Claude Code (Claude Fable 5.1)

Added: 
    mlir/test/Pass/remark-final-verify-diagnostics.mlir

Modified: 
    mlir/include/mlir/IR/Diagnostics.h
    mlir/lib/IR/Diagnostics.cpp
    mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
    mlir/unittests/IR/Diagnostic.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index 276bfd60f0a28..43a80e0ee1c91 100644
--- a/mlir/include/mlir/IR/Diagnostics.h
+++ b/mlir/include/mlir/IR/Diagnostics.h
@@ -694,12 +694,13 @@ class SourceMgrDiagnosticVerifierHandler : public SourceMgrDiagnosticHandler {
   /// verified correctly, failure otherwise.
   LogicalResult verify();
 
-  /// Register this handler with the given context. This is intended for use
-  /// with the splitAndProcessBuffer function.
-  void registerInContext(MLIRContext *ctx);
+  /// Register this handler with `ctx` and return a scoped registration.
+  /// Destroy the returned handle before either this handler or `ctx`.
+  [[nodiscard]] std::unique_ptr<ScopedDiagnosticHandler>
+  registerInContext(MLIRContext *ctx);
 
 private:
-  /// Process a single diagnostic.
+  /// Process a diagnostic and its notes.
   void process(Diagnostic &diag);
 
   /// Process a LocationAttr diagnostic.

diff  --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 0dae9e4c90cf7..a3ddb76a66313 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -879,7 +879,10 @@ SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
   for (unsigned i = 0, e = mgr.getNumBuffers(); i != e; ++i)
     (void)impl->computeExpectedDiags(out, mgr, mgr.getMemoryBuffer(i + 1));
 
-  registerInContext(ctx);
+  // The base class registered a handler that prints every diagnostic. The
+  // verifier takes its place: it consumes diagnostics and reports unexpected
+  // ones itself.
+  setHandler([this](Diagnostic &diag) { process(diag); });
 }
 
 SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
@@ -912,21 +915,18 @@ LogicalResult SourceMgrDiagnosticVerifierHandler::verify() {
   return impl->status;
 }
 
-void SourceMgrDiagnosticVerifierHandler::registerInContext(MLIRContext *ctx) {
-  ctx->getDiagEngine().registerHandler([&](Diagnostic &diag) {
-    // Process the main diagnostics.
-    process(diag);
-
-    // Process each of the notes.
-    for (auto &note : diag.getNotes())
-      process(note);
-  });
+std::unique_ptr<ScopedDiagnosticHandler>
+SourceMgrDiagnosticVerifierHandler::registerInContext(MLIRContext *ctx) {
+  return std::make_unique<ScopedDiagnosticHandler>(
+      ctx, [this](Diagnostic &diag) { process(diag); });
 }
 
-/// Process a single diagnostic.
+/// Process a diagnostic and its notes.
 void SourceMgrDiagnosticVerifierHandler::process(Diagnostic &diag) {
   for (const std::string &str : diag.strs())
     process(diag.getLocation(), str, diag.getSeverity());
+  for (auto &note : diag.getNotes())
+    process(note);
 }
 
 /// Process a diagnostic at a certain location.

diff  --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index 8b49258e135d2..d4812aec50fc1 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -34,6 +34,7 @@
 #include "mlir/Tools/ParseUtilities.h"
 #include "mlir/Tools/Plugins/DialectPlugin.h"
 #include "mlir/Tools/Plugins/PassPlugin.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Remarks/RemarkFormat.h"
 #include "llvm/Support/CommandLine.h"
@@ -580,6 +581,12 @@ performActions(raw_ostream &os,
     break;
   }
   }
+  // Emit the remarks the policy deferred on every exit, while the caller's
+  // diagnostic handlers are still registered; ~MLIRContext would be too late.
+  llvm::scope_exit finalizeRemarks([&ctx] {
+    if (remark::detail::RemarkEngine *engine = ctx.getRemarkEngine())
+      engine->getRemarkEmittingPolicy()->finalize();
+  });
 
   // Prepare the pass manager, applying command-line and reproducer options.
   PassManager pm(op.get()->getName(), PassManager::Nesting::Implicit);
@@ -629,12 +636,6 @@ performActions(raw_ostream &os,
                         config.shouldVerifyPasses() && !pm.empty()),
                     /*locationMap=*/nullptr, &fallbackResourceMap);
   os << OpWithState(op.get(), asmState) << '\n';
-
-  // This is required if the remark policy is final. Otherwise, the remarks are
-  // not emitted.
-  if (remark::detail::RemarkEngine *engine = ctx.getRemarkEngine())
-    engine->getRemarkEmittingPolicy()->finalize();
-
   return success();
 }
 
@@ -661,8 +662,10 @@ processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
   if (threadPool)
     context.setThreadPool(*threadPool);
+  // Keep the registration within the lifetimes of the context and verifier.
+  std::unique_ptr<ScopedDiagnosticHandler> verifierRegistration;
   if (verifyHandler)
-    verifyHandler->registerInContext(&context);
+    verifierRegistration = verifyHandler->registerInContext(&context);
 
   StringRef irdlFile = config.getIrdlFile();
   if (!irdlFile.empty() && failed(loadIRDLDialects(irdlFile, context)))

diff  --git a/mlir/test/Pass/remark-final-verify-diagnostics.mlir b/mlir/test/Pass/remark-final-verify-diagnostics.mlir
new file mode 100644
index 0000000000000..5e6ef7103ed5d
--- /dev/null
+++ b/mlir/test/Pass/remark-final-verify-diagnostics.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s --verify-diagnostics --remarks-filter="category-1-passed" \
+// RUN:   --remark-policy=final \
+// RUN:   --pass-pipeline='builtin.module(test-remark,test-pass-failure)'
+
+// The failing pipeline returns before the textual output; the remark deferred
+// by the final policy must still reach the diagnostic verifier.
+module {
+  // expected-remark at below {{[Passed] test-remark | Category:category-1-passed}}
+  "test.op"() : () -> ()
+}

diff  --git a/mlir/unittests/IR/Diagnostic.cpp b/mlir/unittests/IR/Diagnostic.cpp
index 96e09d3330926..6390cdfe85acc 100644
--- a/mlir/unittests/IR/Diagnostic.cpp
+++ b/mlir/unittests/IR/Diagnostic.cpp
@@ -8,8 +8,12 @@
 
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/Support/TypeID.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
 #include "gtest/gtest.h"
 
+#include <memory>
+
 using namespace mlir;
 using namespace mlir::detail;
 
@@ -60,4 +64,62 @@ TEST(DiagnosticLifetime, TestLazyCopyStringLiteral) {
   ASSERT_STREQ("^rror 1, mutate this", resultMessage.c_str());
 }
 
+// Register counting handlers first so they receive diagnostics only after
+// the verifier's registration is erased.
+
+TEST(SourceMgrDiagnosticVerifierHandler, ScopedRegistration) {
+  MLIRContext own, other;
+  unsigned seen = 0;
+  ScopedDiagnosticHandler counter(&other, [&](Diagnostic &) { ++seen; });
+  llvm::SourceMgr sourceMgr;
+  llvm::raw_null_ostream out;
+  SourceMgrDiagnosticVerifierHandler verifier(
+      sourceMgr, &own, out,
+      SourceMgrDiagnosticVerifierHandler::Level::OnlyExpected);
+  {
+    std::unique_ptr<ScopedDiagnosticHandler> registration =
+        verifier.registerInContext(&other);
+    emitRemark(UnknownLoc::get(&other), "while registered");
+    EXPECT_EQ(seen, 0u);
+  }
+  emitRemark(UnknownLoc::get(&other), "after the handle");
+  EXPECT_EQ(seen, 1u);
+  EXPECT_TRUE(succeeded(verifier.verify()));
+}
+
+TEST(SourceMgrDiagnosticVerifierHandler, ErasesOwnRegistration) {
+  MLIRContext own;
+  unsigned seen = 0;
+  ScopedDiagnosticHandler counter(&own, [&](Diagnostic &) { ++seen; });
+  {
+    llvm::SourceMgr sourceMgr;
+    llvm::raw_null_ostream out;
+    SourceMgrDiagnosticVerifierHandler verifier(
+        sourceMgr, &own, out,
+        SourceMgrDiagnosticVerifierHandler::Level::OnlyExpected);
+    emitRemark(UnknownLoc::get(&own), "while alive");
+    EXPECT_EQ(seen, 0u);
+  }
+  emitRemark(UnknownLoc::get(&own), "after the verifier");
+  EXPECT_EQ(seen, 1u);
+}
+
+TEST(SourceMgrDiagnosticVerifierHandler, OutlivesRegisteredContext) {
+  // Match mlir-opt's per-buffer context lifetime.
+  MLIRContext context;
+  llvm::SourceMgr sourceMgr;
+  llvm::raw_null_ostream out;
+  SourceMgrDiagnosticVerifierHandler verifier(
+      sourceMgr, &context, out,
+      SourceMgrDiagnosticVerifierHandler::Level::OnlyExpected);
+  {
+    MLIRContext perBuffer;
+    std::unique_ptr<ScopedDiagnosticHandler> registration =
+        verifier.registerInContext(&perBuffer);
+    emitRemark(UnknownLoc::get(&perBuffer), "in the temporary context");
+  }
+  emitRemark(UnknownLoc::get(&context), "verifier still usable");
+  EXPECT_TRUE(succeeded(verifier.verify()));
+}
+
 } // namespace


        


More information about the Mlir-commits mailing list