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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/183797

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

>From a2e7eeaf4d0deee56ea1b2e518dfc229ab327260 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Feb 2026 10:32:43 -0800
Subject: [PATCH] [mlir][cf] Fix crash in simplifyBrToBlockWithSinglePred when
 branch operand is a block argument of its successor

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
---
 .../Dialect/ControlFlow/IR/ControlFlowOps.cpp |  8 ++++
 .../Dialect/ControlFlow/canonicalize.mlir     | 39 +++++++++++++++++++
 2 files changed, 47 insertions(+)

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
+}