[Mlir-commits] [mlir] [MLIR] Fix crash in RemoveDeadValues when branch is in dead block (PR #186455)

[Mehdi Amini]

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

When a branch operation is in a statically unreachable (dead) block, its forwarded operands are also dead. processBranchOp was propagating this deadness to the successor block arguments unconditionally, even when those block arguments are live because they are also forwarded from other reachable predecessor branches.

This caused an assertion failure in processSimpleOp when an op had a dead operand (the incorrectly-marked-dead block argument) but a live result.

Fix by checking, for each successor block argument that would be marked dead, whether the liveness analysis considers it live. If so, do not add it to the non-live set.

Fixes #182263

Assisted-by: Claude Code

>From c90736865ba23c5825ea09a5ed39899be729e647 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 6 Mar 2026 11:25:17 -0800
Subject: [PATCH] [MLIR] Fix crash in RemoveDeadValues when branch is in dead
 block

When a branch operation is in a statically unreachable (dead) block,
its forwarded operands are also dead. processBranchOp was propagating
this deadness to the successor block arguments unconditionally, even
when those block arguments are live because they are also forwarded
from other reachable predecessor branches.

This caused an assertion failure in processSimpleOp when an op had
a dead operand (the incorrectly-marked-dead block argument) but a
live result.

Fix by checking, for each successor block argument that would be
marked dead, whether the liveness analysis considers it live. If so,
do not add it to the non-live set.

Fixes #182263

Assisted-by: Claude Code
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 10 +++++++
 mlir/test/Transforms/remove-dead-values.mlir | 30 ++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 957621c16bf2b..2610fa3d1cffe 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -517,6 +517,16 @@ static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
     // Do (2)
     BitVector successorNonLive =
         markLives(operandValues, nonLiveSet, la).flip();
+    // A block argument should not be considered dead if the liveness analysis
+    // determines it is live. This can happen when a branch is in a statically
+    // unreachable (dead) block: the forwarded operand appears dead because it
+    // is in the dead block, but the successor block argument may still be live
+    // because it is also forwarded from other live predecessor branches.
+    for (auto [index, blockArg] :
+         llvm::enumerate(successorBlock->getArguments())) {
+      if (successorNonLive[index] && hasLive({blockArg}, nonLiveSet, la))
+        successorNonLive.reset(index);
+    }
     collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),
                          successorNonLive);
 
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 19bc6b2fddd66..343f7616f73fb 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -847,3 +847,33 @@ module @func_with_non_call_users {
   }
   spirv.EntryPoint "GLCompute" @callee
 }
+
+
+// -----
+
+// Verify that RemoveDeadValues does not crash when a branch in a statically
+// dead block forwards a dead value to a block argument that is live (because
+// another, reachable predecessor also forwards a value to it).
+// The constant comparison %cmp is always true (2 < 5), so ^bb2 is unreachable.
+// The block argument %arg0 in ^bb3 must not be incorrectly marked dead.
+// CHECK-LABEL: func.func @branch_in_dead_block_live_successor_arg
+func.func @branch_in_dead_block_live_successor_arg() -> i64 {
+  %c0 = arith.constant 0 : i64
+  %c1 = arith.constant 1 : i64
+  %sum = arith.addi %c0, %c1 : i64
+  %c2 = arith.constant 2 : i64
+  %mul = arith.muli %sum, %c2 : i64
+  %c5 = arith.constant 5 : i64
+  %cmp = arith.cmpi slt, %mul, %c5 : i64
+  cf.cond_br %cmp, ^bb1, ^bb2
+^bb1:
+  %c10 = arith.constant 10 : i64
+  cf.br ^bb3(%c10 : i64)
+^bb2:
+  %c20 = arith.constant 20 : i64
+  cf.br ^bb3(%c20 : i64)
+^bb3(%arg0: i64):
+  // CHECK: arith.addi
+  %result = arith.addi %arg0, %sum : i64
+  func.return %result : i64
+}