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 <[email protected]> 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 <[email protected]> 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(); _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
