[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:

Here is a comparison of the three approaches on the same input:

```
// 1. main branch (old direct removal):
  func.func @acceptable_ir_has_cleanable_loop_of_conditional_and_branch_op(%arg0: i1) {
    cf.br ^bb1
  ^bb1:  // 2 preds: ^bb0, ^bb2
    cf.br ^bb2
  ^bb2:  // pred: ^bb1
    cf.cond_br %arg0, ^bb1, ^bb3
  ^bb3:  // pred: ^bb2
    return
  }

// -----

// 2. module-level traversal (this PR):
  func.func @acceptable_ir_has_cleanable_loop_of_conditional_and_branch_op(%arg0: i1) {
    cf.br ^bb1
  ^bb1:  // 2 preds: ^bb0, ^bb1
    cf.cond_br %arg0, ^bb1, ^bb2
  ^bb2:  // pred: ^bb1
    return
  }

// -----

// 3. fixed op list (BranchOpInterface + RegionBranchOpInterface only):
  func.func @acceptable_ir_has_cleanable_loop_of_conditional_and_branch_op(%arg0: i1) {
    %0 = ub.poison : i32
    cf.br ^bb1(%0 : i32)
  ^bb1(%1: i32):  // 2 preds: ^bb0, ^bb1
    %2 = ub.poison : i32
    %3 = ub.poison : i32
    cf.cond_br %arg0, ^bb1(%2 : i32), ^bb2(%3 : i32)
  ^bb2(%4: i32):  // pred: ^bb1
    return
  }
```

Option 3 is a regression on that way. Are we ok with that, relying on a follow-up `--canonicalize` to clean it up?

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