[clang] [CIR] Prune CIR type subtrees when collecting symbol uses (PR #234811)
Bruno Cardoso Lopes via cfe-commits
cfe-commits at lists.llvm.org
Fri Oct 23 19:24:17 PDT 2026
https://github.com/bcardosolopes updated https://github.com/llvm/llvm-project/pull/234811
>From 3a983d789d185e90f1c4380e44cd6b630708331b Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Thu, 22 Oct 2026 17:08:30 -0700
Subject: [PATCH 1/2] [CIR] Prune CIR type subtrees when collecting symbol uses
applyReplacements builds a symbol user map over the module, and MLIR's
generic walk descends into every nested record of every function
signature even though no symbol reference can hide in a CIR type. Add
CIRSymbolUserMap, a CIR-aware fork of mlir::SymbolUserMap that skips
CIR-dialect type subtrees and otherwise walks exactly like upstream,
and use it for the collection.
The fork keeps upstream's traversal shape (post-order table walk, no
descent into nested symbol tables, per-op visited set) and drops only
what applyReplacements never queries (visibility, unknown users).
Debug builds assert the prune premise on every skipped type, so a CIR
type that ever gains a symbol reference fails loudly instead of
silently dropping uses.
Measured on the investigation branch: SemaExpr.cpp -fclangir 51.4s to
50.6s, .text byte-identical.
Assisted-by: Muse Code (Muse).
---
clang/lib/CIR/CodeGen/CIRGenModule.cpp | 139 ++++++++++++++++++++++++-
1 file changed, 137 insertions(+), 2 deletions(-)
diff --git a/clang/lib/CIR/CodeGen/CIRGenModule.cpp b/clang/lib/CIR/CodeGen/CIRGenModule.cpp
index 8801a730e8d67..a1555d2b61fd4 100644
--- a/clang/lib/CIR/CodeGen/CIRGenModule.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenModule.cpp
@@ -1874,9 +1874,141 @@ void CIRGenModule::addReplacement(StringRef name, mlir::Operation *op) {
replacements[name] = op;
}
+namespace {
+// Like mlir::SymbolUserMap, but prunes CIR type subtrees when collecting
+// symbol uses. CIR types (see CIRTypes.td) only contain nested types, name
+// strings, memory-space attributes and bit-field descriptors, so no symbol
+// reference can hide in them. MLIR's generic walk still descends into every
+// nested record of every function signature. Types from any other dialect
+// are walked fully, so this stays correct if non-CIR types with attribute
+// parameters ever show up here. The prune relies on CIR pointer address
+// spaces staying free of symbol references (see PointerType::verify).
+// Unlike mlir::SymbolUserMap this tracks no visibility or unknown users;
+// applyReplacements only needs getUsers and replaceAllUsesWith.
+class CIRSymbolUserMap {
+public:
+ CIRSymbolUserMap(mlir::SymbolTableCollection &symbolTables,
+ mlir::Operation *module)
+ : symbolTables(symbolTables) {
+ cirDialect = module->getContext()->getLoadedDialect<cir::CIRDialect>();
+ llvm::SmallVector<mlir::Operation *> symbols;
+ auto walkFn = [&](mlir::Operation *tableOp, bool /*allUsesVisible*/) {
+ for (mlir::Operation &op : tableOp->getRegion(0).getOps())
+ collectFromOp(&op, tableOp, symbols);
+ };
+ // A root with no containing block has no symbol users in enclosing IR.
+ mlir::SymbolTable::walkSymbolTables(module, !module->getBlock(), walkFn);
+ }
+
+ llvm::ArrayRef<mlir::Operation *> getUsers(mlir::Operation *symbol) const {
+ auto it = symbolToUsers.find(symbol);
+ return it != symbolToUsers.end() ? it->second.getArrayRef()
+ : llvm::ArrayRef<mlir::Operation *>();
+ }
+
+ void replaceAllUsesWith(mlir::Operation *symbol,
+ mlir::StringAttr newSymbolName) {
+ auto it = symbolToUsers.find(symbol);
+ if (it == symbolToUsers.end())
+ return;
+
+ // Replace the uses within the users of `symbol`.
+ for (mlir::Operation *user : it->second)
+ (void)mlir::SymbolTable::replaceAllSymbolUses(symbol, newSymbolName,
+ user);
+
+ // Move the current users of `symbol` to the new symbol if it is in the
+ // symbol table.
+ mlir::Operation *newSymbol =
+ symbolTables.lookupSymbolIn(symbol->getParentOp(), newSymbolName);
+ if (newSymbol != symbol) {
+ // Transfer over the users to the new symbol. The reference to the old
+ // one is fetched again as the iterator is invalidated during the
+ // insertion.
+ auto newIt = symbolToUsers.try_emplace(newSymbol);
+ auto oldIt = symbolToUsers.find(symbol);
+ assert(oldIt != symbolToUsers.end() && "missing old users list");
+ if (newIt.second)
+ newIt.first->second = std::move(oldIt->second);
+ else
+ newIt.first->second.set_union(oldIt->second);
+ symbolToUsers.erase(oldIt);
+ }
+ }
+
+private:
+ void collectFromOp(mlir::Operation *op, mlir::Operation *tableOp,
+ llvm::SmallVectorImpl<mlir::Operation *> &symbols) {
+ // Like mlir::walkSymbolUses: an unregistered op with a region might be
+ // a symbol table, whose uses belong to another scope. CIRGen output is
+ // fully registered, so this only guards the fork.
+ assert((op->getNumRegions() != 1 || op->getDialect()) &&
+ "unregistered op with a region");
+ currentOp = op;
+ visited.clear();
+ auto visit = [&](mlir::Attribute attr) { visitAttr(attr); };
+ visit(op->getRawDictionaryAttrs());
+ op->getName().walkInherentAttrs(
+ op, [&](llvm::StringRef, mlir::Attribute &attr) { visit(attr); });
+ for (auto [userOp, ref] : pendingUses) {
+ symbols.clear();
+ (void)symbolTables.lookupSymbolIn(tableOp, ref, symbols);
+ for (mlir::Operation *symbolOp : symbols)
+ symbolToUsers[symbolOp].insert(userOp);
+ }
+ pendingUses.clear();
+
+ // Recurse like mlir::walkSymbolUses: a symbol table defines a new scope.
+ if (op->hasTrait<mlir::OpTrait::SymbolTable>())
+ return;
+ for (mlir::Region ®ion : op->getRegions())
+ for (mlir::Operation &nested : region.getOps())
+ collectFromOp(&nested, tableOp, symbols);
+ }
+
+ void visitAttr(mlir::Attribute attr) {
+ // Visit each sub-element once, like the walker MLIR uses.
+ if (!visited.insert(attr.getAsOpaquePointer()).second)
+ return;
+ if (auto ref = mlir::dyn_cast<mlir::SymbolRefAttr>(attr)) {
+ pendingUses.emplace_back(currentOp, ref);
+ // Don't walk nested references.
+ return;
+ }
+ attr.walkImmediateSubElements([&](mlir::Attribute sub) { visitAttr(sub); },
+ [&](mlir::Type type) { visitType(type); });
+ }
+
+ void visitType(mlir::Type type) {
+ // Prune: audited above, CIR types hold no symbol references.
+ if (cirDialect && &type.getDialect() == cirDialect) {
+#ifndef NDEBUG
+ bool hasSymbolRef = false;
+ type.walk([&](mlir::SymbolRefAttr) { hasSymbolRef = true; });
+ assert(!hasSymbolRef && "CIR type contains a symbol reference");
+#endif
+ return;
+ }
+ if (!visited.insert(type.getAsOpaquePointer()).second)
+ return;
+ type.walkImmediateSubElements([&](mlir::Attribute sub) { visitAttr(sub); },
+ [&](mlir::Type sub) { visitType(sub); });
+ }
+
+ mlir::SymbolTableCollection &symbolTables;
+ cir::CIRDialect *cirDialect = nullptr;
+ mlir::Operation *currentOp = nullptr;
+ llvm::SmallPtrSet<const void *, 32> visited;
+ llvm::SmallVector<std::pair<mlir::Operation *, mlir::SymbolRefAttr>>
+ pendingUses;
+ llvm::DenseMap<mlir::Operation *, llvm::SetVector<mlir::Operation *>>
+ symbolToUsers;
+};
+} // namespace
+
#ifndef NDEBUG
static bool verifyPointerTypeArgs(cir::FuncOp oldF, cir::FuncOp newF,
- mlir::SymbolUserMap &userMap) {
+ CIRSymbolUserMap &userMap) {
for (mlir::Operation *user : userMap.getUsers(oldF)) {
auto call = mlir::dyn_cast<cir::CallOp>(user);
if (!call)
@@ -1900,8 +2032,11 @@ void CIRGenModule::applyReplacements() {
// Build a symbol user map once — this walks the module O(M) one time.
// Previously, each replaceAllSymbolUses call walked the entire module,
// giving O(R × M) quadratic behavior for R replacements.
+ // CIRSymbolUserMap additionally prunes CIR type subtrees (see above):
+ // MLIR's generic walk descends into every nested record of every function
+ // signature even though no symbol reference can hide there.
mlir::SymbolTableCollection symbolTableCollection;
- mlir::SymbolUserMap userMap(symbolTableCollection, getModule());
+ CIRSymbolUserMap userMap(symbolTableCollection, getModule());
for (auto &i : replacements) {
StringRef mangledName = i.first;
>From aa79d151c739707598ea0e7f419e1b991805e269 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Fri, 23 Oct 2026 19:23:21 -0700
Subject: [PATCH 2/2] Address review: collect users with Attribute::walk, drop
the SymbolUserMap fork
---
clang/lib/CIR/CodeGen/CIRGenModule.cpp | 195 ++++++-------------------
1 file changed, 48 insertions(+), 147 deletions(-)
diff --git a/clang/lib/CIR/CodeGen/CIRGenModule.cpp b/clang/lib/CIR/CodeGen/CIRGenModule.cpp
index a1555d2b61fd4..230c65ebd698e 100644
--- a/clang/lib/CIR/CodeGen/CIRGenModule.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenModule.cpp
@@ -1874,142 +1874,10 @@ void CIRGenModule::addReplacement(StringRef name, mlir::Operation *op) {
replacements[name] = op;
}
-namespace {
-// Like mlir::SymbolUserMap, but prunes CIR type subtrees when collecting
-// symbol uses. CIR types (see CIRTypes.td) only contain nested types, name
-// strings, memory-space attributes and bit-field descriptors, so no symbol
-// reference can hide in them. MLIR's generic walk still descends into every
-// nested record of every function signature. Types from any other dialect
-// are walked fully, so this stays correct if non-CIR types with attribute
-// parameters ever show up here. The prune relies on CIR pointer address
-// spaces staying free of symbol references (see PointerType::verify).
-// Unlike mlir::SymbolUserMap this tracks no visibility or unknown users;
-// applyReplacements only needs getUsers and replaceAllUsesWith.
-class CIRSymbolUserMap {
-public:
- CIRSymbolUserMap(mlir::SymbolTableCollection &symbolTables,
- mlir::Operation *module)
- : symbolTables(symbolTables) {
- cirDialect = module->getContext()->getLoadedDialect<cir::CIRDialect>();
- llvm::SmallVector<mlir::Operation *> symbols;
- auto walkFn = [&](mlir::Operation *tableOp, bool /*allUsesVisible*/) {
- for (mlir::Operation &op : tableOp->getRegion(0).getOps())
- collectFromOp(&op, tableOp, symbols);
- };
- // A root with no containing block has no symbol users in enclosing IR.
- mlir::SymbolTable::walkSymbolTables(module, !module->getBlock(), walkFn);
- }
-
- llvm::ArrayRef<mlir::Operation *> getUsers(mlir::Operation *symbol) const {
- auto it = symbolToUsers.find(symbol);
- return it != symbolToUsers.end() ? it->second.getArrayRef()
- : llvm::ArrayRef<mlir::Operation *>();
- }
-
- void replaceAllUsesWith(mlir::Operation *symbol,
- mlir::StringAttr newSymbolName) {
- auto it = symbolToUsers.find(symbol);
- if (it == symbolToUsers.end())
- return;
-
- // Replace the uses within the users of `symbol`.
- for (mlir::Operation *user : it->second)
- (void)mlir::SymbolTable::replaceAllSymbolUses(symbol, newSymbolName,
- user);
-
- // Move the current users of `symbol` to the new symbol if it is in the
- // symbol table.
- mlir::Operation *newSymbol =
- symbolTables.lookupSymbolIn(symbol->getParentOp(), newSymbolName);
- if (newSymbol != symbol) {
- // Transfer over the users to the new symbol. The reference to the old
- // one is fetched again as the iterator is invalidated during the
- // insertion.
- auto newIt = symbolToUsers.try_emplace(newSymbol);
- auto oldIt = symbolToUsers.find(symbol);
- assert(oldIt != symbolToUsers.end() && "missing old users list");
- if (newIt.second)
- newIt.first->second = std::move(oldIt->second);
- else
- newIt.first->second.set_union(oldIt->second);
- symbolToUsers.erase(oldIt);
- }
- }
-
-private:
- void collectFromOp(mlir::Operation *op, mlir::Operation *tableOp,
- llvm::SmallVectorImpl<mlir::Operation *> &symbols) {
- // Like mlir::walkSymbolUses: an unregistered op with a region might be
- // a symbol table, whose uses belong to another scope. CIRGen output is
- // fully registered, so this only guards the fork.
- assert((op->getNumRegions() != 1 || op->getDialect()) &&
- "unregistered op with a region");
- currentOp = op;
- visited.clear();
- auto visit = [&](mlir::Attribute attr) { visitAttr(attr); };
- visit(op->getRawDictionaryAttrs());
- op->getName().walkInherentAttrs(
- op, [&](llvm::StringRef, mlir::Attribute &attr) { visit(attr); });
- for (auto [userOp, ref] : pendingUses) {
- symbols.clear();
- (void)symbolTables.lookupSymbolIn(tableOp, ref, symbols);
- for (mlir::Operation *symbolOp : symbols)
- symbolToUsers[symbolOp].insert(userOp);
- }
- pendingUses.clear();
-
- // Recurse like mlir::walkSymbolUses: a symbol table defines a new scope.
- if (op->hasTrait<mlir::OpTrait::SymbolTable>())
- return;
- for (mlir::Region ®ion : op->getRegions())
- for (mlir::Operation &nested : region.getOps())
- collectFromOp(&nested, tableOp, symbols);
- }
-
- void visitAttr(mlir::Attribute attr) {
- // Visit each sub-element once, like the walker MLIR uses.
- if (!visited.insert(attr.getAsOpaquePointer()).second)
- return;
- if (auto ref = mlir::dyn_cast<mlir::SymbolRefAttr>(attr)) {
- pendingUses.emplace_back(currentOp, ref);
- // Don't walk nested references.
- return;
- }
- attr.walkImmediateSubElements([&](mlir::Attribute sub) { visitAttr(sub); },
- [&](mlir::Type type) { visitType(type); });
- }
-
- void visitType(mlir::Type type) {
- // Prune: audited above, CIR types hold no symbol references.
- if (cirDialect && &type.getDialect() == cirDialect) {
-#ifndef NDEBUG
- bool hasSymbolRef = false;
- type.walk([&](mlir::SymbolRefAttr) { hasSymbolRef = true; });
- assert(!hasSymbolRef && "CIR type contains a symbol reference");
-#endif
- return;
- }
- if (!visited.insert(type.getAsOpaquePointer()).second)
- return;
- type.walkImmediateSubElements([&](mlir::Attribute sub) { visitAttr(sub); },
- [&](mlir::Type sub) { visitType(sub); });
- }
-
- mlir::SymbolTableCollection &symbolTables;
- cir::CIRDialect *cirDialect = nullptr;
- mlir::Operation *currentOp = nullptr;
- llvm::SmallPtrSet<const void *, 32> visited;
- llvm::SmallVector<std::pair<mlir::Operation *, mlir::SymbolRefAttr>>
- pendingUses;
- llvm::DenseMap<mlir::Operation *, llvm::SetVector<mlir::Operation *>>
- symbolToUsers;
-};
-} // namespace
-
#ifndef NDEBUG
-static bool verifyPointerTypeArgs(cir::FuncOp oldF, cir::FuncOp newF,
- CIRSymbolUserMap &userMap) {
- for (mlir::Operation *user : userMap.getUsers(oldF)) {
+static bool verifyPointerTypeArgs(llvm::ArrayRef<mlir::Operation *> users,
+ cir::FuncOp newF) {
+ for (mlir::Operation *user : users) {
auto call = mlir::dyn_cast<cir::CallOp>(user);
if (!call)
continue;
@@ -2029,14 +1897,38 @@ void CIRGenModule::applyReplacements() {
if (replacements.empty())
return;
- // Build a symbol user map once — this walks the module O(M) one time.
- // Previously, each replaceAllSymbolUses call walked the entire module,
- // giving O(R × M) quadratic behavior for R replacements.
- // CIRSymbolUserMap additionally prunes CIR type subtrees (see above):
- // MLIR's generic walk descends into every nested record of every function
- // signature even though no symbol reference can hide there.
- mlir::SymbolTableCollection symbolTableCollection;
- CIRSymbolUserMap userMap(symbolTableCollection, getModule());
+ // Collect the users of the functions being replaced in one walk. Skip CIR
+ // types. None can hold a symbol reference (see CIRTypes.td), and walking
+ // them descends into every nested struct of every signature.
+ llvm::DenseMap<mlir::StringAttr, llvm::SetVector<mlir::Operation *>> users;
+ for (auto &i : replacements)
+ users.try_emplace(mlir::StringAttr::get(&getMLIRContext(), i.first));
+
+ mlir::ModuleOp module = getModule();
+ module->walk<mlir::WalkOrder::PreOrder>([&](mlir::Operation *op) {
+ auto collect = [&](mlir::Attribute attr) {
+ attr.walk<mlir::WalkOrder::PreOrder>(
+ [&](mlir::SymbolRefAttr ref) {
+ auto it = users.find(ref.getRootReference());
+ if (it != users.end())
+ it->second.insert(op);
+ return mlir::WalkResult::skip();
+ },
+ [&](mlir::Type type) {
+ return mlir::isa<cir::CIRDialect>(type.getDialect())
+ ? mlir::WalkResult::skip()
+ : mlir::WalkResult::advance();
+ });
+ };
+ collect(op->getRawDictionaryAttrs());
+ op->getName().walkInherentAttrs(
+ op, [&](llvm::StringRef, mlir::Attribute &attr) { collect(attr); });
+ // Like mlir::SymbolUserMap, don't look into nested symbol tables.
+ if (op != module.getOperation() &&
+ op->hasTrait<mlir::OpTrait::SymbolTable>())
+ return mlir::WalkResult::skip();
+ return mlir::WalkResult::advance();
+ });
for (auto &i : replacements) {
StringRef mangledName = i.first;
@@ -2053,12 +1945,21 @@ void CIRGenModule::applyReplacements() {
continue;
}
- assert(verifyPointerTypeArgs(oldF, newF, userMap) &&
+ auto oldIt = users.find(oldF.getSymNameAttr());
+ assert(oldIt != users.end() && "replaced function was not collected");
+ llvm::ArrayRef<mlir::Operation *> oldUsers = oldIt->second.getArrayRef();
+ assert(verifyPointerTypeArgs(oldUsers, newF) &&
"call argument types do not match replacement function");
- // Replace old with new, but keep the old order. Uses
- // SymbolUserMap to touch only actual users, not the whole module.
- userMap.replaceAllUsesWith(oldF, newF.getSymNameAttr());
+ // Replace old with new, but keep the old order.
+ for (mlir::Operation *user : oldUsers)
+ (void)mlir::SymbolTable::replaceAllSymbolUses(oldF, newF.getSymNameAttr(),
+ user);
+ // These users now name newF. Keep them in case newF is replaced later.
+ auto newIt = users.find(newF.getSymNameAttr());
+ if (newIt != users.end())
+ newIt->second.set_union(oldIt->second);
+
newF->moveBefore(oldF);
eraseGlobalSymbol(oldF);
oldF->erase();
More information about the cfe-commits
mailing list