Skip to content

[MLIR][RemoveDeadValues] Mark arguments of a public function Live #160242

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

[MLIR][RemoveDeadValues] Mark arguments of a public function Live #160242

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b83d8dc
[MLIR][RemoveDeadValues] Mark arguments of a public function Live
Sep 22, 2025
03a8ef6
Merge branch 'llvm:main' into backward_remove_dead_values
navyxliu Sep 23, 2025
84cb2b4
Fix formatter issue.
navyxliu Sep 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions mlir/include/mlir/IR/Visitors.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
54 changes: 48 additions & 6 deletions mlir/lib/Transforms/RemoveDeadValues.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,8 +116,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,
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;
Expand Down Expand Up @@ -257,8 +264,9 @@ static SmallVector<OpOperand *> operandsToOpOperands(OperandRange operands) {
/// - Return-like
static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,
DenseSet<Value> &nonLiveSet,
RDVFinalCleanupList &cl) {
if (!isMemoryEffectFree(op) || hasLive(op->getResults(), nonLiveSet, la)) {
DenseSet<Value> &liveSet, RDVFinalCleanupList &cl) {
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());
Expand Down Expand Up @@ -376,6 +384,31 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
}
}

static void processCallOp(CallOpInterface callOp, Operation *module,
RunLivenessAnalysis &la, DenseSet<Value> &liveSet) {
auto callable = callOp.getCallableForCallee();
Copy link
Collaborator

@joker-eph joker-eph Sep 23, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

  1. interprocedural mode. It depends on interprocedural liveness and operate module.
  2. 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?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

@joker-eph joker-eph Sep 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. 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.
  2. 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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is an example of a fix where I detected whether we're working on a symbol table or not: 13ae9ea


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:
///
Expand Down Expand Up @@ -408,6 +441,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());
Expand Down Expand Up @@ -616,7 +650,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;
}
Expand Down Expand Up @@ -834,16 +868,19 @@ 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>()) {
Expand All @@ -852,8 +889,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);
}
});

Expand Down
18 changes: 18 additions & 0 deletions mlir/test/Transforms/remove-dead-values.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}

// -----
Expand Down