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

If I do smth like this:
```
module->walk([&](Operation *op) {
  if (isa<RegionBranchOpInterface, BranchOpInterface>(op))
    opsToCanonicalize.push_back(op);
});
```
then block args will not be removed, only uses will be replaced. I guess it's not right, so I choose module instead specific ops.

Another way, if I add successor block implicitly inside the loop, smth like this:
```
  for (Block *succ : branch->getSuccessors())
    if (Operation *term = succ->getTerminator())
      opsToCanonicalize.push_back(term);
```
example mlir where it doesn't work safely without additional bfs
```
func.func @ex(%cond: i1, %a: i32) -> i32 {
  cf.cond_br %cond, ^bb1, ^bb2
^bb1:
  cf.br ^bb3(%a : i32)       // branchOp - in the list
^bb2:
  cf.br ^bb3(%a : i32)       // in the list as well
^bb3(%arg0: i32):            // successor - terminator 
  cf.br ^bb4(%arg0 : i32) // will be added as well like a terminator 
^bb4(%arg1: i32):            // but ^bb4 is not ...
  return %arg1 : i32
}
```

The module traversal seemed like prose at first glance, although it was necessary to change the approach. Apparently, there is better way to do it

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