[Mlir-commits] [mlir] [mlir][opt][RemoveDeadValues] Fix crash when processing ops with regions that don't implement RegionBranchOpInterface (PR #182711)
Fedor Nikolaev
llvmlistbot at llvm.org
Mon Feb 23 08:28:20 PST 2026
https://github.com/felichita updated https://github.com/llvm/llvm-project/pull/182711
>From 4abe4e0ec0798c7e455ab8b8e10e546aef7445c9 Mon Sep 17 00:00:00 2001
From: Fedor Nikolaev <fridrixnm at gmail.com>
Date: Sun, 22 Feb 2026 00:01:54 +0100
Subject: [PATCH] [mlir][RemoveDeadValues] Fix crash when processing ops with
regions that don't implement RegionBranchOpInterface
The RemoveDeadValues pass was crashing with an assertion failure when
processing IR containing gpu.launch operations. The root cause was that
the pass and liveness analysis do not handle ops with regions that don't
implement RegionBranchOpInterface - the DataFlow solver doesn't propagate
liveness through such regions, causing block arguments inside them to be
incorrectly marked as dead.
Fixes #182263
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 65 ++++++++++++++++-
mlir/test/Transforms/remove-dead-values.mlir | 74 ++++++++++++++++++++
2 files changed, 136 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 12a47ba2fb65a..a6a7cf33757d0 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -511,8 +511,10 @@ static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
// Do (2)
BitVector successorNonLive =
markLives(operandValues, nonLiveSet, la).flip();
- collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),
- successorNonLive);
+ if (successorBlock->hasNoPredecessors() ||
+ llvm::hasSingleElement(successorBlock->getPredecessors()))
+ collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),
+ successorNonLive);
// Do (3)
cl.blocks.push_back({successorBlock, successorNonLive});
@@ -558,9 +560,15 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
TrackingListener listener;
IRRewriter rewriter(ctx, &listener);
+ auto hasMultiplePreds = [](Block *b) {
+ return !b->hasNoPredecessors() &&
+ !llvm::hasSingleElement(b->getPredecessors());
+ };
+
// 1. 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";
+ DenseSet<Block *> processedBlocks;
for (auto &b : list.blocks) {
// blocks that are accessed via multiple codepaths processed once
if (b.b->getNumArguments() != b.nonLiveArgs.size())
@@ -573,6 +581,26 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
<< OpWithFlags(b.b->getParent()->getParentOp(),
OpPrintingFlags().skipRegions().printGenericOpForm());
});
+ // Skip if already processed
+ if (processedBlocks.count(b.b))
+ continue;
+ // Skip entry blocks of functions - handled by processFuncOp
+ if (b.b->isEntryBlock() && isa<FunctionOpInterface>(b.b->getParentOp()))
+ continue;
+ // Only protect blocks with multiple predecessors
+ if (hasMultiplePreds(b.b)) {
+ // Check if any entry has this block with live args
+ bool hasLiveFromAnyPred = false;
+ for (auto &other : list.blocks) {
+ if (other.b == b.b && other.nonLiveArgs.none()) {
+ hasLiveFromAnyPred = true;
+ break;
+ }
+ }
+ if (hasLiveFromAnyPred)
+ continue;
+ }
+ processedBlocks.insert(b.b);
// 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) {
@@ -599,6 +627,21 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
<< OpWithFlags(op.branch.getOperation(),
OpPrintingFlags().skipRegions().printGenericOpForm());
});
+
+ // Only protect blocks with multiple predecessors
+ Block *succBlock = op.branch->getSuccessor(op.successorIndex);
+ if (hasMultiplePreds(succBlock)) {
+ bool otherLivePred = false;
+ for (auto &other : list.successorOperands) {
+ Block *otherSucc = other.branch->getSuccessor(other.successorIndex);
+ if (otherSucc == succBlock && other.nonLiveOperands.none()) {
+ otherLivePred = true;
+ break;
+ }
+ }
+ if (otherLivePred)
+ continue;
+ }
// it iterates backwards because erase invalidates all successor indexes
for (int i = successorOperands.size() - 1; i >= 0; --i) {
if (!op.nonLiveOperands[i])
@@ -772,7 +815,22 @@ void RemoveDeadValues::runOnOperation() {
// end of this pass.
RDVFinalCleanupList finalCleanupList;
- module->walk([&](Operation *op) {
+ module->walk([&](Operation *op) -> WalkResult {
+ Operation *parent = op->getParentOp();
+ if (parent && !isa<RegionBranchOpInterface>(parent) &&
+ !isa<FunctionOpInterface>(parent) && !isa<ModuleOp>(parent) &&
+ parent->getNumRegions() > 0) {
+ // Op is inside a region not analyzed by the dataflow solver (e.g.
+ // gpu.launch). We cannot use liveness analysis here, so conservatively
+ // erase only ops that are guaranteed dead: no memory effects, no uses of
+ // results, not a terminator, and not a region-containing op.
+ if (!op->hasTrait<::mlir::OpTrait::IsTerminator>() &&
+ !op->getNumRegions() && !isa<BranchOpInterface>(op) &&
+ isMemoryEffectFree(op) && op->use_empty()) {
+ finalCleanupList.operations.push_back(op);
+ }
+ return WalkResult::advance();
+ }
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
processFuncOp(funcOp, module, la, deadVals, finalCleanupList);
} else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
@@ -788,6 +846,7 @@ void RemoveDeadValues::runOnOperation() {
} else {
processSimpleOp(op, la, deadVals, finalCleanupList);
}
+ return WalkResult::advance();
});
MLIRContext *context = module->getContext();
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 87e77b2eb700f..c26078ba640e0 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -476,6 +476,80 @@ func.func @kernel(%arg0: memref<18xf32>) {
// -----
+// CHECK-LABEL: func.func @gpu_launch_mixed_dead_and_live
+// CHECK: gpu.block_id x
+// CHECK-NOT: gpu.thread_id x
+// CHECK: memref.store
+func.func @gpu_launch_mixed_dead_and_live(%buf: memref<8xf32>, %val: f32) {
+ %c1 = arith.constant 1 : index
+ gpu.launch
+ blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%bsx = %c1, %bsy = %c1, %bsz = %c1) {
+ %blk_x = gpu.block_id x // live
+ %thr_x = gpu.thread_id x // dead
+ memref.store %val, %buf[%blk_x] : memref<8xf32>
+ gpu.terminator
+ }
+ return
+}
+
+// -----
+
+// Value used inside gpu.launch feeds into public function return - must not be removed
+// CHECK-LABEL: func.func @gpu_launch_value_used_for_public_return
+// CHECK: gpu.thread_id x
+// CHECK: memref.store
+func.func @gpu_launch_value_used_for_public_return(%buf: memref<1xindex>) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ gpu.launch
+ blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%bsx = %c1, %bsy = %c1, %bsz = %c1) {
+ %thr_x = gpu.thread_id x
+ memref.store %thr_x, %buf[%c0] : memref<1xindex>
+ gpu.terminator
+ }
+ %val = memref.load %buf[%c0] : memref<1xindex>
+ return %val : index
+}
+
+// -----
+
+// CHECK-LABEL: func.func @gpu_launch_thread_id_no_use
+// CHECK-NOT: gpu.thread_id
+// CHECK: return
+func.func @gpu_launch_thread_id_no_use() {
+ %c1 = arith.constant 1 : index
+ gpu.launch
+ blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%bsx = %c1, %bsy = %c1, %bsz = %c1) {
+ %thr_x = gpu.thread_id x
+ gpu.terminator
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @gpu_launch_inside_scf_for
+// CHECK: gpu.block_id x
+// CHECK: memref.store
+func.func @gpu_launch_inside_scf_for(%buf: memref<8xf32>, %val: f32, %lb: index, %ub: index, %step: index) {
+ scf.for %iv = %lb to %ub step %step {
+ %c1 = arith.constant 1 : index
+ gpu.launch
+ blocks(%bx, %by, %bz) in (%gx = %c1, %gy = %c1, %gz = %c1)
+ threads(%tx, %ty, %tz) in (%bsx = %c1, %bsy = %c1, %bsz = %c1) {
+ %blk_x = gpu.block_id x
+ memref.store %val, %buf[%blk_x] : memref<8xf32>
+ gpu.terminator
+ }
+ }
+ return
+}
+
+// -----
+
// CHECK-LABEL: llvm_unreachable
// CHECK-LABEL: @fn_with_llvm_unreachable
More information about the Mlir-commits
mailing list