[Mlir-commits] [mlir] [mlir] Coordinate branch argument cleanup (PR #208940)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jul 11 12:37:07 PDT 2026
https://github.com/mygitljf created https://github.com/llvm/llvm-project/pull/208940
I made branch argument cleanup consistent across all incoming control-flow paths, so unreachable paths no longer leave malformed IR. The change also preserves internally produced arguments and keeps surviving uses valid with poison values.
Fixes #205984
>From b22ea028f463f0da3fdf71aa211bca11f6cf4f07 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Sun, 12 Jul 2026 03:36:50 +0000
Subject: [PATCH] [mlir] Coordinate branch argument cleanup
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 80 ++++++++++----------
mlir/test/Transforms/remove-dead-values.mlir | 63 +++++++++++++++
2 files changed, 101 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 6e55bc390be23..6c0f2f8b7bf1e 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -96,15 +96,10 @@ struct OperandsToCleanup {
bool replaceWithPoison = false;
};
-struct BlockArgsToCleanup {
- Block *b;
- BitVector nonLiveArgs;
-};
-
struct SuccessorOperandsToCleanup {
BranchOpInterface branch;
unsigned successorIndex;
- BitVector nonLiveOperands;
+ Block *successorBlock;
};
struct RDVFinalCleanupList {
@@ -112,7 +107,7 @@ struct RDVFinalCleanupList {
SmallVector<FunctionToCleanUp> functions;
SmallVector<OperandsToCleanup> operands;
SmallVector<ResultsToCleanup> results;
- SmallVector<BlockArgsToCleanup> blocks;
+ DenseMap<Block *, BitVector> branchBlockArgs;
SmallVector<SuccessorOperandsToCleanup> successorOperands;
};
@@ -465,11 +460,10 @@ 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
-/// as well as their corresponding arguments.
+/// (1) Fetch the successor block arguments' liveness analysis data.
+/// (2) Exclude arguments produced internally by the branch operation.
+/// (3) Collect dead forwarded operands and their corresponding block arguments
+/// for future removal.
static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
DenseSet<Value> &nonLiveSet,
@@ -494,25 +488,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)
BitVector successorNonLive =
- markLives(operandValues, nonLiveSet, la).flip();
- collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),
- successorNonLive);
-
- // Do (3)
- cl.blocks.push_back({successorBlock, successorNonLive});
- cl.successorOperands.push_back({branchOp, succIdx, successorNonLive});
+ markLives(successorBlock->getArguments(), nonLiveSet, la).flip();
+ successorNonLive.reset(0, successorOperands.getProducedOperandCount());
+ auto [it, inserted] =
+ cl.branchBlockArgs.try_emplace(successorBlock, successorNonLive);
+ if (!inserted)
+ it->second &= successorNonLive;
+
+ cl.successorOperands.push_back({branchOp, succIdx, successorBlock});
}
}
@@ -576,26 +563,33 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
// 2. Blocks, We must remove the block arguments and successor operands before
// deleting the operation, as they may reside in the region operation.
- LDBG() << "Cleaning up " << list.blocks.size() << " block argument lists";
- for (auto &b : list.blocks) {
- // blocks that are accessed via multiple codepaths processed once
- if (b.b->getNumArguments() != b.nonLiveArgs.size())
+ LDBG() << "Cleaning up " << list.branchBlockArgs.size()
+ << " block argument lists";
+ for (auto &entry : list.branchBlockArgs) {
+ Block *block = entry.first;
+ BitVector &nonLiveArgs = entry.second;
+ if (block->getNumArguments() != nonLiveArgs.size())
continue;
LDBG_OS([&](raw_ostream &os) {
os << "Erasing non-live arguments [";
- llvm::interleaveComma(b.nonLiveArgs.set_bits(), os);
- os << "] from block #" << b.b->computeBlockNumber() << " in region #"
- << b.b->getParent()->getRegionNumber() << " of operation "
- << OpWithFlags(b.b->getParent()->getParentOp(),
+ llvm::interleaveComma(nonLiveArgs.set_bits(), os);
+ os << "] from block #" << block->computeBlockNumber() << " in region #"
+ << block->getParent()->getRegionNumber() << " of operation "
+ << OpWithFlags(block->getParent()->getParentOp(),
OpPrintingFlags().skipRegions().printGenericOpForm());
});
// Note: Iterate from the end to make sure that that indices of not yet
// processes arguments do not change.
- for (int i = b.nonLiveArgs.size() - 1; i >= 0; --i) {
- if (!b.nonLiveArgs[i])
+ for (int i = nonLiveArgs.size() - 1; i >= 0; --i) {
+ if (!nonLiveArgs[i])
continue;
- b.b->getArgument(i).dropAllUses();
- b.b->eraseArgument(i);
+ BlockArgument argument = block->getArgument(i);
+ if (!argument.use_empty()) {
+ rewriter.setInsertionPointToStart(block);
+ Value poison = createPoisonedValue(rewriter, argument);
+ rewriter.replaceAllUsesWith(argument, poison);
+ }
+ block->eraseArgument(i);
}
}
@@ -605,19 +599,21 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
for (auto &op : list.successorOperands) {
SuccessorOperands successorOperands =
op.branch.getSuccessorOperands(op.successorIndex);
+ const BitVector &nonLiveOperands =
+ list.branchBlockArgs.find(op.successorBlock)->second;
// blocks that are accessed via multiple codepaths processed once
- if (successorOperands.size() != op.nonLiveOperands.size())
+ if (successorOperands.size() != nonLiveOperands.size())
continue;
LDBG_OS([&](raw_ostream &os) {
os << "Erasing non-live successor operands [";
- llvm::interleaveComma(op.nonLiveOperands.set_bits(), os);
+ llvm::interleaveComma(nonLiveOperands.set_bits(), os);
os << "] from successor " << op.successorIndex << " of branch: "
<< OpWithFlags(op.branch.getOperation(),
OpPrintingFlags().skipRegions().printGenericOpForm());
});
// it iterates backwards because erase invalidates all successor indexes
for (int i = successorOperands.size() - 1; i >= 0; --i) {
- if (!op.nonLiveOperands[i])
+ if (!nonLiveOperands[i])
continue;
successorOperands.erase(i);
}
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 390a448060b7f..ed835d2d5fd91 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -895,3 +895,66 @@ module @func_with_non_call_users {
}
spirv.EntryPoint "GLCompute" @callee
}
+
+// -----
+
+// Regression test for #205984. The loop header arguments are live on the entry
+// edge, even though the constant comparison makes the backedge unreachable.
+// CHECK-LABEL: func.func @keep_live_loop_header_args_with_unreachable_backedge
+// CHECK: cf.br ^[[HEADER:bb[0-9]+]](%[[ZERO:.*]], %[[LIMIT:.*]] : index, index)
+// CHECK: ^[[HEADER]](%[[IV:.*]]: index, %[[BOUND:.*]]: index):
+// CHECK-NEXT: %[[CMP:.*]] = arith.cmpi slt, %[[IV]], %[[BOUND]] : index
+// CHECK-NEXT: cf.cond_br %[[CMP]], ^[[BACKEDGE:bb[0-9]+]], ^[[EXIT:bb[0-9]+]]
+// CHECK: ^[[BACKEDGE]]:
+// CHECK-NEXT: %[[NEXT:.*]] = ub.poison : index
+// CHECK-NEXT: %[[NEXT_BOUND:.*]] = ub.poison : index
+// CHECK-NEXT: cf.br ^[[HEADER]](%[[NEXT]], %[[NEXT_BOUND]] : index, index)
+// CHECK: ^[[EXIT]]:
+
+// CHECK-CANONICALIZE-LABEL: func.func @keep_live_loop_header_args_with_unreachable_backedge
+// CHECK-CANONICALIZE: cf.br ^[[HEADER:bb[0-9]+]](%[[ZERO:.*]], %[[LIMIT:.*]] : index, index)
+// CHECK-CANONICALIZE: ^[[HEADER]](%[[IV:.*]]: index, %[[BOUND:.*]]: index):
+// CHECK-CANONICALIZE-NEXT: %[[CMP:.*]] = arith.cmpi slt, %[[IV]], %[[BOUND]] : index
+// CHECK-CANONICALIZE-NEXT: cf.cond_br %[[CMP]], ^[[BACKEDGE:bb[0-9]+]], ^[[EXIT:bb[0-9]+]]
+// CHECK-CANONICALIZE: ^[[BACKEDGE]]:
+// CHECK-CANONICALIZE-NEXT: %[[NEXT:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT: %[[NEXT_BOUND:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT: cf.br ^[[HEADER]](%[[NEXT]], %[[NEXT_BOUND]] : index, index)
+// CHECK-CANONICALIZE: ^[[EXIT]]:
+func.func @keep_live_loop_header_args_with_unreachable_backedge() {
+ %zero = arith.constant 0 : index
+ %limit = arith.constant 0 : index
+ cf.br ^bb1(%zero, %limit : index, index)
+^bb1(%iv: index, %bound: index):
+ %cmp = arith.cmpi slt, %iv, %bound : index
+ cf.cond_br %cmp, ^bb2(%limit, %limit : index, index), ^bb3
+^bb2(%next: index, %next_bound: index):
+ cf.br ^bb1(%next, %next_bound : index, index)
+^bb3:
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @keep_produced_successor_argument
+// CHECK: "test.internal_br"()[^[[NORMAL:.*]], ^[[ERROR:.*]]]
+// CHECK: ^[[NORMAL]]:
+// CHECK-NEXT: "test.terminator"() : () -> ()
+// CHECK: ^[[ERROR]](%{{.*}}: i32):
+// CHECK-NEXT: "test.terminator"() : () -> ()
+
+// CHECK-CANONICALIZE-LABEL: func.func @keep_produced_successor_argument
+// CHECK-CANONICALIZE: "test.internal_br"()[^[[NORMAL:.*]], ^[[ERROR:.*]]]
+// CHECK-CANONICALIZE: ^[[NORMAL]]:
+// CHECK-CANONICALIZE-NEXT: "test.terminator"() : () -> ()
+// CHECK-CANONICALIZE: ^[[ERROR]](%{{.*}}: i32):
+// CHECK-CANONICALIZE-NEXT: "test.terminator"() : () -> ()
+func.func @keep_produced_successor_argument() {
+ "test.internal_br"() [^bb0, ^bb1] {
+ operandSegmentSizes = array<i32: 0, 0>
+ } : () -> ()
+^bb0:
+ "test.terminator"() : () -> ()
+^bb1(%produced: i32):
+ "test.terminator"() : () -> ()
+}
More information about the Mlir-commits
mailing list