Guard against invalid block erasure. Support forwarding to null listeners.

Author: kuhar

mlir/docs/PatternRewriter.md

@@ -331,12 +331,17 @@ The driver performs a post-order traversal. Note that it walks regions of the
 given op but does not visit the op.
 
 This driver does not (re)visit modified or newly replaced ops, and does not
-allow for progressive rewrites of the same op. Op erasure is only supported for
-the currently matched op. If your pattern set requires these, consider using the
-Greedy Pattern Rewrite Driver instead, at the expense of extra overhead.
+allow for progressive rewrites of the same op. Op and block erasure is only
+supported for the currently matched op and its descendant. If your pattern
+set requires these, consider using the Greedy Pattern Rewrite Driver instead,
+at the expense of extra overhead.
 
 This driver is exposed using the `walkAndApplyPatterns` function.
 
+Note: This driver listens for IR changes via the callbacks provided by
+`RewriterBase`. It is important that patterns announce all IR changes to the
+rewriter and do not bypass the rewriter API by modifying ops directly.
+
 #### Debugging
 
 You can debug the Walk Pattern Rewrite Driver by passing the

mlir/include/mlir/IR/PatternMatch.h

@@ -461,54 +461,60 @@ class RewriterBase : public OpBuilder {
   /// struct can be used as a base to create listener chains, so that multiple
   /// listeners can be notified of IR changes.
   struct ForwardingListener : public RewriterBase::Listener {
-    ForwardingListener(OpBuilder::Listener *listener) : listener(listener) {}
+    ForwardingListener(OpBuilder::Listener *listener)
+        : listener(listener),
+          rewriteListener(
+              dyn_cast_if_present<RewriterBase::Listener>(listener)) {}
 
     void notifyOperationInserted(Operation *op, InsertPoint previous) override {
-      listener->notifyOperationInserted(op, previous);
+      if (listener)
+        listener->notifyOperationInserted(op, previous);
     }
     void notifyBlockInserted(Block *block, Region *previous,
                              Region::iterator previousIt) override {
-      listener->notifyBlockInserted(block, previous, previousIt);
+      if (listener)
+        listener->notifyBlockInserted(block, previous, previousIt);
     }
     void notifyBlockErased(Block *block) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyBlockErased(block);
     }
     void notifyOperationModified(Operation *op) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyOperationModified(op);
     }
     void notifyOperationReplaced(Operation *op, Operation *newOp) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyOperationReplaced(op, newOp);
     }
     void notifyOperationReplaced(Operation *op,
                                  ValueRange replacement) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyOperationReplaced(op, replacement);
     }
     void notifyOperationErased(Operation *op) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyOperationErased(op);
     }
     void notifyPatternBegin(const Pattern &pattern, Operation *op) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyPatternBegin(pattern, op);
     }
     void notifyPatternEnd(const Pattern &pattern,
                           LogicalResult status) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyPatternEnd(pattern, status);
     }
     void notifyMatchFailure(
         Location loc,
         function_ref<void(Diagnostic &)> reasonCallback) override {
-      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+      if (rewriteListener)
         rewriteListener->notifyMatchFailure(loc, reasonCallback);
     }
 
   private:
     OpBuilder::Listener *listener;
+    RewriterBase::Listener *rewriteListener;
   };
 
   /// Move the blocks that belong to "region" before the given position in

mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp

@@ -35,21 +35,34 @@ struct WalkAndApplyPatternsAction final
 };
 
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
-// Forwarding listener to guard against unsupported erasures. Because we use
-// walk-based pattern application, erasing the op from the *next* iteration
-// (e.g., a user of the visited op) is not valid.
-// Note that this is only used with expensive pattern API checks.
+// Forwarding listener to guard against unsupported erasures of non-descendant
+// ops/blocks. Because we use walk-based pattern application, erasing the
+// op/block from the *next* iteration (e.g., a user of the visited op) is not
+// valid. Note that this is only used with expensive pattern API checks.
 struct ErasedOpsListener final : RewriterBase::ForwardingListener {
   using RewriterBase::ForwardingListener::ForwardingListener;
 
   void notifyOperationErased(Operation *op) override {
-    if (op != visitedOp)
-      llvm::report_fatal_error("unsupported op erased in WalkPatternRewriter; "
-                               "erasure is only supported for matched ops");
-
+    checkErasure(op);
     ForwardingListener::notifyOperationErased(op);
   }
 
+  void notifyBlockErased(Block *block) override {
+    checkErasure(block->getParentOp());
+    ForwardingListener::notifyBlockErased(block);
+  }
+
+  void checkErasure(Operation *op) const {
+    Operation *ancestorOp = op;
+    while (ancestorOp && ancestorOp != visitedOp)
+      ancestorOp = ancestorOp->getParentOp();
+
+    if (ancestorOp != visitedOp)
+      llvm::report_fatal_error(
+          "unsupported erased in WalkPatternRewriter; "
+          "erasure is only supported for matched ops and their descendants");
+  }
+
   Operation *visitedOp = nullptr;
 };
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS

mlir/test/IR/test-walk-pattern-rewrite-driver.mlir

@@ -105,3 +105,17 @@ func.func @replace_with_new_op() -> i32 {
   %res = arith.addi %a, %a : i32
   return %res : i32
 }
+
+// Check that we can erase nested blocks.
+// CHECK-LABEL: func.func @erase_nested_block
+// CHECK:         %[[RES:.+]] = "test.erase_first_block"
+// CHECK-NEXT:    foo.bar
+// CHECK:         return %[[RES]]
+func.func @erase_nested_block() -> i32 {
+  %a = "test.erase_first_block"() ({
+    "foo.foo"() : () -> ()
+    ^bb1:
+    "foo.bar"() : () -> ()
+  }): () -> (i32)
+  return %a : i32
+}

mlir/test/lib/Dialect/Test/TestPatterns.cpp

@@ -342,6 +342,28 @@ class ReplaceWithNewOp : public RewritePattern {
   }
 };
 
+/// Erases the first child block of the matched "test.erase_first_block"
+/// operation.
+class EraseFirstBlock : public RewritePattern {
+public:
+  EraseFirstBlock(MLIRContext *context)
+      : RewritePattern("test.erase_first_block", /*benefit=*/1, context) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    llvm::errs() << "Num regions: " << op->getNumRegions() << "\n";
+    for (Region &r : op->getRegions()) {
+      for (Block &b : r.getBlocks()) {
+        rewriter.eraseBlock(&b);
+        llvm::errs() << "Erasing block: " << b << "\n";
+        return success();
+      }
+    }
+
+    return failure();
+  }
+};
+
 struct TestGreedyPatternDriver
     : public PassWrapper<TestGreedyPatternDriver, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGreedyPatternDriver)
@@ -608,7 +630,8 @@ struct TestWalkPatternDriver final
 
     // Patterns for testing the WalkPatternRewriteDriver.
     patterns.add<IncrementIntAttribute<3>, MoveBeforeParentOp,
-                 MoveAfterParentOp, CloneOp, ReplaceWithNewOp>(&getContext());
+                 MoveAfterParentOp, CloneOp, ReplaceWithNewOp, EraseFirstBlock>(
+        &getContext());
 
     DumpNotifications dumpListener;
     walkAndApplyPatterns(getOperation(), std::move(patterns),