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

[Matthias Springer]

================
@@ -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)))) {
----------------
matthias-springer wrote:

OK, then what about running canonicalization patterns on **all** `RegionBranchOpInterface` ops and `BranchOpInterface` ops (kind of what this PR is doing in its current state)? That's already much more lightweight.

There's one problem with narrowing it down even further to only branch ops that are also somehow related to dropped uses of block arguments. By applying canonicalization, we get a few extra simplifications such as [folding pass-through branches](https://github.com/llvm/llvm-project/blob/33be2d0e7ab2c8874f53df8cbbda07381020b1c4/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp#L190). It would be somewhat odd if these simplifications are applied only to basic block from which this pass happened to drop block arguments. (Makes it hard to explain what exactly this pass is doing.)

> Note: using ub ops instead of performing direct simplifications in the IR is already a bit unfortunate from this perspective.

Would be nice, but not practical IMO. How many bugs did we have in this pass over the last year? I have given up on that approach, it's too complicated. I'd rather pay the extra compilation cost and have it correct in return...


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