[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Add missing erasure notifications (PR #145030)

2025-06-20 Thread Jeremy Kun via llvm-branch-commits

https://github.com/j2kun approved this pull request.

LGTM!

https://github.com/llvm/llvm-project/pull/145030
___
llvm-branch-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Add missing erasure notifications (PR #145030)

2025-06-20 Thread via llvm-branch-commits

llvmbot wrote:




@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)


Changes

Add missing listener notifications when erasing nested blocks/operations.

This commit also moves some of the functionality from 
`ConversionPatternRewriter` to `ConversionPatternRewriterImpl`. This is in 
preparation of the One-Shot Dialect Conversion refactoring: The implementations 
in `ConversionPatternRewriter` should be as simple as possible, so that a 
switch between "rollback allowed" and "rollback not allowed" can be inserted at 
that level. (In the latter case, `ConversionPatternRewriterImpl` can be 
bypassed to some degree, and `PatternRewriter::eraseBlock` etc. can be used.)



---
Full diff: https://github.com/llvm/llvm-project/pull/145030.diff


2 Files Affected:

- (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+43-19) 
- (modified) mlir/test/Transforms/test-legalizer.mlir (+16-2) 


``diff
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp 
b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index ff48647f43305..7419d79cd8856 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -274,6 +274,26 @@ struct RewriterState {
 // IR rewrites
 
//===--===//
 
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op);
+
+/// Notify the listener that the given block and its contents are being erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Block &b) {
+  for (Operation &op : b)
+notifyIRErased(listener, op);
+  listener->notifyBlockErased(&b);
+}
+
+/// Notify the listener that the given operation and its contents are being
+/// erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op) {
+  for (Region &r : op.getRegions()) {
+for (Block &b : r) {
+  notifyIRErased(listener, b);
+}
+  }
+  listener->notifyOperationErased(&op);
+}
+
 /// An IR rewrite that can be committed (upon success) or rolled back (upon
 /// failure).
 ///
@@ -422,17 +442,20 @@ class EraseBlockRewrite : public BlockRewrite {
   }
 
   void commit(RewriterBase &rewriter) override {
-// Erase the block.
 assert(block && "expected block");
-assert(block->empty() && "expected empty block");
 
-// Notify the listener that the block is about to be erased.
+// Notify the listener that the block and its contents are being erased.
 if (auto *listener =
 dyn_cast_or_null(rewriter.getListener()))
-  listener->notifyBlockErased(block);
+  notifyIRErased(listener, *block);
   }
 
   void cleanup(RewriterBase &rewriter) override {
+// Erase the contents of the block.
+for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block)))
+  rewriter.eraseOp(&op);
+assert(block->empty() && "expected empty block");
+
 // Erase the block.
 block->dropAllDefinedValueUses();
 delete block;
