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

[Mehdi Amini]

================
@@ -802,18 +737,42 @@ void RemoveDeadValues::runOnOperation() {
   if (!canonicalize)
     return;
 
-  // Canonicalize all region branch ops.
-  SmallVector<Operation *> opsToCanonicalize;
-  module->walk([&](RegionBranchOpInterface regionBranchOp) {
-    opsToCanonicalize.push_back(regionBranchOp.getOperation());
+  // Collect ops to canonicalize via BFS over successor blocks.
+  // Seed: all branch ops, region branch ops, and return-like ops.
+  // Then transitively follow successor block terminators to cover
+  // reachable blocks. Note: block arguments are not removed here;
+  // a follow-up --canonicalize pass is needed for full cleanup.
+  SmallVector<Operation *> worklist;
+  DenseSet<Operation *> visited;
+
+  module->walk([&](Operation *op) {
+    if (!isa<RegionBranchOpInterface, BranchOpInterface>(op) &&
+        !op->hasTrait<OpTrait::ReturnLike>())
+      return;
+    if (visited.insert(op).second)
+      worklist.push_back(op);
   });
-  // Collect all canonicalization patterns for region branch ops.
+
+  // BFS: follow successor block terminators transitively.
+  SmallVector<Operation *> opsToCanonicalize;
+  while (!worklist.empty()) {
+    Operation *op = worklist.pop_back_val();
+    opsToCanonicalize.push_back(op);
+    for (Block *succ : op->getSuccessors()) {
+      Operation *term = succ->getTerminator();
+      if (term && visited.insert(term).second)
+        worklist.push_back(term);
+    }
+  }
----------------
joker-eph wrote:

Can you extract this logic into a helper function taking the `opsToCanonicalize` to populate as argument? This should also be documented better into what it collects: right now you wrote "Collect ops to canonicalize via BFS over successor blocks" but we don't know which ops you're collecting (seems like it's all reachable terminators? Could we save the BFS and just collect all terminators in the region regardless of reachability?)

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