[MLIR][RemoveDeadValues] Mark arguments of a public function Live #160242
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This diff also changes traversal order from forward to backward for region/block/ops. This order guanratees Liveness updates at a callsite can propagates to the defs of arguments. ``` ./bin/llvm-lit -v ../mlir/test/Transforms/remove-dead-values.mlir ```
|
@llvm/pr-subscribers-mlir-core Author: xin liu (navyxliu) ChangesProblemLiveness analysis is inter-procedural. If there are some unused arguments in a public function, they propagate to callers. From the perspective of RemoveDeadValues, a public function is immutable. It can't cope with this situation. One side, it deletes outgoing arguments, on the other side it keeps the function intact. SolutionThis diff attempts to add another DenseSet called 'liveSet'. The initial values are from public functions. It propagates liveness backward just like Liveness analysis. Test planbefore: after: pass AlternativesI also consider to update liveness analysis instead. For a public function, we mark its arguments live just like External. It can also fix this bug. However, it actually drops information. With a non-live argument from callee, we could replace it with a dummy placeholder like arith.constant 0 : i32. It would be impossible if we mark all arguments of a public function live. Another idea is that we force the liveness analysis of RemoveDeadValues intra-procedural. Full diff: https://github.com/llvm/llvm-project/pull/160242.diff 3 Files Affected:
diff --git a/mlir/include/mlir/IR/Visitors.h b/mlir/include/mlir/IR/Visitors.h
index 893f66ae33deb..5766d262796d6 100644
--- a/mlir/include/mlir/IR/Visitors.h
+++ b/mlir/include/mlir/IR/Visitors.h
@@ -39,6 +39,20 @@ struct ForwardIterator {
}
};
+/// This iterator enumerates the elements in "backward" order.
+struct BackwardIterator {
+ template <typename T>
+ static auto makeIterable(T &range) {
+ if constexpr (std::is_same<T, Operation>()) {
+ /// Make operations iterable: return the list of regions.
+ return llvm::reverse(range.getRegions());
+ } else {
+ /// Regions and block are already iterable.
+ return llvm::reverse(range);
+ }
+ }
+};
+
/// A utility class to encode the current walk stage for "generic" walkers.
/// When walking an operation, we can either choose a Pre/Post order walker
/// which invokes the callback on an operation before/after all its attached
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 0e84b6dd17f29..4d1cd991af6e5 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -115,9 +115,15 @@ struct RDVFinalCleanupList {
/// Return true iff at least one value in `values` is live, given the liveness
/// information in `la`.
-static bool hasLive(ValueRange values, const DenseSet<Value> &nonLiveSet,
+static bool hasLive(ValueRange values, const DenseSet<Value> &nonLiveSet, const DenseSet<Value> &liveSet,
+
RunLivenessAnalysis &la) {
for (Value value : values) {
+ if (liveSet.contains(value)) {
+ LDBG() << "Value " << value << " is marked live by CallOp";
+ return true;
+ }
+
if (nonLiveSet.contains(value)) {
LDBG() << "Value " << value << " is already marked non-live (dead)";
continue;
@@ -257,8 +263,9 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
/// - Return-like
static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,
DenseSet<Value> &nonLiveSet,
+ DenseSet<Value> &liveSet,
RDVFinalCleanupList &cl) {
- if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, la)) {
+ if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, liveSet, la)) {
LDBG() << "Simple op is not memory effect free or has live results, "
"preserving it: "
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
@@ -376,6 +383,31 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
}
}
+static void processCallOp(CallOpInterface callOp, Operation *module,
+ RunLivenessAnalysis &la,
+ DenseSet<Value> &liveSet) {
+ auto callable = callOp.getCallableForCallee();
+
+ if (auto symbolRef = callable.dyn_cast<SymbolRefAttr>()) {
+ Operation *calleeOp = SymbolTable::lookupSymbolIn(module, symbolRef);
+
+ if (auto funcOp = llvm::dyn_cast_or_null<mlir::FunctionOpInterface>(calleeOp)) {
+ // Ensure the outgoing arguments of PUBLIC functions are live
+ // because processFuncOp can not process them.
+ //
+ // Liveness treats the external function as a blackbox.
+ if (funcOp.isPublic()) {
+ for (Value arg: callOp.getArgOperands()) {
+ const Liveness *liveness = la.getLiveness(arg);
+ if (liveness && !liveness->isLive) {
+ liveSet.insert(arg);
+ }
+ }
+ }
+ }
+ }
+}
+
/// Process a region branch operation `regionBranchOp` using the liveness
/// information in `la`. The processing involves two scenarios:
///
@@ -408,6 +440,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
RunLivenessAnalysis &la,
DenseSet<Value> &nonLiveSet,
+ DenseSet<Value> &liveSet,
RDVFinalCleanupList &cl) {
LDBG() << "Processing region branch op: "
<< OpWithFlags(regionBranchOp, OpPrintingFlags().skipRegions());
@@ -616,7 +649,7 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
// attributed to something else.
// Do (1') and (2').
if (isMemoryEffectFree(regionBranchOp.getOperation()) &&
- !hasLive(regionBranchOp->getResults(), nonLiveSet, la)) {
+ !hasLive(regionBranchOp->getResults(), nonLiveSet, liveSet, la)) {
cl.operations.push_back(regionBranchOp.getOperation());
return;
}
@@ -834,16 +867,18 @@ void RemoveDeadValues::runOnOperation() {
// Tracks values eligible for erasure - complements liveness analysis to
// identify "droppable" values.
DenseSet<Value> deadVals;
+ // mark outgoing arguments to a public function LIVE.
+ DenseSet<Value> liveVals;
// Maintains a list of Ops, values, branches, etc., slated for cleanup at the
// end of this pass.
RDVFinalCleanupList finalCleanupList;
- module->walk([&](Operation *op) {
+ module->walk<WalkOrder::PostOrder, BackwardIterator>([&](Operation *op) {
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
processFuncOp(funcOp, module, la, deadVals, finalCleanupList);
} else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
- processRegionBranchOp(regionBranchOp, la, deadVals, finalCleanupList);
+ processRegionBranchOp(regionBranchOp, la, deadVals, liveVals, finalCleanupList);
} else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
processBranchOp(branchOp, la, deadVals, finalCleanupList);
} else if (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {
@@ -852,8 +887,13 @@ void RemoveDeadValues::runOnOperation() {
} else if (isa<CallOpInterface>(op)) {
// Nothing to do because this op is associated with a function op and gets
// cleaned when the latter is cleaned.
+ //
+ // The only exception is public callee. By default, Liveness analysis is inter-procedural.
+ // Unused arguments of a public function nonLive and are propagated to the caller.
+ // processCallOp puts them to liveVals.
+ processCallOp(cast<CallOpInterface>(op), module, la, liveVals);
} else {
- processSimpleOp(op, la, deadVals, finalCleanupList);
+ processSimpleOp(op, la, deadVals, liveVals, finalCleanupList);
}
});
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index fa2c145bd3701..1580009c74d4d 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -569,6 +569,24 @@ module @return_void_with_unused_argument {
call @fn_return_void_with_unused_argument(%arg0, %unused) : (i32, memref<4xi32>) -> ()
return %unused : memref<4xi32>
}
+
+ // the function is immutable because it is public.
+ func.func public @immutable_fn_return_void_with_unused_argument(%arg0: i32, %unused: i32) -> () {
+ %sum = arith.addi %arg0, %arg0 : i32
+ %c0 = arith.constant 0 : index
+ %buf = memref.alloc() : memref<1xi32>
+ memref.store %sum, %buf[%c0] : memref<1xi32>
+ return
+ }
+ // CHECK-LABEL: func.func @main2
+ // CHECK-SAME: (%[[ARG0_MAIN:.*]]: i32)
+ // CHECK: %[[UNUSED:.*]] = arith.constant 0 : i32
+ // CHECK: call @immutable_fn_return_void_with_unused_argument(%[[ARG0_MAIN]], %[[UNUSED]]) : (i32, i32) -> ()
+ func.func @main2(%arg0: i32) -> () {
+ %zero = arith.constant 0 : i32
+ call @immutable_fn_return_void_with_unused_argument(%arg0, %zero) : (i32, i32) -> ()
+ return
+ }
}
// -----
|
|
@llvm/pr-subscribers-mlir Author: xin liu (navyxliu) ChangesProblemLiveness analysis is inter-procedural. If there are some unused arguments in a public function, they propagate to callers. From the perspective of RemoveDeadValues, a public function is immutable. It can't cope with this situation. One side, it deletes outgoing arguments, on the other side it keeps the function intact. SolutionThis diff attempts to add another DenseSet called 'liveSet'. The initial values are from public functions. It propagates liveness backward just like Liveness analysis. Test planbefore: after: pass AlternativesI also consider to update liveness analysis instead. For a public function, we mark its arguments live just like External. It can also fix this bug. However, it actually drops information. With a non-live argument from callee, we could replace it with a dummy placeholder like arith.constant 0 : i32. It would be impossible if we mark all arguments of a public function live. Another idea is that we force the liveness analysis of RemoveDeadValues intra-procedural. Full diff: https://github.com/llvm/llvm-project/pull/160242.diff 3 Files Affected:
diff --git a/mlir/include/mlir/IR/Visitors.h b/mlir/include/mlir/IR/Visitors.h
index 893f66ae33deb..5766d262796d6 100644
--- a/mlir/include/mlir/IR/Visitors.h
+++ b/mlir/include/mlir/IR/Visitors.h
@@ -39,6 +39,20 @@ struct ForwardIterator {
}
};
+/// This iterator enumerates the elements in "backward" order.
+struct BackwardIterator {
+ template <typename T>
+ static auto makeIterable(T &range) {
+ if constexpr (std::is_same<T, Operation>()) {
+ /// Make operations iterable: return the list of regions.
+ return llvm::reverse(range.getRegions());
+ } else {
+ /// Regions and block are already iterable.
+ return llvm::reverse(range);
+ }
+ }
+};
+
/// A utility class to encode the current walk stage for "generic" walkers.
/// When walking an operation, we can either choose a Pre/Post order walker
/// which invokes the callback on an operation before/after all its attached
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 0e84b6dd17f29..4d1cd991af6e5 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -115,9 +115,15 @@ struct RDVFinalCleanupList {
/// Return true iff at least one value in `values` is live, given the liveness
/// information in `la`.
-static bool hasLive(ValueRange values, const DenseSet<Value> &nonLiveSet,
+static bool hasLive(ValueRange values, const DenseSet<Value> &nonLiveSet, const DenseSet<Value> &liveSet,
+
RunLivenessAnalysis &la) {
for (Value value : values) {
+ if (liveSet.contains(value)) {
+ LDBG() << "Value " << value << " is marked live by CallOp";
+ return true;
+ }
+
if (nonLiveSet.contains(value)) {
LDBG() << "Value " << value << " is already marked non-live (dead)";
continue;
@@ -257,8 +263,9 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
/// - Return-like
static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,
DenseSet<Value> &nonLiveSet,
+ DenseSet<Value> &liveSet,
RDVFinalCleanupList &cl) {
- if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, la)) {
+ if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, liveSet, la)) {
LDBG() << "Simple op is not memory effect free or has live results, "
"preserving it: "
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
@@ -376,6 +383,31 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
}
}
+static void processCallOp(CallOpInterface callOp, Operation *module,
+ RunLivenessAnalysis &la,
+ DenseSet<Value> &liveSet) {
+ auto callable = callOp.getCallableForCallee();
+
+ if (auto symbolRef = callable.dyn_cast<SymbolRefAttr>()) {
+ Operation *calleeOp = SymbolTable::lookupSymbolIn(module, symbolRef);
+
+ if (auto funcOp = llvm::dyn_cast_or_null<mlir::FunctionOpInterface>(calleeOp)) {
+ // Ensure the outgoing arguments of PUBLIC functions are live
+ // because processFuncOp can not process them.
+ //
+ // Liveness treats the external function as a blackbox.
+ if (funcOp.isPublic()) {
+ for (Value arg: callOp.getArgOperands()) {
+ const Liveness *liveness = la.getLiveness(arg);
+ if (liveness && !liveness->isLive) {
+ liveSet.insert(arg);
+ }
+ }
+ }
+ }
+ }
+}
+
/// Process a region branch operation `regionBranchOp` using the liveness
/// information in `la`. The processing involves two scenarios:
///
@@ -408,6 +440,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
RunLivenessAnalysis &la,
DenseSet<Value> &nonLiveSet,
+ DenseSet<Value> &liveSet,
RDVFinalCleanupList &cl) {
LDBG() << "Processing region branch op: "
<< OpWithFlags(regionBranchOp, OpPrintingFlags().skipRegions());
@@ -616,7 +649,7 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp,
// attributed to something else.
// Do (1') and (2').
if (isMemoryEffectFree(regionBranchOp.getOperation()) &&
- !hasLive(regionBranchOp->getResults(), nonLiveSet, la)) {
+ !hasLive(regionBranchOp->getResults(), nonLiveSet, liveSet, la)) {
cl.operations.push_back(regionBranchOp.getOperation());
return;
}
@@ -834,16 +867,18 @@ void RemoveDeadValues::runOnOperation() {
// Tracks values eligible for erasure - complements liveness analysis to
// identify "droppable" values.
DenseSet<Value> deadVals;
+ // mark outgoing arguments to a public function LIVE.
+ DenseSet<Value> liveVals;
// Maintains a list of Ops, values, branches, etc., slated for cleanup at the
// end of this pass.
RDVFinalCleanupList finalCleanupList;
- module->walk([&](Operation *op) {
+ module->walk<WalkOrder::PostOrder, BackwardIterator>([&](Operation *op) {
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
processFuncOp(funcOp, module, la, deadVals, finalCleanupList);
} else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
- processRegionBranchOp(regionBranchOp, la, deadVals, finalCleanupList);
+ processRegionBranchOp(regionBranchOp, la, deadVals, liveVals, finalCleanupList);
} else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
processBranchOp(branchOp, la, deadVals, finalCleanupList);
} else if (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {
@@ -852,8 +887,13 @@ void RemoveDeadValues::runOnOperation() {
} else if (isa<CallOpInterface>(op)) {
// Nothing to do because this op is associated with a function op and gets
// cleaned when the latter is cleaned.
+ //
+ // The only exception is public callee. By default, Liveness analysis is inter-procedural.
+ // Unused arguments of a public function nonLive and are propagated to the caller.
+ // processCallOp puts them to liveVals.
+ processCallOp(cast<CallOpInterface>(op), module, la, liveVals);
} else {
- processSimpleOp(op, la, deadVals, finalCleanupList);
+ processSimpleOp(op, la, deadVals, liveVals, finalCleanupList);
}
});
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index fa2c145bd3701..1580009c74d4d 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -569,6 +569,24 @@ module @return_void_with_unused_argument {
call @fn_return_void_with_unused_argument(%arg0, %unused) : (i32, memref<4xi32>) -> ()
return %unused : memref<4xi32>
}
+
+ // the function is immutable because it is public.
+ func.func public @immutable_fn_return_void_with_unused_argument(%arg0: i32, %unused: i32) -> () {
+ %sum = arith.addi %arg0, %arg0 : i32
+ %c0 = arith.constant 0 : index
+ %buf = memref.alloc() : memref<1xi32>
+ memref.store %sum, %buf[%c0] : memref<1xi32>
+ return
+ }
+ // CHECK-LABEL: func.func @main2
+ // CHECK-SAME: (%[[ARG0_MAIN:.*]]: i32)
+ // CHECK: %[[UNUSED:.*]] = arith.constant 0 : i32
+ // CHECK: call @immutable_fn_return_void_with_unused_argument(%[[ARG0_MAIN]], %[[UNUSED]]) : (i32, i32) -> ()
+ func.func @main2(%arg0: i32) -> () {
+ %zero = arith.constant 0 : i32
+ call @immutable_fn_return_void_with_unused_argument(%arg0, %zero) : (i32, i32) -> ()
+ return
+ }
}
// -----
|
|
cc @victor-eds |
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
|
cc @joker-eph |
joker-eph
reviewed
Sep 23, 2025
joker-eph
reviewed
Sep 23, 2025
| static void processCallOp(CallOpInterface callOp, Operation *module, | ||
| RunLivenessAnalysis &la, | ||
| DenseSet<Value> &liveSet) { | ||
| auto callable = callOp.getCallableForCallee(); |
There was a problem hiding this comment.
We should not do this when the analysis is not inter procedural, instead should just mark things live and return, otherwise we're gonna trigger race conditions here.
There was a problem hiding this comment.
That also points to a need for a test showing that when the pass is scheduled at the function level, we should make everything live and only optimize when scheduled at the module level.
There was a problem hiding this comment.
Hi, @joker-eph
I am new to the MLIR community. I make a tentative change to fix this problem. I'm open to alternatives. Next time, should I post on discourse and get a clear direction beforehand?
Are you suggesting RemoveDeadValues should support 2 modes?
- interprocedural mode. It depends on interprocedural liveness and operate module.
- intraprocedural mode. It depends on intraprocedural liveness and operate on individual Functions.
How does RemoveDeadValues distinct two modes?
I think it's 1) by default. Actually, it caught me a surprise. I didn't realize that until I hit this bug.
In the downstream project, I inserted this code and has used it 3 months. It reconfigures 'Liveness dataflow' intra-procedural. It works to us, but it only eschews the bug, not solve it.
struct RunIntraproceduralLivenessAnalysis {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RunIntraproceduralLivenessAnalysis)
RunIntraproceduralLivenessAnalysis(Operation *op) : solver(DataFlowConfig().setInterprocedural(false)) {
SymbolTableCollection symbolTable;
loadBaselineAnalyses(solver);
solver.load<LivenessAnalysis>(symbolTable);
(void)solver.initializeAndRun(op);
}
const Liveness *getLiveness(Value val) {
return solver.lookupState<Liveness>(val);
}
private:
/// Stores the result of the liveness analysis that was run.
DataFlowSolver solver;
};
using RunLivenessAnalysis = RunIntraproceduralLivenessAnalysis;It just a boolean flag. now it is fixed in static. Should we make it configurable in runtime first?
There was a problem hiding this comment.
Function passes may run concurrently on different functions. Therefore, they should generally avoid looking at adjacent functions as that would cause race conditions. I'm not sure what will happen here as each pass may be computing its own instance of the liveness analysis scoped to the function it operates on so it may have no info at all about other functions. A straightforward fix to this is to make this a module pass so it cannot run on functions.
There was a problem hiding this comment.
- I fixed some issues recently to ensure that data-flow analysis based analyses and transformations honor the scope of the pass. I thought RemoveDeadValue was one of these, but it seems I got confused about it.
- Actually I looked at the code and it may work in a subtle way today, we do in the
runOnFunction():
module->walk([&](Operation *op) {
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
processFuncOp(funcOp, module, la, deadVals, finalCleanupList);
If the pass is scheduled on a func.func, then this walk does not start on the module (despite the variable name) but on the function. And so it'll never visit the func.func itself and the processFuncOp won't be called.
That said, processCallOp is gonna be called.
There was a problem hiding this comment.
Here is an example of a fix where I detected whether we're working on a symbol table or not: 13ae9ea
This was referenced Sep 30, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Problem
Liveness analysis is inter-procedural. If there are some unused arguments in a public function, they propagate to callers. From the perspective of RemoveDeadValues, a public function is immutable. It can't cope with this situation. One side, it deletes outgoing arguments, on the other side it keeps the function intact.
Solution
This diff attempts to add another DenseSet called 'liveSet'. The initial values are from public functions. It propagates liveness backward just like Liveness analysis.
Test plan
before:
after: pass
Alternatives
I also consider to update liveness analysis instead. For a public function, we mark its arguments live just like External. It can also fix this bug. However, it actually drops information. With a non-live argument from callee, we could replace it with a dummy placeholder like arith.constant 0 : i32. It would be impossible if we mark all arguments of a public function live.
Another idea is that we force the liveness analysis of RemoveDeadValues intra-procedural.
it also miss optimization opportunity.