[Mlir-commits] [mlir] [mlir][RemoveDeadValues] Simplify branch op handling using ub.poison (PR #182711)

[Fedor Nikolaev]

================
@@ -808,20 +740,17 @@ void RemoveDeadValues::runOnOperation() {
   if (!canonicalize)
     return;
 
-  // Canonicalize all region branch ops.
-  SmallVector<Operation *> opsToCanonicalize;
-  module->walk([&](RegionBranchOpInterface regionBranchOp) {
-    opsToCanonicalize.push_back(regionBranchOp.getOperation());
-  });
-  // Collect all canonicalization patterns for region branch ops.
+  // Canonicalize all region branch ops and branch ops.
   RewritePatternSet owningPatterns(context);
   DenseSet<RegisteredOperationName> populatedPatterns;
-  for (Operation *op : opsToCanonicalize)
+  module->walk([&](Operation *op) {
+    if (!isa<RegionBranchOpInterface, BranchOpInterface>(op))
+      return;
     if (std::optional<RegisteredOperationName> info = op->getRegisteredInfo())
       if (populatedPatterns.insert(*info).second)
         info->getCanonicalizationPatterns(owningPatterns, context);
-  if (failed(applyOpPatternsGreedily(opsToCanonicalize,
-                                     std::move(owningPatterns)))) {
+  });
+  if (failed(applyPatternsGreedily(module, std::move(owningPatterns)))) {
----------------
felichita wrote:

The old approach collected a list of ops and applied patterns only to those, but at the same time new canonicalizer needs to propagate across block boundaries to eliminate the corresponding block arguments. `applyOpPatternsGreedily` on a fixed op list can't do that `applyPatternsGreedily` on the module allows patterns to canonicalize wherever needed

https://github.com/llvm/llvm-project/pull/182711