[Mlir-commits] [mlir] [mlir] Preserve live branch arguments (PR #208742)

[llvmlistbot at llvm.org]

https://github.com/mygitljf created https://github.com/llvm/llvm-project/pull/208742

I found a case where deadness from one incoming edge was being treated as if the shared successor argument itself was dead. This change keeps that edge-local information from overriding the liveness result, while still allowing dead branch inputs to be cleaned up.
Fixes #206916 

>From 0be1d90d51a20bb0d195a117d731feea625e150a Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 10 Jul 2026 22:54:54 +0000
Subject: [PATCH] [mlir] Preserve live branch arguments

---
 mlir/lib/Transforms/RemoveDeadValues.cpp     |  5 +++
 mlir/test/Transforms/remove-dead-values.mlir | 34 ++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index f0a210a2ededb..9521db2d1aa57 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -511,6 +511,11 @@ static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
     // Do (2)
     BitVector successorNonLive =
         markLives(operandValues, nonLiveSet, la).flip();
+    // A dead incoming edge does not make the successor argument globally dead.
+    for (auto [argIndex, arg] : llvm::enumerate(successorBlock->getArguments()))
+      if (successorNonLive[argIndex] && hasLive(arg, nonLiveSet, la))
+        successorNonLive.reset(argIndex);
+
     collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),
                          successorNonLive);
 
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 64088ce15cd48..437509875df36 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -850,6 +850,40 @@ func.func @replace_dead_operation_results_with_poison(%0: vector<1xindex>) -> ve
 
 // -----
 
+// CHECK-LABEL: func.func @dont_mark_live_successor_arg_dead_from_unreachable_pred
+// CHECK-CANONICALIZE-LABEL: func.func @dont_mark_live_successor_arg_dead_from_unreachable_pred
+func.func @dont_mark_live_successor_arg_dead_from_unreachable_pred(%arg0: memref<?x8x16xf32>, %arg1: memref<4x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %cst = arith.constant 1.000000e+00 : f32
+  cf.br ^bb1(%c0 : index)
+^bb1(%0: index):
+  %1 = arith.cmpi slt, %0, %c4 : index
+  cf.cond_br %1, ^bb2, ^bb7
+^bb2:
+  cf.br ^bb3(%c0 : index)
+^bb3(%2: index):
+  %3 = arith.cmpi slt, %2, %c4 : index
+  cf.cond_br %3, ^bb4, ^bb6
+^bb4:
+  %4 = memref.load %arg0[%0, %2, %c1] : memref<?x8x16xf32>
+  %5 = arith.addf %4, %cst : f32
+  memref.store %5, %arg1[%0, %2] : memref<4x?xf32>
+  %6 = arith.addi %2, %c1 : index
+  cf.br ^bb3(%6 : index)
+^bb5:
+  %7 = arith.addi %0, %c1 : index
+  cf.br ^bb1(%7 : index)
+^bb6:
+  %8 = arith.addi %0, %c1 : index
+  cf.br ^bb1(%8 : index)
+^bb7:
+  return
+}
+
+// -----
+
 // Verify that a referenced by a non-call op (spirv.EntryPoint),
 // while still having another usual call site is preserved as-is
 // since the pass cannot analyse non-call users.