Add support to GreedyPatternRewriter for erasing unreachable blocks.

Rewrite patterns may make modifications to the CFG, including dropping edges between blocks. This change adds a simple unreachable block elimination run at the end of each iteration to ensure that the CFG remains valid.

PiperOrigin-RevId: 277545805

1 files changed, 56 insertions, 0 deletions

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 56701e39680..8c1fddf3a81 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -22,8 +22,10 @@
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -113,6 +115,56 @@ protected:
   }
 
 private:
+  /// Erase the unreachable blocks within the provided regions. Returns success
+  /// if any blocks were erased, failure otherwise.
+  LogicalResult eraseUnreachableBlocks(MutableArrayRef<Region> regions) {
+    // Set of blocks found to be reachable within a given region.
+    llvm::df_iterator_default_set<Block *, 16> reachable;
+    // If any blocks were found to be dead.
+    bool erasedDeadBlocks = false;
+
+    SmallVector<Region *, 1> worklist;
+    worklist.reserve(regions.size());
+    for (Region &region : regions)
+      worklist.push_back(&region);
+    while (!worklist.empty()) {
+      Region *region = worklist.pop_back_val();
+      if (region->empty())
+        continue;
+
+      // If this is a single block region, just collect the nested regions.
+      if (std::next(region->begin()) == region->end()) {
+        for (Operation &op : region->front())
+          for (Region &region : op.getRegions())
+            worklist.push_back(&region);
+        continue;
+      }
+
+      // Mark all reachable blocks.
+      reachable.clear();
+      for (Block *block : depth_first_ext(&region->front(), reachable))
+        (void)block /* Mark all reachable blocks */;
+
+      // Collect all of the dead blocks and push the live regions onto the
+      // worklist.
+      for (Block &block : llvm::make_early_inc_range(*region)) {
+        if (!reachable.count(&block)) {
+          block.dropAllDefinedValueUses();
+          block.erase();
+          erasedDeadBlocks = true;
+          continue;
+        }
+
+        // Walk any regions within this block.
+        for (Operation &op : block)
+          for (Region &region : op.getRegions())
+            worklist.push_back(&region);
+      }
+    }
+
+    return success(erasedDeadBlocks);
+  }
+
   // Look over the provided operands for any defining operations that should
   // be re-added to the worklist. This function should be called when an
   // operation is modified or removed, as it may trigger further
@@ -209,6 +261,10 @@ bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions,
       // notified of any necessary changes, so there is nothing else to do here.
       changed |= matcher.matchAndRewrite(op, *this);
     }
+
+    // After applying patterns, make sure that the CFG of each of the regions is
+    // kept up to date.
+    changed |= succeeded(eraseUnreachableBlocks(regions));
   } while (changed && ++i < maxIterations);
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;