[Mlir-commits] [mlir] [mlir][cf] Do not merge unique-pred block when branch operand derived from successor argument (PR #225791)

[Scott Manley]

https://github.com/rscottmanley created https://github.com/llvm/llvm-project/pull/225791



simplifyBrToBlockWithSinglePred already refused to merge when a branch operand was itself a destination block argument. That missed the loop latch case where the operand is computed from the argument, e.g. `%next = addi %iter, 1` / `br ^header(%next)`. Merging rewrote the add into a cyclic def (`addi %next, 1`), after which AddIAddConstant unbounded-doubled the step.

Walk the operand's defining chain and refuse the merge if it reaches a successor argument. Adds a canonicalize regression test.

>From 96b5fa4a3b81abc7601741cd60e51d0ea1774762 Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Wed, 23 Sep 2026 07:08:56 -0700
Subject: [PATCH] [mlir][cf] Do not merge a unique-pred block when the branch
 operand is derived from a successor argument

simplifyBrToBlockWithSinglePred already refused to merge when a branch
operand was itself a destination block argument. That missed the loop
latch case where the operand is computed from the argument, e.g.
`%next = addi %iter, 1` / `br ^header(%next)`. Merging rewrote the add
into a cyclic def (`addi %next, 1`), after which AddIAddConstant
unbounded-doubled the step.

Walk the operand's defining chain and refuse the merge if it reaches a
successor argument. Adds a canonicalize regression test.
---
 .../Dialect/ControlFlow/IR/ControlFlowOps.cpp | 37 ++++++++++++++++---
 .../Dialect/ControlFlow/canonicalize.mlir     | 27 ++++++++++++++
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 9016d1929f859c..171e41a27abd37 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -25,6 +25,7 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include <numeric>
 
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOpsDialect.cpp.inc"
@@ -156,6 +157,28 @@ static LogicalResult collapseBranch(Block *&successor,
   return success();
 }
 
+/// Return true if `value` is a block argument of `block`, or is defined by an
+/// operation that transitively uses such an argument.
+static bool isTransitivelyDefinedFrom(Value value, Block *block) {
+  SmallVector<Value, 8> worklist{value};
+  SmallPtrSet<Value, 16> visited;
+  while (!worklist.empty()) {
+    Value current = worklist.pop_back_val();
+    if (!visited.insert(current).second)
+      continue;
+    if (auto arg = dyn_cast<BlockArgument>(current)) {
+      if (arg.getOwner() == block)
+        return true;
+      continue;
+    }
+    Operation *def = current.getDefiningOp();
+    if (!def)
+      continue;
+    llvm::append_range(worklist, def->getOperands());
+  }
+  return false;
+}
+
 /// Simplify a branch to a block that has a single predecessor. This effectively
 /// merges the two blocks.
 static LogicalResult
@@ -166,13 +189,15 @@ simplifyBrToBlockWithSinglePred(BranchOp op, PatternRewriter &rewriter) {
   if (succ == opParent || !llvm::hasSingleElement(succ->getPredecessors()))
     return failure();
 
-  // If any branch operand is itself a block argument of the successor, merging
-  // would call replaceAllUsesWith(arg, arg) — a no-op — leaving dangling uses
-  // of that argument after the successor block is erased.
+  // If any branch operand is defined from a block argument of the successor,
+  // merging replaces that argument with the operand. A direct self-argument
+  // makes replaceAllUsesWith a no-op and leaves dangling uses after the
+  // successor is erased. A value derived from the argument (for example
+  // `addi %arg, 1` on a backedge) becomes a cyclic def after the replacement,
+  // and later arith folds do not terminate.
   for (Value operand : op.getOperands())
-    if (auto ba = dyn_cast<BlockArgument>(operand))
-      if (ba.getOwner() == succ)
-        return failure();
+    if (isTransitivelyDefinedFrom(operand, succ))
+      return failure();
 
   // Merge the successor into the current block and erase the branch.
   SmallVector<Value> brOperands(op.getOperands());
diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
index c8cf5931a1297c..731b4bb4990f07 100644
--- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir
+++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
@@ -687,6 +687,33 @@ func.func @no_merge_self_arg_loop(%step: i1) -> i1 {
   return %result : i1
 }
 
+// -----
+
+// Related to the case above: after folding the constant guard, ^body is the
+// unique predecessor of ^header. The backedge operand is `addi %iter, 1`,
+// which is derived from the destination block argument. Merging would rewrite
+// that add to `addi %next, 1` (a cyclic def) and arith canonicalization would
+// then keep doubling the step.
+
+// CHECK-LABEL: @no_merge_derived_from_dest_arg
+//  CHECK-NEXT:   %[[C0:.*]] = arith.constant 0 : i32
+//  CHECK-NEXT:   return %[[C0]]
+func.func @no_merge_derived_from_dest_arg(%loop: i1, %initial: i32) -> i32 {
+  %true = arith.constant true
+  %c1 = arith.constant 1 : i32
+  %c0 = arith.constant 0 : i32
+  cf.br ^guard(%true : i1)
+^guard(%enter: i1):
+  cf.cond_br %enter, ^exit(%c0 : i32), ^header(%initial : i32)
+^header(%iter: i32):
+  cf.cond_br %loop, ^body, ^exit(%iter : i32)
+^body:
+  %next = arith.addi %iter, %c1 : i32
+  cf.br ^header(%next : i32)
+^exit(%result: i32):
+  return %result : i32
+}
+
 // Verify that block arguments are replaced with a uniform incoming value
 // when all predecessors pass the same SSA value