[Mlir-commits] [mlir] [mlir] Fix RDV poison cleanup (PR #206270)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jun 27 09:50:25 PDT 2026
https://github.com/mygitljf created https://github.com/llvm/llvm-project/pull/206270
I moved the poison replacement step before cleanup paths that can drop uses and leave queued operands temporarily null. This keeps region branch operands valid while the later cleanup removes dead arguments and results.
Fixes #206094
>From cd9b4fbb6547b09fac5914dace79c9f9338c72c2 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Sun, 28 Jun 2026 00:33:21 +0000
Subject: [PATCH] [mlir] Fix RDV poison cleanup
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 67 ++++++++++++--------
mlir/test/Transforms/remove-dead-values.mlir | 27 ++++++++
2 files changed, 67 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index f0a210a2ededb..8e88e84234599 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -520,15 +520,12 @@ static void processBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la,
}
}
-/// Create ub.poison ops for the given values. If a value has no uses, return
-/// an "empty" value.
-static SmallVector<Value> createPoisonedValues(OpBuilder &b,
- ValueRange values) {
- return llvm::map_to_vector(values, [&](Value value) {
- if (value.use_empty())
- return Value();
- return ub::PoisonOp::create(b, value.getLoc(), value.getType()).getResult();
- });
+/// Create a ub.poison op for the given value. If it has no uses, return an
+/// "empty" value.
+static Value createPoisonedValue(OpBuilder &b, Value value) {
+ if (value.use_empty())
+ return Value();
+ return ub::PoisonOp::create(b, value.getLoc(), value.getType()).getResult();
}
namespace {
@@ -558,7 +555,30 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
TrackingListener listener;
IRRewriter rewriter(ctx, &listener);
- // 1. Blocks, We must remove the block arguments and successor operands before
+ // 1. Operands to replace with poison. These rewrites need the original
+ // operand values for their location and type, so they must run before any
+ // cleanup that can drop uses and leave operands temporarily null.
+ LDBG() << "Replacing dead operands with poison in " << list.operands.size()
+ << " operand lists";
+ for (OperandsToCleanup &o : list.operands) {
+ if (!o.replaceWithPoison || !o.nonLive.any())
+ continue;
+ LDBG_OS([&](raw_ostream &os) {
+ os << "Replacing non-live operands [";
+ llvm::interleaveComma(o.nonLive.set_bits(), os);
+ os << "] with poison in operation: "
+ << OpWithFlags(o.op,
+ OpPrintingFlags().skipRegions().printGenericOpForm());
+ });
+ rewriter.setInsertionPoint(o.op);
+ for (auto deadIdx : o.nonLive.set_bits()) {
+ Value operand = o.op->getOperand(deadIdx);
+ assert(operand && "expected non-null operand for poison replacement");
+ o.op->setOperand(deadIdx, createPoisonedValue(rewriter, operand));
+ }
+ }
+
+ // 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) {
@@ -583,7 +603,7 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
}
}
- // 2. Successor Operands
+ // 3. Successor Operands
LDBG() << "Cleaning up " << list.successorOperands.size()
<< " successor operand lists";
for (auto &op : list.successorOperands) {
@@ -607,7 +627,7 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
}
}
- // 3. Functions
+ // 4. Functions
LDBG() << "Cleaning up " << list.functions.size() << " functions";
// Record which function arguments were erased so we can shrink call-site
// argument segments for CallOpInterface operations (e.g. ops using
@@ -641,9 +661,11 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
(void)f.funcOp.eraseResults(f.nonLiveRets);
}
- // 4. Operands
+ // 5. Operands
LDBG() << "Cleaning up " << list.operands.size() << " operand lists";
for (OperandsToCleanup &o : list.operands) {
+ if (o.replaceWithPoison)
+ continue;
// Handle call-specific cleanup only when we have a cached callee reference.
// This avoids expensive symbol lookup and is defensive against future
// changes.
@@ -686,20 +708,11 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
<< OpWithFlags(o.op,
OpPrintingFlags().skipRegions().printGenericOpForm());
});
- if (o.replaceWithPoison) {
- rewriter.setInsertionPoint(o.op);
- for (auto deadIdx : o.nonLive.set_bits()) {
- o.op->setOperand(
- deadIdx, createPoisonedValues(rewriter, o.op->getOperand(deadIdx))
- .front());
- }
- } else {
- o.op->eraseOperands(o.nonLive);
- }
+ o.op->eraseOperands(o.nonLive);
}
}
- // 5. Results
+ // 6. Results
LDBG() << "Cleaning up " << list.results.size() << " result lists";
for (auto &r : list.results) {
LDBG_OS([&](raw_ostream &os) {
@@ -712,7 +725,7 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
dropUsesAndEraseResults(rewriter, r.op, r.nonLive);
}
- // 6. Operations
+ // 7. Operations
LDBG() << "Cleaning up " << list.operations.size() << " operations";
for (Operation *op : list.operations) {
LDBG() << "Erasing operation: "
@@ -741,7 +754,7 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
continue;
rewriter.setInsertionPoint(op);
- Value poisonedValue = createPoisonedValues(rewriter, opResult).front();
+ Value poisonedValue = createPoisonedValue(rewriter, opResult);
rewriter.replaceAllUsesWith(opResult, poisonedValue);
}
@@ -749,7 +762,7 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
rewriter.eraseOp(op);
}
- // 7. Remove all dead poison ops.
+ // 8. Remove all dead poison ops.
for (ub::PoisonOp poisonOp : listener.poisonOps) {
if (poisonOp.use_empty())
poisonOp.erase();
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 64088ce15cd48..390a448060b7f 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -624,6 +624,33 @@ module @return_void_with_unused_argument {
// -----
+// CHECK-LABEL: func.func private @keep_region_branch_operands_valid() {
+func.func private @keep_region_branch_operands_valid(%arg0: memref<f64>) {
+ %false = arith.constant false
+ %cst = arith.constant 2.000000e+00 : f64
+ %0 = memref.load %arg0[] {name = "caller"} : memref<f64>
+ memref.store %cst, %arg0[] {name = "callee"} : memref<f64>
+ // CHECK: %[[COND:.*]] = ub.poison : i1
+ // CHECK: scf.if %[[COND]]
+ // CHECK: %[[YIELD0:.*]] = ub.poison : memref<f64>
+ // CHECK: scf.yield %[[YIELD0]] : memref<f64>
+ // CHECK: %[[YIELD1:.*]] = ub.poison : memref<f64>
+ // CHECK: scf.yield %[[YIELD1]] : memref<f64>
+ %1 = scf.if %false -> (memref<f64>) {
+ scf.yield %arg0 : memref<f64>
+ } else {
+ %2 = bufferization.clone %arg0 : memref<f64> to memref<f64>
+ scf.yield %2 : memref<f64>
+ }
+ %true = arith.constant true
+ %base_buffer, %offset = memref.extract_strided_metadata %arg0 : memref<f64> -> memref<f64>, index
+ %base_buffer_0, %offset_1 = memref.extract_strided_metadata %1 : memref<f64> -> memref<f64>, index
+ bufferization.dealloc (%base_buffer, %base_buffer_0 : memref<f64>, memref<f64>) if (%false, %true)
+ return
+}
+
+// -----
+
// CHECK-LABEL: module @dynamically_unreachable
module @dynamically_unreachable {
func.func @dynamically_unreachable() {
More information about the Mlir-commits
mailing list