[Mlir-commits] [mlir] [mlir][RemoveDeadValues] Use `SymbolUserMap` to avoid quadratic symbol lookups (PR #205448)
Victor Perez
llvmlistbot at llvm.org
Wed Jun 24 07:00:41 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/205448
>From 44bd45f7340bc591b42a333c450eb7865de6ed4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?= <victorperez at fb.com>
Date: Tue, 23 Jun 2026 15:47:02 -0700
Subject: [PATCH 1/3] [mlir][RemoveDeadValues] Use `SymbolUserMap` to avoid
quadratic symbol lookups
`processFuncOp` previously called `funcOp.getSymbolUses(module)` for every
function, which walks the entire module to find that function's callers.
Since `processFuncOp` runs once per function, the pass was effectively
*O(numFunctions * numOperations)*.
Build a `SymbolUserMap` once up front in `runOnOperation()` and look up each
function's callers in *O(1)*, making the collection phase linear in the
size of the module.
This is behavior-preserving: the map is read only during the
mutation-free collection walk, and all IR erasure happens afterwards in
`cleanUpDeadVals`, so the map cannot become stale while it is in use.
Signed-off-by: Victor Perez Carrasco <victor.pc.upm at gmail.com>
Co-authored-by: mlevesquedion <mlevesquedion at meta.com>
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 28 ++++++++++++++----------
1 file changed, 16 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index f0a210a2ededb..7210e4982678a 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -259,7 +259,7 @@ static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,
}
/// Process a function-like operation `funcOp` using the liveness analysis `la`
-/// and the IR in `module`. If it is not public or external:
+/// and `symbolUserMap`. If it is not public or external:
/// (1) Adding its non-live arguments to a list for future removal.
/// (2) Marking their corresponding operands in its callers for removal.
/// (3) Identifying and enqueueing unnecessary terminator operands
@@ -268,7 +268,8 @@ static void processSimpleOp(Operation *op, RunLivenessAnalysis &la,
/// (5) Collecting the uses of these return values in its callers for future
/// removal.
/// (6) Marking all its results as non-live values.
-static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
+static void processFuncOp(FunctionOpInterface funcOp,
+ const SymbolUserMap &symbolUserMap,
RunLivenessAnalysis &la, DenseSet<Value> &nonLiveSet,
RDVFinalCleanupList &cl) {
LDBG() << "Processing function op: "
@@ -279,9 +280,9 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
<< funcOp.getOperation()->getName();
return;
}
- SymbolTable::UseRange uses = *funcOp.getSymbolUses(module);
- if (llvm::any_of(uses, [](SymbolTable::SymbolUse use) {
- return !isa<CallOpInterface>(use.getUser());
+ ArrayRef<Operation *> users = symbolUserMap.getUsers(funcOp);
+ if (llvm::any_of(users, [](Operation *user) {
+ return !isa<CallOpInterface>(user);
})) {
// If a non-call operation references the function (e.g. spirv.EntryPoint),
// we cannot safely remove arguments or return values since we don't know
@@ -301,8 +302,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
// Do (2). (Skip creating generic operand cleanup entries for call ops.
// Call arguments will be removed in the call-site specific segment-aware
// cleanup, avoiding generic eraseOperands bitvector mechanics.)
- for (SymbolTable::SymbolUse use : uses) {
- Operation *callOp = use.getUser();
+ for (Operation *callOp : users) {
// Push an empty operand cleanup entry so that call-site specific logic in
// cleanUpDeadVals runs (it keys off CallOpInterface). The BitVector is
// intentionally all false to avoid generic erasure.
@@ -336,8 +336,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
// since it forwards only to non-live value(s) (%1#1).
size_t numReturns = funcOp.getNumResults();
BitVector nonLiveRets(numReturns, true);
- for (SymbolTable::SymbolUse use : uses) {
- Operation *callOp = use.getUser();
+ for (Operation *callOp : users) {
assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
BitVector liveCallRets = markLives(callOp->getResults(), nonLiveSet, la);
nonLiveRets &= liveCallRets.flip();
@@ -360,8 +359,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
// Do (5) and (6).
if (numReturns == 0)
return;
- for (SymbolTable::SymbolUse use : uses) {
- Operation *callOp = use.getUser();
+ for (Operation *callOp : users) {
assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
cl.results.push_back({callOp, nonLiveRets});
collectNonLiveValues(nonLiveSet, callOp->getResults(), nonLiveRets);
@@ -770,6 +768,12 @@ void RemoveDeadValues::runOnOperation() {
auto &la = getAnalysis<RunLivenessAnalysis>();
Operation *module = getOperation();
+ // Build a symbol user map once up front so that processFuncOp can look up the
+ // callers of each function in O(1). Otherwise, each call would walk the entire
+ // module to find the callers, making the pass O(numFunctions * numOperations).
+ SymbolTableCollection symbolTableCollection;
+ SymbolUserMap symbolUserMap(symbolTableCollection, module);
+
// Tracks values eligible for erasure - complements liveness analysis to
// identify "droppable" values.
DenseSet<Value> deadVals;
@@ -780,7 +784,7 @@ void RemoveDeadValues::runOnOperation() {
module->walk([&](Operation *op) {
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
- processFuncOp(funcOp, module, la, deadVals, finalCleanupList);
+ processFuncOp(funcOp, symbolUserMap, la, deadVals, finalCleanupList);
} else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
processRegionBranchOp(regionBranchOp, la, deadVals, finalCleanupList);
} else if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
>From f845447c9941794099cf725b817ba4cf2db56e46 Mon Sep 17 00:00:00 2001
From: Victor Perez <victor.pc.upm at gmail.com>
Date: Tue, 23 Jun 2026 18:55:32 -0400
Subject: [PATCH 2/3] Refactor lambda formatting in RemoveDeadValues.cpp
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 7210e4982678a..6ba45de8dbee0 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -281,9 +281,8 @@ static void processFuncOp(FunctionOpInterface funcOp,
return;
}
ArrayRef<Operation *> users = symbolUserMap.getUsers(funcOp);
- if (llvm::any_of(users, [](Operation *user) {
- return !isa<CallOpInterface>(user);
- })) {
+ if (llvm::any_of(
+ users, [](Operation *user) { return !isa<CallOpInterface>(user); })) {
// If a non-call operation references the function (e.g. spirv.EntryPoint),
// we cannot safely remove arguments or return values since we don't know
// what the user expects. Skip this function entirely.
@@ -769,8 +768,9 @@ void RemoveDeadValues::runOnOperation() {
Operation *module = getOperation();
// Build a symbol user map once up front so that processFuncOp can look up the
- // callers of each function in O(1). Otherwise, each call would walk the entire
- // module to find the callers, making the pass O(numFunctions * numOperations).
+ // callers of each function in O(1). Otherwise, each call would walk the
+ // entire module to find the callers, making the pass O(numFunctions *
+ // numOperations).
SymbolTableCollection symbolTableCollection;
SymbolUserMap symbolUserMap(symbolTableCollection, module);
>From a2e78c2efb86c922f2c4ef9161cf81a00a9a9bae Mon Sep 17 00:00:00 2001
From: Victor Perez <victor.pc.upm at gmail.com>
Date: Wed, 24 Jun 2026 10:00:29 -0400
Subject: [PATCH 3/3] Correct formatting in RemoveDeadValues.cpp
Fix formatting issues in RemoveDeadValues.cpp.
---
mlir/lib/Transforms/RemoveDeadValues.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 6ba45de8dbee0..671e06bc10b6a 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -282,7 +282,7 @@ static void processFuncOp(FunctionOpInterface funcOp,
}
ArrayRef<Operation *> users = symbolUserMap.getUsers(funcOp);
if (llvm::any_of(
- users, [](Operation *user) { return !isa<CallOpInterface>(user); })) {
+ users, [](Operation *user) { return !isa<CallOpInterface>(user); })) {
// If a non-call operation references the function (e.g. spirv.EntryPoint),
// we cannot safely remove arguments or return values since we don't know
// what the user expects. Skip this function entirely.
More information about the Mlir-commits
mailing list