3 files changed, 35 insertions, 24 deletions

diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 3e33795..776b5c6 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -821,15 +821,7 @@ SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
   for (unsigned i = 0, e = mgr.getNumBuffers(); i != e; ++i)
     (void)impl->computeExpectedDiags(out, mgr, mgr.getMemoryBuffer(i + 1));
 
-  // Register a handler to verify the diagnostics.
-  setHandler([&](Diagnostic &diag) {
-    // Process the main diagnostics.
-    process(diag);
-
-    // Process each of the notes.
-    for (auto &note : diag.getNotes())
-      process(note);
-  });
+  registerInContext(ctx);
 }
 
 SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
@@ -862,6 +854,17 @@ 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);
+  });
+}
+
 /// Process a single diagnostic.
 void SourceMgrDiagnosticVerifierHandler::process(Diagnostic &diag) {
   return process(diag.getLocation(), diag.str(), diag.getSeverity());
diff --git a/mlir/lib/IR/PatternLoggingListener.cpp b/mlir/lib/IR/PatternLoggingListener.cpp
index ce2123a..0db13ab 100644
--- a/mlir/lib/IR/PatternLoggingListener.cpp
+++ b/mlir/lib/IR/PatternLoggingListener.cpp
@@ -1,50 +1,48 @@
 #include "mlir/IR/PatternMatch.h"
-#include "llvm/Support/Debug.h"
+#include "llvm/Support/DebugLog.h"
 
 #define DEBUG_TYPE "pattern-logging-listener"
-#define DBGS() (llvm::dbgs() << "[" << DEBUG_TYPE << "] ")
-#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
 using namespace mlir;
 
 void RewriterBase::PatternLoggingListener::notifyOperationInserted(
     Operation *op, InsertPoint previous) {
-  LDBG(patternName << " | notifyOperationInserted"
-                   << " | " << op->getName());
+  LDBG() << patternName << " | notifyOperationInserted"
+         << " | " << op->getName();
   ForwardingListener::notifyOperationInserted(op, previous);
 }
 
 void RewriterBase::PatternLoggingListener::notifyOperationModified(
     Operation *op) {
-  LDBG(patternName << " | notifyOperationModified"
-                   << " | " << op->getName());
+  LDBG() << patternName << " | notifyOperationModified"
+         << " | " << op->getName();
   ForwardingListener::notifyOperationModified(op);
 }
 
 void RewriterBase::PatternLoggingListener::notifyOperationReplaced(
     Operation *op, Operation *newOp) {
-  LDBG(patternName << " | notifyOperationReplaced (with op)"
-                   << " | " << op->getName() << " | " << newOp->getName());
+  LDBG() << patternName << " | notifyOperationReplaced (with op)"
+         << " | " << op->getName() << " | " << newOp->getName();
   ForwardingListener::notifyOperationReplaced(op, newOp);
 }
 
 void RewriterBase::PatternLoggingListener::notifyOperationReplaced(
     Operation *op, ValueRange replacement) {
-  LDBG(patternName << " | notifyOperationReplaced (with values)"
-                   << " | " << op->getName());
+  LDBG() << patternName << " | notifyOperationReplaced (with values)"
+         << " | " << op->getName();
   ForwardingListener::notifyOperationReplaced(op, replacement);
 }
 
 void RewriterBase::PatternLoggingListener::notifyOperationErased(
     Operation *op) {
-  LDBG(patternName << " | notifyOperationErased"
-                   << " | " << op->getName());
+  LDBG() << patternName << " | notifyOperationErased"
+         << " | " << op->getName();
   ForwardingListener::notifyOperationErased(op);
 }
 
 void RewriterBase::PatternLoggingListener::notifyPatternBegin(
     const Pattern &pattern, Operation *op) {
-  LDBG(patternName << " | notifyPatternBegin"
-                   << " | " << op->getName());
+  LDBG() << patternName << " | notifyPatternBegin"
+         << " | " << op->getName();
   ForwardingListener::notifyPatternBegin(pattern, op);
 }
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 5c98417..9332f55 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -156,6 +156,11 @@ void RewriterBase::eraseOp(Operation *op) {
   assert(op->use_empty() && "expected 'op' to have no uses");
   auto *rewriteListener = dyn_cast_if_present<Listener>(listener);
 
+  // If the current insertion point is before the erased operation, we adjust
+  // the insertion point to be after the operation.
+  if (getInsertionPoint() == op->getIterator())
+    setInsertionPointAfter(op);
+
   // Fast path: If no listener is attached, the op can be dropped in one go.
   if (!rewriteListener) {
     op->erase();
@@ -320,6 +325,11 @@ void RewriterBase::inlineBlockBefore(Block *source, Block *dest,
       moveOpBefore(&source->front(), dest, before);
   }
 
+  // If the current insertion point is within the source block, adjust the
+  // insertion point to the destination block.
+  if (getInsertionBlock() == source)
+    setInsertionPoint(dest, getInsertionPoint());
+
   // Erase the source block.
   assert(source->empty() && "expected 'source' to be empty");
   eraseBlock(source);