[Mlir-commits] [mlir] [MLIR] Fix remove-dead-values missing dead block arguments (PR #216652)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Aayush Shrivastava (iamaayushrivastava)

<details>
<summary>Changes</summary>

Fixes #<!-- -->216644

In `processBranchOp` (`mlir/lib/Transforms/RemoveDeadValues.cpp`), the liveness of a `cf.br` successor's block argument was computed from the liveness of the caller-side operand value being forwarded to it, not from the block argument itself. When the forwarded value happened to be live for an unrelated reason (here, `%arg0` is both passed to `cf.br` and directly returned), the block argument incorrectly inherited that liveness and was never recognized as dead. The fix computes `successorNonLive` from `successorBlock->getArguments()` directly, which correctly reflects whether the argument is used inside its own block.

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


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+11-16) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 6e55bc390be23..3ea0f5f642365 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -465,10 +465,9 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
 /// conditional branch op), the entire operation is dead.
 ///
 /// Otherwise, iterate through each successor block of `branchOp`.
-/// (1) For each successor block, gather all operands from all successors.
-/// (2) Fetch their associated liveness analysis data and collect for future
-///     removal.
-/// (3) Identify and collect the dead operands from the successor block
+/// (1) For each successor block, fetch the liveness analysis data of its
+///     block arguments and collect for future removal.
+/// (2) Identify and collect the dead operands from the successor block
 ///     as well as their corresponding arguments.
 
 static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
@@ -495,22 +494,18 @@ static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
   for (unsigned succIdx = 0; succIdx < numSuccessors; ++succIdx) {
     Block *successorBlock = branchOp->getSuccessor(succIdx);
 
-    // Do (1)
-    SuccessorOperands successorOperands =
-        branchOp.getSuccessorOperands(succIdx);
-    SmallVector<Value> operandValues;
-    for (unsigned operandIdx = 0; operandIdx < successorOperands.size();
-         ++operandIdx) {
-      operandValues.push_back(successorOperands[operandIdx]);
-    }
-
-    // Do (2)
+    // Do (1). Liveness of a successor block argument must be determined
+    // from the argument itself (i.e. whether it is used inside the
+    // successor block), not from the liveness of the operand value being
+    // forwarded to it. The forwarded operand value may well be live for
+    // unrelated reasons (e.g. it is also used elsewhere), which must not
+    // prevent the block argument from being recognized as dead.
     BitVector successorNonLive =
-        markLives(operandValues, nonLiveSet, la).flip();
+        markLives(successorBlock->getArguments(), nonLiveSet, la).flip();
     collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),
                          successorNonLive);
 
-    // Do (3)
+    // Do (2)
     cl.blocks.push_back({successorBlock, successorNonLive});
     cl.successorOperands.push_back({branchOp, succIdx, successorNonLive});
   }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 390a448060b7f..bca692699225c 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -895,3 +895,21 @@ module @func_with_non_call_users {
   }
   spirv.EntryPoint "GLCompute" @callee
 }
+
+// -----
+
+// A dead successor block argument must be removed even when the value
+// forwarded to it is itself live for an unrelated reason (here, %arg0 is
+// also returned directly). Fixes #216644.
+//
+// CHECK-LABEL: func.func @dead_block_arg_with_live_forwarded_operand
+// CHECK-SAME: (%[[arg0:.*]]: i32) -> i32
+// CHECK: cf.br ^[[BB1:bb[0-9]+]]
+// CHECK-NOT: i32
+// CHECK: ^[[BB1]]:
+// CHECK-NEXT: return %[[arg0]] : i32
+func.func @dead_block_arg_with_live_forwarded_operand(%arg0: i32) -> i32 {
+  cf.br ^bb1(%arg0 : i32)
+^bb1(%0: i32):
+  return %arg0 : i32
+}

``````````

</details>


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