@@ -1147,12 +1170,9 @@ void ReplaceOperationRewrite::commit(RewriterBase 
&rewriter) {
   if (getConfig().unlegalizedOps)
 getConfig().unlegalizedOps->erase(op);
 
-  // Notify the listener that the operation (and its nested operations) was
-  // erased.
-  if (listener) {
-op->walk(
-[&](Operation *op) { listener->notifyOperationErased(op); });
-  }
+  // Notify the listener that the operation and its contents are being erased.
+  if (listener)
+notifyIRErased(listener, *op);
 
   // Do not erase the operation yet. It may still be referenced in `mapping`.
   // Just unlink it for now and erase it during cleanup.
@@ -1605,6 +1625,8 @@ void ConversionPatternRewriterImpl::replaceOp(
 }
 
 void ConversionPatternRewriterImpl::eraseBlock(Block *block) {
+  assert(!wasOpReplaced(block->getParentOp()) &&
+ "attempting to erase a block within a replaced/erased op");
   appendRewrite(block);
 
   // Unlink the block from its parent region. The block is kept in the rewrite
@@ -1612,12 +1634,16 @@ void ConversionPatternRewriterImpl::eraseBlock(Block 
*block) {
   // allows us to keep the operations in the block live and undo the removal by
   // re-inserting the block.
   block->getParent()->getBlocks().remove(block);
+
+  // Mark all nested ops as erased.
+  block->walk([&](Operation *op) { replacedOps.insert(op); });
 }
 
 void ConversionPatternRewriterImpl::notifyBlockInserted(
 Block *block, Region *previous, Region::iterator previousIt) {
-  assert(!wasOpReplaced(block->getParentOp()) &&
- "attempting to insert into a region within a replaced/erased op");
+  assert(
+  (!config.allowPatternRollback || !wasOpReplaced(block->getParentOp())) &&
+  "attempting to insert into a region within a replaced/erased op");
   LLVM_DEBUG(
   {
 Operation *parent = block->getParentOp();
@@ -1630,6 +1656,11 @@ void ConversionPatternRewriterImpl::notifyBlockInserted(
 }
   });
 
+  if (!config.allowPa

[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Add missing erasure notifications (PR #145030)

2025-06-20 Thread Matthias Springer via llvm-branch-commits

https://github.com/matthias-springer updated 
https://github.com/llvm/llvm-project/pull/145030

>From d65161f2d8e65512a6924ac96f069ab5acce0fcd Mon Sep 17 00:00:00 2001
From: Matthias Springer 
Date: Fri, 20 Jun 2025 12:25:00 +
Subject: [PATCH] [mlir][Transforms] Dialect conversion: Add missing erasure
 notifications

---
 .../Transforms/Utils/DialectConversion.cpp| 57 ---
 mlir/test/Transforms/test-legalizer.mlir  | 18 +-
 2 files changed, 54 insertions(+), 21 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp 
b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index ff48647f43305..9c75b7436cddc 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -274,6 +274,26 @@ struct RewriterState {
 // IR rewrites
 
//===--===//
 
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op);
+
+/// Notify the listener that the given block and its contents are being erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Block &b) {
+  for (Operation &op : b)
+notifyIRErased(listener, op);
+  listener->notifyBlockErased(&b);
+}
+
+/// Notify the listener that the given operation and its contents are being
+/// erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op) {
+  for (Region &r : op.getRegions()) {
+for (Block &b : r) {
+  notifyIRErased(listener, b);
+}
+  }
+  listener->notifyOperationErased(&op);
+}
+
 /// An IR rewrite that can be committed (upon success) or rolled back (upon
 /// failure).
 ///
@@ -422,17 +442,20 @@ class EraseBlockRewrite : public BlockRewrite {
   }
 
   void commit(RewriterBase &rewriter) override {
-// Erase the block.
 assert(block && "expected block");
-assert(block->empty() && "expected empty block");
 
-// Notify the listener that the block is about to be erased.
+// Notify the listener that the block and its contents are being erased.
 if (auto *listener =
 dyn_cast_or_null(rewriter.getListener()))
-  listener->notifyBlockErased(block);
+  notifyIRErased(listener, *block);
   }
 
   void cleanup(RewriterBase &rewriter) override {
+// Erase the contents of the block.
+for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block)))
+  rewriter.eraseOp(&op);
+assert(block->empty() && "expected empty block");
+
 // Erase the block.
 block->dropAllDefinedValueUses();
 delete block;
@@ -1147,12 +1170,9 @@ void ReplaceOperationRewrite::commit(RewriterBase 
&rewriter) {
   if (getConfig().unlegalizedOps)
 getConfig().unlegalizedOps->erase(op);
 
-  // Notify the listener that the operation (and its nested operations) was
-  // erased.
-  if (listener) {
-op->walk(
-[&](Operation *op) { listener->notifyOperationErased(op); });
-  }
+  // Notify the listener that the operation and its contents are being erased.
+  if (listener)
+notifyIRErased(listener, *op);
 
   // Do not erase the operation yet. It may still be referenced in `mapping`.
   // Just unlink it for now and erase it during cleanup.
@@ -1605,6 +1625,8 @@ void ConversionPatternRewriterImpl::replaceOp(
 }
 
 void ConversionPatternRewriterImpl::eraseBlock(Block *block) {
+  assert(!wasOpReplaced(block->getParentOp()) &&
+ "attempting to erase a block within a replaced/erased op");
   appendRewrite(block);
 
   // Unlink the block from its parent region. The block is kept in the rewrite
@@ -1612,12 +1634,16 @@ void ConversionPatternRewriterImpl::eraseBlock(Block 
*block) {
   // allows us to keep the operations in the block live and undo the removal by
   // re-inserting the block.
   block->getParent()->getBlocks().remove(block);
+
+  // Mark all nested ops as erased.
+  block->walk([&](Operation *op) { replacedOps.insert(op); });
 }
 
 void ConversionPatternRewriterImpl::notifyBlockInserted(
 Block *block, Region *previous, Region::iterator previousIt) {
-  assert(!wasOpReplaced(block->getParentOp()) &&
- "attempting to insert into a region within a replaced/erased op");
+  assert(
+  (!config.allowPatternRollback || !wasOpReplaced(block->getParentOp())) &&
+  "attempting to insert into a region within a replaced/erased op");
   LLVM_DEBUG(
   {
 Operation *parent = block->getParentOp();
@@ -1709,13 +1735,6 @@ void ConversionPatternRewriter::eraseOp(Operation *op) {
 }
 
 void ConversionPatternRewriter::eraseBlock(Block *block) {
-  assert(!impl->wasOpReplaced(block->getParentOp()) &&
- "attempting to erase a block within a replaced/erased op");
-
-  // Mark all ops for erasure.
-  for (Operation &op : *block)
-eraseOp(&op);
-
   impl->eraseBlock(block);
 }
 
diff --git a/mlir/test/Transforms/test-legalizer.mlir 
b/mlir/test/Transforms/test-legalizer.mlir
index 34948ae685f0a..204c8c1456826 100644
--- a/mlir/test/Transforms/tes

[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Add missing erasure notifications (PR #145030)

2025-06-20 Thread Matthias Springer via llvm-branch-commits

https://github.com/matthias-springer updated 
https://github.com/llvm/llvm-project/pull/145030

>From edb49ecf11faa51847b324d6e43336845e71fcf4 Mon Sep 17 00:00:00 2001
From: Matthias Springer 
Date: Fri, 20 Jun 2025 12:25:00 +
Subject: [PATCH] [mlir][Transforms] Dialect conversion: Add missing erasure
 notifications

---
 .../Transforms/Utils/DialectConversion.cpp| 52 +--
 mlir/test/Transforms/test-legalizer.mlir  | 18 ++-
 2 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp 
b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index ff48647f43305..ad82a007b7996 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -274,6 +274,26 @@ struct RewriterState {
 // IR rewrites
 
//===--===//
 
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op);
+
+/// Notify the listener that the given block and its contents are being erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Block &b) {
+  for (Operation &op : b)
+notifyIRErased(listener, op);
+  listener->notifyBlockErased(&b);
+}
+
+/// Notify the listener that the given operation and its contents are being
+/// erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op) {
+  for (Region &r : op.getRegions()) {
+for (Block &b : r) {
+  notifyIRErased(listener, b);
+}
+  }
+  listener->notifyOperationErased(&op);
+}
+
 /// An IR rewrite that can be committed (upon success) or rolled back (upon
 /// failure).
 ///
@@ -422,17 +442,20 @@ class EraseBlockRewrite : public BlockRewrite {
   }
 
   void commit(RewriterBase &rewriter) override {
-// Erase the block.
 assert(block && "expected block");
-assert(block->empty() && "expected empty block");
 
-// Notify the listener that the block is about to be erased.
+// Notify the listener that the block and its contents are being erased.
 if (auto *listener =
 dyn_cast_or_null(rewriter.getListener()))
-  listener->notifyBlockErased(block);
+  notifyIRErased(listener, *block);
   }
 
   void cleanup(RewriterBase &rewriter) override {
+// Erase the contents of the block.
+for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block)))
+  rewriter.eraseOp(&op);
+assert(block->empty() && "expected empty block");
+
 // Erase the block.
 block->dropAllDefinedValueUses();
 delete block;
@@ -1147,12 +1170,9 @@ void ReplaceOperationRewrite::commit(RewriterBase 
&rewriter) {
   if (getConfig().unlegalizedOps)
 getConfig().unlegalizedOps->erase(op);
 
-  // Notify the listener that the operation (and its nested operations) was
-  // erased.
-  if (listener) {
-op->walk(
-[&](Operation *op) { listener->notifyOperationErased(op); });
-  }
+  // Notify the listener that the operation and its contents are being erased.
+  if (listener)
+notifyIRErased(listener, *op);
 
   // Do not erase the operation yet. It may still be referenced in `mapping`.
   // Just unlink it for now and erase it during cleanup.
@@ -1605,6 +1625,8 @@ void ConversionPatternRewriterImpl::replaceOp(
 }
 
 void ConversionPatternRewriterImpl::eraseBlock(Block *block) {
+  assert(!wasOpReplaced(block->getParentOp()) &&
+ "attempting to erase a block within a replaced/erased op");
   appendRewrite(block);
 
   // Unlink the block from its parent region. The block is kept in the rewrite
@@ -1612,6 +1634,9 @@ void ConversionPatternRewriterImpl::eraseBlock(Block 
*block) {
   // allows us to keep the operations in the block live and undo the removal by
   // re-inserting the block.
   block->getParent()->getBlocks().remove(block);
+
+  // Mark all nested ops as erased.
+  block->walk([&](Operation *op) { replacedOps.insert(op); });
 }
 
 void ConversionPatternRewriterImpl::notifyBlockInserted(
@@ -1709,13 +1734,6 @@ void ConversionPatternRewriter::eraseOp(Operation *op) {
 }
 
 void ConversionPatternRewriter::eraseBlock(Block *block) {
-  assert(!impl->wasOpReplaced(block->getParentOp()) &&
- "attempting to erase a block within a replaced/erased op");
-
-  // Mark all ops for erasure.
-  for (Operation &op : *block)
-eraseOp(&op);
-
   impl->eraseBlock(block);
 }
 
diff --git a/mlir/test/Transforms/test-legalizer.mlir 
b/mlir/test/Transforms/test-legalizer.mlir
index 34948ae685f0a..204c8c1456826 100644
--- a/mlir/test/Transforms/test-legalizer.mlir
+++ b/mlir/test/Transforms/test-legalizer.mlir
@@ -461,12 +461,26 @@ func.func @convert_detached_signature() {
 
 // -
 
+// CHECK: notifyOperationReplaced: test.erase_op
+// CHECK: notifyOperationErased: test.dummy_op_lvl_2
+// CHECK: notifyBlockErased
+// CHECK: notifyOperationErased: test.dummy_op_lvl_1
+// CHECK: notifyBlockErased
+// CHECK: notifyOperationErased: test.erase_op
+// CHECK: notifyOperationInser

[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Add missing erasure notifications (PR #145030)

2025-06-20 Thread Matthias Springer via llvm-branch-commits

https://github.com/matthias-springer edited 
https://github.com/llvm/llvm-project/pull/145030
___
llvm-branch-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Add missing erasure notifications (PR #145030)

2025-06-20 Thread via llvm-branch-commits

llvmbot wrote:




@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)


Changes

Add missing listener notifications when erasing nested blocks/operations.

This commit also moves some of the functionality from 
`ConversionPatternRewriter` to `ConversionPatternRewriterImpl`. This is in 
preparation of the One-Shot Dialect Conversion refactoring: The implementations 
in `ConversionPatternRewriter` should be as simple as possible, so that a 
switch between "rollback allowed" and "rollback not allowed" can be inserted at 
that level. (In the latter case, `ConversionPatternRewriterImpl` can be 
bypassed to some degree, and `PatternRewriter::eraseBlock` etc. can be used.)



---
Full diff: https://github.com/llvm/llvm-project/pull/145030.diff


2 Files Affected:

- (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+43-19) 
- (modified) mlir/test/Transforms/test-legalizer.mlir (+16-2) 


``diff
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp 
b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index ff48647f43305..7419d79cd8856 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -274,6 +274,26 @@ struct RewriterState {
 // IR rewrites
 
//===--===//
 
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op);
+
+/// Notify the listener that the given block and its contents are being erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Block &b) {
+  for (Operation &op : b)
+notifyIRErased(listener, op);
+  listener->notifyBlockErased(&b);
+}
+
+/// Notify the listener that the given operation and its contents are being
+/// erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op) {
+  for (Region &r : op.getRegions()) {
+for (Block &b : r) {
+  notifyIRErased(listener, b);
+}
+  }
+  listener->notifyOperationErased(&op);
+}
+
 /// An IR rewrite that can be committed (upon success) or rolled back (upon
 /// failure).
 ///
@@ -422,17 +442,20 @@ class EraseBlockRewrite : public BlockRewrite {
   }
 
   void commit(RewriterBase &rewriter) override {
-// Erase the block.
 assert(block && "expected block");
-assert(block->empty() && "expected empty block");
 
-// Notify the listener that the block is about to be erased.
+// Notify the listener that the block and its contents are being erased.
 if (auto *listener =
 dyn_cast_or_null(rewriter.getListener()))
-  listener->notifyBlockErased(block);
+  notifyIRErased(listener, *block);
   }
 
   void cleanup(RewriterBase &rewriter) override {
+// Erase the contents of the block.
+for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block)))
+  rewriter.eraseOp(&op);
+assert(block->empty() && "expected empty block");
+
 // Erase the block.
 block->dropAllDefinedValueUses();
 delete block;
@@ -1147,12 +1170,9 @@ void ReplaceOperationRewrite::commit(RewriterBase 
&rewriter) {
   if (getConfig().unlegalizedOps)
 getConfig().unlegalizedOps->erase(op);
 
-  // Notify the listener that the operation (and its nested operations) was
-  // erased.
-  if (listener) {
-op->walk(
-[&](Operation *op) { listener->notifyOperationErased(op); });
-  }
+  // Notify the listener that the operation and its contents are being erased.
+  if (listener)
+notifyIRErased(listener, *op);
 
   // Do not erase the operation yet. It may still be referenced in `mapping`.
   // Just unlink it for now and erase it during cleanup.
@@ -1605,6 +1625,8 @@ void ConversionPatternRewriterImpl::replaceOp(
 }
 
 void ConversionPatternRewriterImpl::eraseBlock(Block *block) {
+  assert(!wasOpReplaced(block->getParentOp()) &&
+ "attempting to erase a block within a replaced/erased op");
   appendRewrite(block);
 
   // Unlink the block from its parent region. The block is kept in the rewrite
@@ -1612,12 +1634,16 @@ void ConversionPatternRewriterImpl::eraseBlock(Block 
*block) {
   // allows us to keep the operations in the block live and undo the removal by
   // re-inserting the block.
   block->getParent()->getBlocks().remove(block);
+
+  // Mark all nested ops as erased.
+  block->walk([&](Operation *op) { replacedOps.insert(op); });
 }
 
 void ConversionPatternRewriterImpl::notifyBlockInserted(
 Block *block, Region *previous, Region::iterator previousIt) {
-  assert(!wasOpReplaced(block->getParentOp()) &&
- "attempting to insert into a region within a replaced/erased op");
+  assert(
+  (!config.allowPatternRollback || !wasOpReplaced(block->getParentOp())) &&
+  "attempting to insert into a region within a replaced/erased op");
   LLVM_DEBUG(
   {
 Operation *parent = block->getParentOp();
@@ -1630,6 +1656,11 @@ void ConversionPatternRewriterImpl::notifyBlockInserted(
 }
   });
 
+  if (!config.allowPattern

[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Add missing erasure notifications (PR #145030)

2025-06-20 Thread Matthias Springer via llvm-branch-commits

https://github.com/matthias-springer created 
https://github.com/llvm/llvm-project/pull/145030

Add missing listener notifications when erasing nested blocks/operations.

This commit also moves some of the functionality from 
`ConversionPatternRewriter` to `ConversionPatternRewriterImpl`. This is in 
preparation of the One-Shot Dialect Conversion refactoring: The implementations 
in `ConversionPatternRewriter` should be as simple as possible, so that a 
switch between "rollback allowed" and "rollback not allowed" can be inserted at 
that level. (In the latter case, `ConversionPatternRewriterImpl` can be 
bypassed to some degree, and `PatternRewriter::eraseBlock` etc. can be used.)



>From 40dea2a59a6fcc49976488b106109f22df8707f0 Mon Sep 17 00:00:00 2001
From: Matthias Springer 
Date: Fri, 20 Jun 2025 12:25:00 +
Subject: [PATCH] [mlir][Transforms] Dialect conversion: Add missing erasure
 notifications

---
 .../Transforms/Utils/DialectConversion.cpp| 62 +--
 mlir/test/Transforms/test-legalizer.mlir  | 18 +-
 2 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp 
b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index ff48647f43305..7419d79cd8856 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -274,6 +274,26 @@ struct RewriterState {
 // IR rewrites
 
//===--===//
 
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op);
+
+/// Notify the listener that the given block and its contents are being erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Block &b) {
+  for (Operation &op : b)
+notifyIRErased(listener, op);
+  listener->notifyBlockErased(&b);
+}
+
+/// Notify the listener that the given operation and its contents are being
+/// erased.
+static void notifyIRErased(RewriterBase::Listener *listener, Operation &op) {
+  for (Region &r : op.getRegions()) {
+for (Block &b : r) {
+  notifyIRErased(listener, b);
+}
+  }
+  listener->notifyOperationErased(&op);
+}
+
 /// An IR rewrite that can be committed (upon success) or rolled back (upon
 /// failure).
 ///
@@ -422,17 +442,20 @@ class EraseBlockRewrite : public BlockRewrite {
   }
 
   void commit(RewriterBase &rewriter) override {
-// Erase the block.
 assert(block && "expected block");
-assert(block->empty() && "expected empty block");
 
-// Notify the listener that the block is about to be erased.
+// Notify the listener that the block and its contents are being erased.
 if (auto *listener =
 dyn_cast_or_null(rewriter.getListener()))
-  listener->notifyBlockErased(block);
+  notifyIRErased(listener, *block);
   }
 
   void cleanup(RewriterBase &rewriter) override {
+// Erase the contents of the block.
+for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block)))
+  rewriter.eraseOp(&op);
+assert(block->empty() && "expected empty block");
+
 // Erase the block.
 block->dropAllDefinedValueUses();
 delete block;
@@ -1147,12 +1170,9 @@ void ReplaceOperationRewrite::commit(RewriterBase 
&rewriter) {
   if (getConfig().unlegalizedOps)
 getConfig().unlegalizedOps->erase(op);
 
-  // Notify the listener that the operation (and its nested operations) was
-  // erased.
-  if (listener) {
-op->walk(
-[&](Operation *op) { listener->notifyOperationErased(op); });
-  }
+  // Notify the listener that the operation and its contents are being erased.
+  if (listener)
+notifyIRErased(listener, *op);
 
   // Do not erase the operation yet. It may still be referenced in `mapping`.
   // Just unlink it for now and erase it during cleanup.
@@ -1605,6 +1625,8 @@ void ConversionPatternRewriterImpl::replaceOp(
 }
 
 void ConversionPatternRewriterImpl::eraseBlock(Block *block) {
+  assert(!wasOpReplaced(block->getParentOp()) &&
+ "attempting to erase a block within a replaced/erased op");
   appendRewrite(block);
 
   // Unlink the block from its parent region. The block is kept in the rewrite
@@ -1612,12 +1634,16 @@ void ConversionPatternRewriterImpl::eraseBlock(Block 
*block) {
   // allows us to keep the operations in the block live and undo the removal by
   // re-inserting the block.
   block->getParent()->getBlocks().remove(block);
+
+  // Mark all nested ops as erased.
+  block->walk([&](Operation *op) { replacedOps.insert(op); });
 }
 
 void ConversionPatternRewriterImpl::notifyBlockInserted(
 Block *block, Region *previous, Region::iterator previousIt) {
-  assert(!wasOpReplaced(block->getParentOp()) &&
- "attempting to insert into a region within a replaced/erased op");
+  assert(
+  (!config.allowPatternRollback || !wasOpReplaced(block->getParentOp())) &&
+  "attempting to insert into a region within a replaced/erased op");
   LLVM_DEBUG(
   {
 Operat