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

[Fedor Nikolaev]

================
@@ -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>())
----------------
felichita wrote:

@matthias-springer For the constant attribute equality follow-up: when all incoming values match the same constant attribute, what is the best way to materialize the replacement value  clone one of the existing constant ops from the IR or use rewriter.createOrFold?

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