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

[Matthias Springer]

================
@@ -196,9 +196,85 @@ static LogicalResult simplifyPassThroughBr(BranchOp op,
   return success();
 }
 
+/// If all incoming values for a block argument from all predecessors are the
+/// same value, replace uses of the block argument with that value. This allows
+/// the block argument to be removed by other canonicalization patterns.
+///
+/// Example:
+///   cf.br ^bb1(%poison : i32)      // pred 1
+///   cf.br ^bb1(%poison : i32)      // pred 2
+/// ^bb1(%arg0: i32):
+///   use(%arg0)
+/// ->
+/// ^bb1(%arg0: i32):
+///   use(%poison)                   // %arg0 now unused, folds away
+///
+static bool simplifyUniformBlockArgs(Block *dest, PatternRewriter &rewriter) {
+  bool changed = false;
+  for (BlockArgument arg : dest->getArguments()) {
+    if (arg.use_empty())
+      continue;
+
+    Value commonValue;
+    bool allSame = true;
+    for (Block *pred : dest->getPredecessors()) {
+      auto branch = dyn_cast<BranchOpInterface>(pred->getTerminator());
+      if (!branch) {
+        allSame = false;
+        break;
+      }
+
+      Value incoming;
+      for (unsigned i = 0; i < branch->getNumSuccessors(); ++i) {
+        if (branch->getSuccessor(i) != dest)
+          continue;
+        SuccessorOperands succOps = branch.getSuccessorOperands(i);
+        if (arg.getArgNumber() >= succOps.size()) {
+          allSame = false;
+          break;
+        }
+        incoming = succOps[arg.getArgNumber()];
+        break;
+      }
+      if (!incoming || (commonValue && commonValue != incoming)) {
+        allSame = false;
+        break;
+      }
+      commonValue = incoming;
+    }
+
+    if (allSame && commonValue && commonValue != arg) {
+      if (!commonValue.getDefiningOp<ub::PoisonOp>())
----------------
matthias-springer wrote:

There are possible ways:
1. Clone one of the existing ops.
2. Use the constant materializer. (`Dialect::materializeConstant`)

I believe the second option is safer. Cloning could be problematic because the op that defines the constant could have a side effect.

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