[Mlir-commits] [mlir] [mlir][cf] Fix crash in simplifyBrToBlockWithSinglePred when branch operand is a block argument of its successor (PR #183797)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-cf

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When `simplifyBrToBlockWithSinglePred` merges a block into its sole predecessor, it calls `inlineBlockBefore` which replaces each block argument with the corresponding value passed by the branch. If one of those values is itself a block argument of the successor block, the call `replaceAllUsesWith(arg, arg)` is a no-op. Any uses of that argument outside the block (e.g. in a downstream block) are therefore not replaced, and when the successor block is erased the argument is destroyed while those uses are still live, triggering the assertion `use_empty() && "Cannot destroy a value that still has uses\!"` in `IRObjectWithUseList::~IRObjectWithUseList`.

Guard against this by returning early when any branch operand is a block argument owned by the destination block.

Fixes #<!-- -->126213

---
Full diff: https://github.com/llvm/llvm-project/pull/183797.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp (+8) 
- (modified) mlir/test/Dialect/ControlFlow/canonicalize.mlir (+39) 


``````````diff
diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index d2078d8ab5ca5..0ce0d55f4397c 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -166,6 +166,14 @@ 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.
+  for (Value operand : op.getOperands())
+    if (auto ba = dyn_cast<BlockArgument>(operand))
+      if (ba.getOwner() == succ)
+        return failure();
+
   // Merge the successor into the current block and erase the branch.
   SmallVector<Value> brOperands(op.getOperands());
   rewriter.eraseOp(op);
diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
index 21a16784b81b2..3bf4ca006f34c 100644
--- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir
+++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
@@ -656,3 +656,42 @@ func.func @drop_unreachable_branch_2(%c: i1) {
 ^bb2:
   ub.unreachable
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/126213:
+// simplifyBrToBlockWithSinglePred must not merge a block into its sole
+// predecessor when a branch operand is itself a block argument of the
+// destination — replaceAllUsesWith would be a no-op for that argument,
+// leaving the block argument live after the block is erased and crashing.
+
+// CHECK-LABEL: @no_merge_self_arg_loop
+func.func @no_merge_self_arg_loop() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %true = arith.constant true
+  cf.cond_br %true, ^bb1, ^bb2
+^bb1:
+  cf.br ^bb3(%true : i1)
+^bb2:
+  cf.br ^bb3(%true : i1)
+^bb3(%0: i1):
+  cf.br ^bb4
+^bb4:
+  cf.br ^bb5
+^bb5:
+  cf.cond_br %0, ^bb6, ^bb7
+^bb6:
+  cf.br ^bb5
+^bb7:
+  cf.br ^bb8(%c0, %true : index, i1)
+^bb8(%1: index, %2: i1):
+  %3 = arith.cmpi slt, %1, %c1 : index
+  cf.cond_br %3, ^bb9, ^bb10
+^bb9:
+  %4 = arith.addi %1, %c1 : index
+  cf.br ^bb8(%4, %2 : index, i1)
+^bb10:
+  vector.print %2 : i1
+  return
+}

``````````

</details>


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