1135 lines · cpp
1//===- SymbolTable.cpp - MLIR Symbol Table Class --------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/IR/SymbolTable.h"10#include "mlir/IR/Builders.h"11#include "mlir/IR/OpImplementation.h"12#include "llvm/ADT/SetVector.h"13#include "llvm/ADT/SmallString.h"14#include "llvm/ADT/StringSwitch.h"15#include <optional>16 17using namespace mlir;18 19/// Return true if the given operation is unknown and may potentially define a20/// symbol table.21static bool isPotentiallyUnknownSymbolTable(Operation *op) {22 return op->getNumRegions() == 1 && !op->getDialect();23}24 25/// Returns the string name of the given symbol, or null if this is not a26/// symbol.27static StringAttr getNameIfSymbol(Operation *op) {28 return op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());29}30static StringAttr getNameIfSymbol(Operation *op, StringAttr symbolAttrNameId) {31 return op->getAttrOfType<StringAttr>(symbolAttrNameId);32}33 34/// Computes the nested symbol reference attribute for the symbol 'symbolName'35/// that are usable within the symbol table operations from 'symbol' as far up36/// to the given operation 'within', where 'within' is an ancestor of 'symbol'.37/// Returns success if all references up to 'within' could be computed.38static LogicalResult39collectValidReferencesFor(Operation *symbol, StringAttr symbolName,40 Operation *within,41 SmallVectorImpl<SymbolRefAttr> &results) {42 assert(within->isAncestor(symbol) && "expected 'within' to be an ancestor");43 MLIRContext *ctx = symbol->getContext();44 45 auto leafRef = FlatSymbolRefAttr::get(symbolName);46 results.push_back(leafRef);47 48 // Early exit for when 'within' is the parent of 'symbol'.49 Operation *symbolTableOp = symbol->getParentOp();50 if (within == symbolTableOp)51 return success();52 53 // Collect references until 'symbolTableOp' reaches 'within'.54 SmallVector<FlatSymbolRefAttr, 1> nestedRefs(1, leafRef);55 StringAttr symbolNameId =56 StringAttr::get(ctx, SymbolTable::getSymbolAttrName());57 do {58 // Each parent of 'symbol' should define a symbol table.59 if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())60 return failure();61 // Each parent of 'symbol' should also be a symbol.62 StringAttr symbolTableName = getNameIfSymbol(symbolTableOp, symbolNameId);63 if (!symbolTableName)64 return failure();65 results.push_back(SymbolRefAttr::get(symbolTableName, nestedRefs));66 67 symbolTableOp = symbolTableOp->getParentOp();68 if (symbolTableOp == within)69 break;70 nestedRefs.insert(nestedRefs.begin(),71 FlatSymbolRefAttr::get(symbolTableName));72 } while (true);73 return success();74}75 76/// Walk all of the operations within the given set of regions, without77/// traversing into any nested symbol tables. Stops walking if the result of the78/// callback is anything other than `WalkResult::advance`.79static std::optional<WalkResult>80walkSymbolTable(MutableArrayRef<Region> regions,81 function_ref<std::optional<WalkResult>(Operation *)> callback) {82 SmallVector<Region *, 1> worklist(llvm::make_pointer_range(regions));83 while (!worklist.empty()) {84 for (Operation &op : worklist.pop_back_val()->getOps()) {85 std::optional<WalkResult> result = callback(&op);86 if (result != WalkResult::advance())87 return result;88 89 // If this op defines a new symbol table scope, we can't traverse. Any90 // symbol references nested within 'op' are different semantically.91 if (!op.hasTrait<OpTrait::SymbolTable>()) {92 for (Region ®ion : op.getRegions())93 worklist.push_back(®ion);94 }95 }96 }97 return WalkResult::advance();98}99 100/// Walk all of the operations nested under, and including, the given operation,101/// without traversing into any nested symbol tables. Stops walking if the102/// result of the callback is anything other than `WalkResult::advance`.103static std::optional<WalkResult>104walkSymbolTable(Operation *op,105 function_ref<std::optional<WalkResult>(Operation *)> callback) {106 std::optional<WalkResult> result = callback(op);107 if (result != WalkResult::advance() || op->hasTrait<OpTrait::SymbolTable>())108 return result;109 return walkSymbolTable(op->getRegions(), callback);110}111 112//===----------------------------------------------------------------------===//113// SymbolTable114//===----------------------------------------------------------------------===//115 116/// Build a symbol table with the symbols within the given operation.117SymbolTable::SymbolTable(Operation *symbolTableOp)118 : symbolTableOp(symbolTableOp) {119 assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>() &&120 "expected operation to have SymbolTable trait");121 assert(symbolTableOp->getNumRegions() == 1 &&122 "expected operation to have a single region");123 assert(symbolTableOp->getRegion(0).hasOneBlock() &&124 "expected operation to have a single block");125 126 StringAttr symbolNameId = StringAttr::get(symbolTableOp->getContext(),127 SymbolTable::getSymbolAttrName());128 for (auto &op : symbolTableOp->getRegion(0).front()) {129 StringAttr name = getNameIfSymbol(&op, symbolNameId);130 if (!name)131 continue;132 133 auto inserted = symbolTable.insert({name, &op});134 (void)inserted;135 assert(inserted.second &&136 "expected region to contain uniquely named symbol operations");137 }138}139 140/// Look up a symbol with the specified name, returning null if no such name141/// exists. Names never include the @ on them.142Operation *SymbolTable::lookup(StringRef name) const {143 return lookup(StringAttr::get(symbolTableOp->getContext(), name));144}145Operation *SymbolTable::lookup(StringAttr name) const {146 return symbolTable.lookup(name);147}148 149void SymbolTable::remove(Operation *op) {150 StringAttr name = getNameIfSymbol(op);151 assert(name && "expected valid 'name' attribute");152 assert(op->getParentOp() == symbolTableOp &&153 "expected this operation to be inside of the operation with this "154 "SymbolTable");155 156 auto it = symbolTable.find(name);157 if (it != symbolTable.end() && it->second == op)158 symbolTable.erase(it);159}160 161void SymbolTable::erase(Operation *symbol) {162 remove(symbol);163 symbol->erase();164}165 166// TODO: Consider if this should be renamed to something like insertOrUpdate167/// Insert a new symbol into the table and associated operation if not already168/// there and rename it as necessary to avoid collisions. Return the name of169/// the symbol after insertion as attribute.170StringAttr SymbolTable::insert(Operation *symbol, Block::iterator insertPt) {171 // The symbol cannot be the child of another op and must be the child of the172 // symbolTableOp after this.173 //174 // TODO: consider if SymbolTable's constructor should behave the same.175 if (!symbol->getParentOp()) {176 auto &body = symbolTableOp->getRegion(0).front();177 if (insertPt == Block::iterator()) {178 insertPt = Block::iterator(body.end());179 } else {180 assert((insertPt == body.end() ||181 insertPt->getParentOp() == symbolTableOp) &&182 "expected insertPt to be in the associated module operation");183 }184 // Insert before the terminator, if any.185 if (insertPt == Block::iterator(body.end()) && !body.empty() &&186 std::prev(body.end())->hasTrait<OpTrait::IsTerminator>())187 insertPt = std::prev(body.end());188 189 body.getOperations().insert(insertPt, symbol);190 }191 assert(symbol->getParentOp() == symbolTableOp &&192 "symbol is already inserted in another op");193 194 // Add this symbol to the symbol table, uniquing the name if a conflict is195 // detected.196 StringAttr name = getSymbolName(symbol);197 if (symbolTable.insert({name, symbol}).second)198 return name;199 // If the symbol was already in the table, also return.200 if (symbolTable.lookup(name) == symbol)201 return name;202 203 MLIRContext *context = symbol->getContext();204 SmallString<128> nameBuffer = generateSymbolName<128>(205 name.getValue(),206 [&](StringRef candidate) {207 return !symbolTable208 .insert({StringAttr::get(context, candidate), symbol})209 .second;210 },211 uniquingCounter);212 setSymbolName(symbol, nameBuffer);213 return getSymbolName(symbol);214}215 216LogicalResult SymbolTable::rename(StringAttr from, StringAttr to) {217 Operation *op = lookup(from);218 return rename(op, to);219}220 221LogicalResult SymbolTable::rename(Operation *op, StringAttr to) {222 StringAttr from = getNameIfSymbol(op);223 (void)from;224 225 assert(from && "expected valid 'name' attribute");226 assert(op->getParentOp() == symbolTableOp &&227 "expected this operation to be inside of the operation with this "228 "SymbolTable");229 assert(lookup(from) == op && "current name does not resolve to op");230 assert(lookup(to) == nullptr && "new name already exists");231 232 if (failed(SymbolTable::replaceAllSymbolUses(op, to, getOp())))233 return failure();234 235 // Remove op with old name, change name, add with new name. The order is236 // important here due to how `remove` and `insert` rely on the op name.237 remove(op);238 setSymbolName(op, to);239 insert(op);240 241 assert(lookup(to) == op && "new name does not resolve to renamed op");242 assert(lookup(from) == nullptr && "old name still exists");243 244 return success();245}246 247LogicalResult SymbolTable::rename(StringAttr from, StringRef to) {248 auto toAttr = StringAttr::get(getOp()->getContext(), to);249 return rename(from, toAttr);250}251 252LogicalResult SymbolTable::rename(Operation *op, StringRef to) {253 auto toAttr = StringAttr::get(getOp()->getContext(), to);254 return rename(op, toAttr);255}256 257FailureOr<StringAttr>258SymbolTable::renameToUnique(StringAttr oldName,259 ArrayRef<SymbolTable *> others) {260 261 // Determine new name that is unique in all symbol tables.262 StringAttr newName;263 {264 MLIRContext *context = oldName.getContext();265 SmallString<64> prefix = oldName.getValue();266 int uniqueId = 0;267 prefix.push_back('_');268 while (true) {269 newName = StringAttr::get(context, prefix + Twine(uniqueId++));270 auto lookupNewName = [&](SymbolTable *st) { return st->lookup(newName); };271 if (!lookupNewName(this) && llvm::none_of(others, lookupNewName)) {272 break;273 }274 }275 }276 277 // Apply renaming.278 if (failed(rename(oldName, newName)))279 return failure();280 return newName;281}282 283FailureOr<StringAttr>284SymbolTable::renameToUnique(Operation *op, ArrayRef<SymbolTable *> others) {285 StringAttr from = getNameIfSymbol(op);286 assert(from && "expected valid 'name' attribute");287 return renameToUnique(from, others);288}289 290/// Returns the name of the given symbol operation.291StringAttr SymbolTable::getSymbolName(Operation *symbol) {292 StringAttr name = getNameIfSymbol(symbol);293 assert(name && "expected valid symbol name");294 return name;295}296 297/// Sets the name of the given symbol operation.298void SymbolTable::setSymbolName(Operation *symbol, StringAttr name) {299 symbol->setAttr(getSymbolAttrName(), name);300}301 302/// Returns the visibility of the given symbol operation.303SymbolTable::Visibility SymbolTable::getSymbolVisibility(Operation *symbol) {304 // If the attribute doesn't exist, assume public.305 StringAttr vis = symbol->getAttrOfType<StringAttr>(getVisibilityAttrName());306 if (!vis)307 return Visibility::Public;308 309 // Otherwise, switch on the string value.310 return StringSwitch<Visibility>(vis.getValue())311 .Case("private", Visibility::Private)312 .Case("nested", Visibility::Nested)313 .Case("public", Visibility::Public);314}315/// Sets the visibility of the given symbol operation.316void SymbolTable::setSymbolVisibility(Operation *symbol, Visibility vis) {317 MLIRContext *ctx = symbol->getContext();318 319 // If the visibility is public, just drop the attribute as this is the320 // default.321 if (vis == Visibility::Public) {322 symbol->removeAttr(StringAttr::get(ctx, getVisibilityAttrName()));323 return;324 }325 326 // Otherwise, update the attribute.327 assert((vis == Visibility::Private || vis == Visibility::Nested) &&328 "unknown symbol visibility kind");329 330 StringRef visName = vis == Visibility::Private ? "private" : "nested";331 symbol->setAttr(getVisibilityAttrName(), StringAttr::get(ctx, visName));332}333 334/// Returns the nearest symbol table from a given operation `from`. Returns335/// nullptr if no valid parent symbol table could be found.336Operation *SymbolTable::getNearestSymbolTable(Operation *from) {337 assert(from && "expected valid operation");338 if (isPotentiallyUnknownSymbolTable(from))339 return nullptr;340 341 while (!from->hasTrait<OpTrait::SymbolTable>()) {342 from = from->getParentOp();343 344 // Check that this is a valid op and isn't an unknown symbol table.345 if (!from || isPotentiallyUnknownSymbolTable(from))346 return nullptr;347 }348 return from;349}350 351/// Walks all symbol table operations nested within, and including, `op`. For352/// each symbol table operation, the provided callback is invoked with the op353/// and a boolean signifying if the symbols within that symbol table can be354/// treated as if all uses are visible. `allSymUsesVisible` identifies whether355/// all of the symbol uses of symbols within `op` are visible.356void SymbolTable::walkSymbolTables(357 Operation *op, bool allSymUsesVisible,358 function_ref<void(Operation *, bool)> callback) {359 bool isSymbolTable = op->hasTrait<OpTrait::SymbolTable>();360 if (isSymbolTable) {361 SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);362 allSymUsesVisible |= !symbol || symbol.isPrivate();363 } else {364 // Otherwise if 'op' is not a symbol table, any nested symbols are365 // guaranteed to be hidden.366 allSymUsesVisible = true;367 }368 369 for (Region ®ion : op->getRegions())370 for (Block &block : region)371 for (Operation &nestedOp : block)372 walkSymbolTables(&nestedOp, allSymUsesVisible, callback);373 374 // If 'op' had the symbol table trait, visit it after any nested symbol375 // tables.376 if (isSymbolTable)377 callback(op, allSymUsesVisible);378}379 380/// Returns the operation registered with the given symbol name with the381/// regions of 'symbolTableOp'. 'symbolTableOp' is required to be an operation382/// with the 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol383/// was found.384Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,385 StringAttr symbol) {386 assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());387 Region ®ion = symbolTableOp->getRegion(0);388 if (region.empty())389 return nullptr;390 391 // Look for a symbol with the given name.392 StringAttr symbolNameId = StringAttr::get(symbolTableOp->getContext(),393 SymbolTable::getSymbolAttrName());394 for (auto &op : region.front())395 if (getNameIfSymbol(&op, symbolNameId) == symbol)396 return &op;397 return nullptr;398}399Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,400 SymbolRefAttr symbol) {401 SmallVector<Operation *, 4> resolvedSymbols;402 if (failed(lookupSymbolIn(symbolTableOp, symbol, resolvedSymbols)))403 return nullptr;404 return resolvedSymbols.back();405}406 407/// Internal implementation of `lookupSymbolIn` that allows for specialized408/// implementations of the lookup function.409static LogicalResult lookupSymbolInImpl(410 Operation *symbolTableOp, SymbolRefAttr symbol,411 SmallVectorImpl<Operation *> &symbols,412 function_ref<Operation *(Operation *, StringAttr)> lookupSymbolFn) {413 assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());414 415 // Lookup the root reference for this symbol.416 symbolTableOp = lookupSymbolFn(symbolTableOp, symbol.getRootReference());417 if (!symbolTableOp)418 return failure();419 symbols.push_back(symbolTableOp);420 421 // If there are no nested references, just return the root symbol directly.422 ArrayRef<FlatSymbolRefAttr> nestedRefs = symbol.getNestedReferences();423 if (nestedRefs.empty())424 return success();425 426 // Verify that the root is also a symbol table.427 if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())428 return failure();429 430 // Otherwise, lookup each of the nested non-leaf references and ensure that431 // each corresponds to a valid symbol table.432 for (FlatSymbolRefAttr ref : nestedRefs.drop_back()) {433 symbolTableOp = lookupSymbolFn(symbolTableOp, ref.getAttr());434 if (!symbolTableOp || !symbolTableOp->hasTrait<OpTrait::SymbolTable>())435 return failure();436 symbols.push_back(symbolTableOp);437 }438 symbols.push_back(lookupSymbolFn(symbolTableOp, symbol.getLeafReference()));439 return success(symbols.back());440}441 442LogicalResult443SymbolTable::lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr symbol,444 SmallVectorImpl<Operation *> &symbols) {445 auto lookupFn = [](Operation *symbolTableOp, StringAttr symbol) {446 return lookupSymbolIn(symbolTableOp, symbol);447 };448 return lookupSymbolInImpl(symbolTableOp, symbol, symbols, lookupFn);449}450 451/// Returns the operation registered with the given symbol name within the452/// closes parent operation with the 'OpTrait::SymbolTable' trait. Returns453/// nullptr if no valid symbol was found.454Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,455 StringAttr symbol) {456 Operation *symbolTableOp = getNearestSymbolTable(from);457 return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;458}459Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,460 SymbolRefAttr symbol) {461 Operation *symbolTableOp = getNearestSymbolTable(from);462 return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;463}464 465raw_ostream &mlir::operator<<(raw_ostream &os,466 SymbolTable::Visibility visibility) {467 switch (visibility) {468 case SymbolTable::Visibility::Public:469 return os << "public";470 case SymbolTable::Visibility::Private:471 return os << "private";472 case SymbolTable::Visibility::Nested:473 return os << "nested";474 }475 llvm_unreachable("Unexpected visibility");476}477 478//===----------------------------------------------------------------------===//479// SymbolTable Trait Types480//===----------------------------------------------------------------------===//481 482LogicalResult detail::verifySymbolTable(Operation *op) {483 if (op->getNumRegions() != 1)484 return op->emitOpError()485 << "Operations with a 'SymbolTable' must have exactly one region";486 if (!op->getRegion(0).hasOneBlock())487 return op->emitOpError()488 << "Operations with a 'SymbolTable' must have exactly one block";489 490 // Check that all symbols are uniquely named within child regions.491 DenseMap<Attribute, Location> nameToOrigLoc;492 for (auto &block : op->getRegion(0)) {493 for (auto &op : block) {494 // Check for a symbol name attribute.495 auto nameAttr =496 op.getAttrOfType<StringAttr>(mlir::SymbolTable::getSymbolAttrName());497 if (!nameAttr)498 continue;499 500 // Try to insert this symbol into the table.501 auto it = nameToOrigLoc.try_emplace(nameAttr, op.getLoc());502 if (!it.second)503 return op.emitError()504 .append("redefinition of symbol named '", nameAttr.getValue(), "'")505 .attachNote(it.first->second)506 .append("see existing symbol definition here");507 }508 }509 510 // Verify any nested symbol user operations.511 SymbolTableCollection symbolTable;512 auto verifySymbolUserFn = [&](Operation *op) -> std::optional<WalkResult> {513 if (SymbolUserOpInterface user = dyn_cast<SymbolUserOpInterface>(op))514 return WalkResult(user.verifySymbolUses(symbolTable));515 return WalkResult::advance();516 };517 518 std::optional<WalkResult> result =519 walkSymbolTable(op->getRegions(), verifySymbolUserFn);520 return success(result && !result->wasInterrupted());521}522 523LogicalResult detail::verifySymbol(Operation *op) {524 // Verify the name attribute.525 if (!op->getAttrOfType<StringAttr>(mlir::SymbolTable::getSymbolAttrName()))526 return op->emitOpError() << "requires string attribute '"527 << mlir::SymbolTable::getSymbolAttrName() << "'";528 529 // Verify the visibility attribute.530 if (Attribute vis = op->getAttr(mlir::SymbolTable::getVisibilityAttrName())) {531 StringAttr visStrAttr = llvm::dyn_cast<StringAttr>(vis);532 if (!visStrAttr)533 return op->emitOpError() << "requires visibility attribute '"534 << mlir::SymbolTable::getVisibilityAttrName()535 << "' to be a string attribute, but got " << vis;536 537 if (!llvm::is_contained(ArrayRef<StringRef>{"public", "private", "nested"},538 visStrAttr.getValue()))539 return op->emitOpError()540 << "visibility expected to be one of [\"public\", \"private\", "541 "\"nested\"], but got "542 << visStrAttr;543 }544 return success();545}546 547//===----------------------------------------------------------------------===//548// Symbol Use Lists549//===----------------------------------------------------------------------===//550 551/// Walk all of the symbol references within the given operation, invoking the552/// provided callback for each found use. The callbacks takes the use of the553/// symbol.554static WalkResult555walkSymbolRefs(Operation *op,556 function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {557 return op->getAttrDictionary().walk<WalkOrder::PreOrder>(558 [&](SymbolRefAttr symbolRef) {559 if (callback({op, symbolRef}).wasInterrupted())560 return WalkResult::interrupt();561 562 // Don't walk nested references.563 return WalkResult::skip();564 });565}566 567/// Walk all of the uses, for any symbol, that are nested within the given568/// regions, invoking the provided callback for each. This does not traverse569/// into any nested symbol tables.570static std::optional<WalkResult>571walkSymbolUses(MutableArrayRef<Region> regions,572 function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {573 return walkSymbolTable(regions,574 [&](Operation *op) -> std::optional<WalkResult> {575 // Check that this isn't a potentially unknown symbol576 // table.577 if (isPotentiallyUnknownSymbolTable(op))578 return std::nullopt;579 580 return walkSymbolRefs(op, callback);581 });582}583/// Walk all of the uses, for any symbol, that are nested within the given584/// operation 'from', invoking the provided callback for each. This does not585/// traverse into any nested symbol tables.586static std::optional<WalkResult>587walkSymbolUses(Operation *from,588 function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {589 // If this operation has regions, and it, as well as its dialect, isn't590 // registered then conservatively fail. The operation may define a591 // symbol table, so we can't opaquely know if we should traverse to find592 // nested uses.593 if (isPotentiallyUnknownSymbolTable(from))594 return std::nullopt;595 596 // Walk the uses on this operation.597 if (walkSymbolRefs(from, callback).wasInterrupted())598 return WalkResult::interrupt();599 600 // Only recurse if this operation is not a symbol table. A symbol table601 // defines a new scope, so we can't walk the attributes from within the symbol602 // table op.603 if (!from->hasTrait<OpTrait::SymbolTable>())604 return walkSymbolUses(from->getRegions(), callback);605 return WalkResult::advance();606}607 608namespace {609/// This class represents a single symbol scope. A symbol scope represents the610/// set of operations nested within a symbol table that may reference symbols611/// within that table. A symbol scope does not contain the symbol table612/// operation itself, just its contained operations. A scope ends at leaf613/// operations or another symbol table operation.614struct SymbolScope {615 /// Walk the symbol uses within this scope, invoking the given callback.616 /// This variant is used when the callback type matches that expected by617 /// 'walkSymbolUses'.618 template <typename CallbackT,619 std::enable_if_t<!std::is_same<620 typename llvm::function_traits<CallbackT>::result_t,621 void>::value> * = nullptr>622 std::optional<WalkResult> walk(CallbackT cback) {623 if (Region *region = llvm::dyn_cast_if_present<Region *>(limit))624 return walkSymbolUses(*region, cback);625 return walkSymbolUses(cast<Operation *>(limit), cback);626 }627 /// This variant is used when the callback type matches a stripped down type:628 /// void(SymbolTable::SymbolUse use)629 template <typename CallbackT,630 std::enable_if_t<std::is_same<631 typename llvm::function_traits<CallbackT>::result_t,632 void>::value> * = nullptr>633 std::optional<WalkResult> walk(CallbackT cback) {634 return walk([=](SymbolTable::SymbolUse use) {635 return cback(use), WalkResult::advance();636 });637 }638 639 /// Walk all of the operations nested under the current scope without640 /// traversing into any nested symbol tables.641 template <typename CallbackT>642 std::optional<WalkResult> walkSymbolTable(CallbackT &&cback) {643 if (Region *region = llvm::dyn_cast_if_present<Region *>(limit))644 return ::walkSymbolTable(*region, cback);645 return ::walkSymbolTable(cast<Operation *>(limit), cback);646 }647 648 /// The representation of the symbol within this scope.649 SymbolRefAttr symbol;650 651 /// The IR unit representing this scope.652 llvm::PointerUnion<Operation *, Region *> limit;653};654} // namespace655 656/// Collect all of the symbol scopes from 'symbol' to (inclusive) 'limit'.657static SmallVector<SymbolScope, 2> collectSymbolScopes(Operation *symbol,658 Operation *limit) {659 StringAttr symName = SymbolTable::getSymbolName(symbol);660 assert(!symbol->hasTrait<OpTrait::SymbolTable>() || symbol != limit);661 662 // Compute the ancestors of 'limit'.663 SetVector<Operation *, SmallVector<Operation *, 4>,664 SmallPtrSet<Operation *, 4>>665 limitAncestors;666 Operation *limitAncestor = limit;667 do {668 // Check to see if 'symbol' is an ancestor of 'limit'.669 if (limitAncestor == symbol) {670 // Check that the nearest symbol table is 'symbol's parent. SymbolRefAttr671 // doesn't support parent references.672 if (SymbolTable::getNearestSymbolTable(limit->getParentOp()) ==673 symbol->getParentOp())674 return {{SymbolRefAttr::get(symName), limit}};675 return {};676 }677 678 limitAncestors.insert(limitAncestor);679 } while ((limitAncestor = limitAncestor->getParentOp()));680 681 // Try to find the first ancestor of 'symbol' that is an ancestor of 'limit'.682 Operation *commonAncestor = symbol->getParentOp();683 do {684 if (limitAncestors.count(commonAncestor))685 break;686 } while ((commonAncestor = commonAncestor->getParentOp()));687 assert(commonAncestor && "'limit' and 'symbol' have no common ancestor");688 689 // Compute the set of valid nested references for 'symbol' as far up to the690 // common ancestor as possible.691 SmallVector<SymbolRefAttr, 2> references;692 bool collectedAllReferences = succeeded(693 collectValidReferencesFor(symbol, symName, commonAncestor, references));694 695 // Handle the case where the common ancestor is 'limit'.696 if (commonAncestor == limit) {697 SmallVector<SymbolScope, 2> scopes;698 699 // Walk each of the ancestors of 'symbol', calling the compute function for700 // each one.701 Operation *limitIt = symbol->getParentOp();702 for (size_t i = 0, e = references.size(); i != e;703 ++i, limitIt = limitIt->getParentOp()) {704 assert(limitIt->hasTrait<OpTrait::SymbolTable>());705 scopes.push_back({references[i], &limitIt->getRegion(0)});706 }707 return scopes;708 }709 710 // Otherwise, we just need the symbol reference for 'symbol' that will be711 // used within 'limit'. This is the last reference in the list we computed712 // above if we were able to collect all references.713 if (!collectedAllReferences)714 return {};715 return {{references.back(), limit}};716}717static SmallVector<SymbolScope, 2> collectSymbolScopes(Operation *symbol,718 Region *limit) {719 auto scopes = collectSymbolScopes(symbol, limit->getParentOp());720 721 // If we collected some scopes to walk, make sure to constrain the one for722 // limit to the specific region requested.723 if (!scopes.empty())724 scopes.back().limit = limit;725 return scopes;726}727static SmallVector<SymbolScope, 1> collectSymbolScopes(StringAttr symbol,728 Region *limit) {729 return {{SymbolRefAttr::get(symbol), limit}};730}731 732static SmallVector<SymbolScope, 1> collectSymbolScopes(StringAttr symbol,733 Operation *limit) {734 SmallVector<SymbolScope, 1> scopes;735 auto symbolRef = SymbolRefAttr::get(symbol);736 for (auto ®ion : limit->getRegions())737 scopes.push_back({symbolRef, ®ion});738 return scopes;739}740 741/// Returns true if the given reference 'SubRef' is a sub reference of the742/// reference 'ref', i.e. 'ref' is a further qualified reference.743static bool isReferencePrefixOf(SymbolRefAttr subRef, SymbolRefAttr ref) {744 if (ref == subRef)745 return true;746 747 // If the references are not pointer equal, check to see if `subRef` is a748 // prefix of `ref`.749 if (llvm::isa<FlatSymbolRefAttr>(ref) ||750 ref.getRootReference() != subRef.getRootReference())751 return false;752 753 auto refLeafs = ref.getNestedReferences();754 auto subRefLeafs = subRef.getNestedReferences();755 return subRefLeafs.size() < refLeafs.size() &&756 subRefLeafs == refLeafs.take_front(subRefLeafs.size());757}758 759//===----------------------------------------------------------------------===//760// SymbolTable::getSymbolUses761//===----------------------------------------------------------------------===//762 763/// The implementation of SymbolTable::getSymbolUses below.764template <typename FromT>765static std::optional<SymbolTable::UseRange> getSymbolUsesImpl(FromT from) {766 std::vector<SymbolTable::SymbolUse> uses;767 auto walkFn = [&](SymbolTable::SymbolUse symbolUse) {768 uses.push_back(symbolUse);769 return WalkResult::advance();770 };771 auto result = walkSymbolUses(from, walkFn);772 return result ? std::optional<SymbolTable::UseRange>(std::move(uses))773 : std::nullopt;774}775 776/// Get an iterator range for all of the uses, for any symbol, that are nested777/// within the given operation 'from'. This does not traverse into any nested778/// symbol tables, and will also only return uses on 'from' if it does not779/// also define a symbol table. This is because we treat the region as the780/// boundary of the symbol table, and not the op itself. This function returns781/// std::nullopt if there are any unknown operations that may potentially be782/// symbol tables.783auto SymbolTable::getSymbolUses(Operation *from) -> std::optional<UseRange> {784 return getSymbolUsesImpl(from);785}786auto SymbolTable::getSymbolUses(Region *from) -> std::optional<UseRange> {787 return getSymbolUsesImpl(MutableArrayRef<Region>(*from));788}789 790//===----------------------------------------------------------------------===//791// SymbolTable::getSymbolUses792//===----------------------------------------------------------------------===//793 794/// The implementation of SymbolTable::getSymbolUses below.795template <typename SymbolT, typename IRUnitT>796static std::optional<SymbolTable::UseRange> getSymbolUsesImpl(SymbolT symbol,797 IRUnitT *limit) {798 std::vector<SymbolTable::SymbolUse> uses;799 for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {800 if (!scope.walk([&](SymbolTable::SymbolUse symbolUse) {801 if (isReferencePrefixOf(scope.symbol, symbolUse.getSymbolRef()))802 uses.push_back(symbolUse);803 }))804 return std::nullopt;805 }806 return SymbolTable::UseRange(std::move(uses));807}808 809/// Get all of the uses of the given symbol that are nested within the given810/// operation 'from'. This does not traverse into any nested symbol tables.811/// This function returns std::nullopt if there are any unknown operations that812/// may potentially be symbol tables.813auto SymbolTable::getSymbolUses(StringAttr symbol, Operation *from)814 -> std::optional<UseRange> {815 return getSymbolUsesImpl(symbol, from);816}817auto SymbolTable::getSymbolUses(Operation *symbol, Operation *from)818 -> std::optional<UseRange> {819 return getSymbolUsesImpl(symbol, from);820}821auto SymbolTable::getSymbolUses(StringAttr symbol, Region *from)822 -> std::optional<UseRange> {823 return getSymbolUsesImpl(symbol, from);824}825auto SymbolTable::getSymbolUses(Operation *symbol, Region *from)826 -> std::optional<UseRange> {827 return getSymbolUsesImpl(symbol, from);828}829 830//===----------------------------------------------------------------------===//831// SymbolTable::symbolKnownUseEmpty832//===----------------------------------------------------------------------===//833 834/// The implementation of SymbolTable::symbolKnownUseEmpty below.835template <typename SymbolT, typename IRUnitT>836static bool symbolKnownUseEmptyImpl(SymbolT symbol, IRUnitT *limit) {837 for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {838 // Walk all of the symbol uses looking for a reference to 'symbol'.839 if (scope.walk([&](SymbolTable::SymbolUse symbolUse) {840 return isReferencePrefixOf(scope.symbol, symbolUse.getSymbolRef())841 ? WalkResult::interrupt()842 : WalkResult::advance();843 }) != WalkResult::advance())844 return false;845 }846 return true;847}848 849/// Return if the given symbol is known to have no uses that are nested within850/// the given operation 'from'. This does not traverse into any nested symbol851/// tables. This function will also return false if there are any unknown852/// operations that may potentially be symbol tables.853bool SymbolTable::symbolKnownUseEmpty(StringAttr symbol, Operation *from) {854 return symbolKnownUseEmptyImpl(symbol, from);855}856bool SymbolTable::symbolKnownUseEmpty(Operation *symbol, Operation *from) {857 return symbolKnownUseEmptyImpl(symbol, from);858}859bool SymbolTable::symbolKnownUseEmpty(StringAttr symbol, Region *from) {860 return symbolKnownUseEmptyImpl(symbol, from);861}862bool SymbolTable::symbolKnownUseEmpty(Operation *symbol, Region *from) {863 return symbolKnownUseEmptyImpl(symbol, from);864}865 866//===----------------------------------------------------------------------===//867// SymbolTable::replaceAllSymbolUses868//===----------------------------------------------------------------------===//869 870/// Generates a new symbol reference attribute with a new leaf reference.871static SymbolRefAttr generateNewRefAttr(SymbolRefAttr oldAttr,872 FlatSymbolRefAttr newLeafAttr) {873 if (llvm::isa<FlatSymbolRefAttr>(oldAttr))874 return newLeafAttr;875 auto nestedRefs = llvm::to_vector<2>(oldAttr.getNestedReferences());876 nestedRefs.back() = newLeafAttr;877 return SymbolRefAttr::get(oldAttr.getRootReference(), nestedRefs);878}879 880/// The implementation of SymbolTable::replaceAllSymbolUses below.881template <typename SymbolT, typename IRUnitT>882static LogicalResult883replaceAllSymbolUsesImpl(SymbolT symbol, StringAttr newSymbol, IRUnitT *limit) {884 // Generate a new attribute to replace the given attribute.885 FlatSymbolRefAttr newLeafAttr = FlatSymbolRefAttr::get(newSymbol);886 for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {887 SymbolRefAttr oldAttr = scope.symbol;888 SymbolRefAttr newAttr = generateNewRefAttr(scope.symbol, newLeafAttr);889 AttrTypeReplacer replacer;890 replacer.addReplacement(891 [&](SymbolRefAttr attr) -> std::pair<Attribute, WalkResult> {892 // Regardless of the match, don't walk nested SymbolRefAttrs, we don't893 // want to accidentally replace an inner reference.894 if (attr == oldAttr)895 return {newAttr, WalkResult::skip()};896 // Handle prefix matches.897 if (isReferencePrefixOf(oldAttr, attr)) {898 auto oldNestedRefs = oldAttr.getNestedReferences();899 auto nestedRefs = attr.getNestedReferences();900 if (oldNestedRefs.empty())901 return {SymbolRefAttr::get(newSymbol, nestedRefs),902 WalkResult::skip()};903 904 auto newNestedRefs = llvm::to_vector<4>(nestedRefs);905 newNestedRefs[oldNestedRefs.size() - 1] = newLeafAttr;906 return {SymbolRefAttr::get(attr.getRootReference(), newNestedRefs),907 WalkResult::skip()};908 }909 return {attr, WalkResult::skip()};910 });911 912 auto walkFn = [&](Operation *op) -> std::optional<WalkResult> {913 replacer.replaceElementsIn(op);914 return WalkResult::advance();915 };916 if (!scope.walkSymbolTable(walkFn))917 return failure();918 }919 return success();920}921 922/// Attempt to replace all uses of the given symbol 'oldSymbol' with the923/// provided symbol 'newSymbol' that are nested within the given operation924/// 'from'. This does not traverse into any nested symbol tables. If there are925/// any unknown operations that may potentially be symbol tables, no uses are926/// replaced and failure is returned.927LogicalResult SymbolTable::replaceAllSymbolUses(StringAttr oldSymbol,928 StringAttr newSymbol,929 Operation *from) {930 return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);931}932LogicalResult SymbolTable::replaceAllSymbolUses(Operation *oldSymbol,933 StringAttr newSymbol,934 Operation *from) {935 return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);936}937LogicalResult SymbolTable::replaceAllSymbolUses(StringAttr oldSymbol,938 StringAttr newSymbol,939 Region *from) {940 return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);941}942LogicalResult SymbolTable::replaceAllSymbolUses(Operation *oldSymbol,943 StringAttr newSymbol,944 Region *from) {945 return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);946}947 948//===----------------------------------------------------------------------===//949// SymbolTableCollection950//===----------------------------------------------------------------------===//951 952Operation *SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,953 StringAttr symbol) {954 return getSymbolTable(symbolTableOp).lookup(symbol);955}956Operation *SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,957 SymbolRefAttr name) {958 SmallVector<Operation *, 4> symbols;959 if (failed(lookupSymbolIn(symbolTableOp, name, symbols)))960 return nullptr;961 return symbols.back();962}963/// A variant of 'lookupSymbolIn' that returns all of the symbols referenced by964/// a given SymbolRefAttr. Returns failure if any of the nested references could965/// not be resolved.966LogicalResult967SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,968 SymbolRefAttr name,969 SmallVectorImpl<Operation *> &symbols) {970 auto lookupFn = [this](Operation *symbolTableOp, StringAttr symbol) {971 return lookupSymbolIn(symbolTableOp, symbol);972 };973 return lookupSymbolInImpl(symbolTableOp, name, symbols, lookupFn);974}975 976/// Returns the operation registered with the given symbol name within the977/// closest parent operation of, or including, 'from' with the978/// 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol was979/// found.980Operation *SymbolTableCollection::lookupNearestSymbolFrom(Operation *from,981 StringAttr symbol) {982 Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(from);983 return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;984}985Operation *986SymbolTableCollection::lookupNearestSymbolFrom(Operation *from,987 SymbolRefAttr symbol) {988 Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(from);989 return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;990}991 992/// Lookup, or create, a symbol table for an operation.993SymbolTable &SymbolTableCollection::getSymbolTable(Operation *op) {994 auto it = symbolTables.try_emplace(op, nullptr);995 if (it.second)996 it.first->second = std::make_unique<SymbolTable>(op);997 return *it.first->second;998}999 1000void SymbolTableCollection::invalidateSymbolTable(Operation *op) {1001 symbolTables.erase(op);1002}1003 1004//===----------------------------------------------------------------------===//1005// LockedSymbolTableCollection1006//===----------------------------------------------------------------------===//1007 1008Operation *LockedSymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,1009 StringAttr symbol) {1010 return getSymbolTable(symbolTableOp).lookup(symbol);1011}1012 1013Operation *1014LockedSymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,1015 FlatSymbolRefAttr symbol) {1016 return lookupSymbolIn(symbolTableOp, symbol.getAttr());1017}1018 1019Operation *LockedSymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,1020 SymbolRefAttr name) {1021 SmallVector<Operation *> symbols;1022 if (failed(lookupSymbolIn(symbolTableOp, name, symbols)))1023 return nullptr;1024 return symbols.back();1025}1026 1027LogicalResult LockedSymbolTableCollection::lookupSymbolIn(1028 Operation *symbolTableOp, SymbolRefAttr name,1029 SmallVectorImpl<Operation *> &symbols) {1030 auto lookupFn = [this](Operation *symbolTableOp, StringAttr symbol) {1031 return lookupSymbolIn(symbolTableOp, symbol);1032 };1033 return lookupSymbolInImpl(symbolTableOp, name, symbols, lookupFn);1034}1035 1036SymbolTable &1037LockedSymbolTableCollection::getSymbolTable(Operation *symbolTableOp) {1038 assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());1039 // Try to find an existing symbol table.1040 {1041 llvm::sys::SmartScopedReader<true> lock(mutex);1042 auto it = collection.symbolTables.find(symbolTableOp);1043 if (it != collection.symbolTables.end())1044 return *it->second;1045 }1046 // Create a symbol table for the operation. Perform construction outside of1047 // the critical section.1048 auto symbolTable = std::make_unique<SymbolTable>(symbolTableOp);1049 // Insert the constructed symbol table.1050 llvm::sys::SmartScopedWriter<true> lock(mutex);1051 return *collection.symbolTables1052 .insert({symbolTableOp, std::move(symbolTable)})1053 .first->second;1054}1055 1056//===----------------------------------------------------------------------===//1057// SymbolUserMap1058//===----------------------------------------------------------------------===//1059 1060SymbolUserMap::SymbolUserMap(SymbolTableCollection &symbolTable,1061 Operation *symbolTableOp)1062 : symbolTable(symbolTable) {1063 // Walk each of the symbol tables looking for discardable callgraph nodes.1064 SmallVector<Operation *> symbols;1065 auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {1066 for (Operation &nestedOp : symbolTableOp->getRegion(0).getOps()) {1067 auto symbolUses = SymbolTable::getSymbolUses(&nestedOp);1068 assert(symbolUses && "expected uses to be valid");1069 1070 for (const SymbolTable::SymbolUse &use : *symbolUses) {1071 symbols.clear();1072 (void)symbolTable.lookupSymbolIn(symbolTableOp, use.getSymbolRef(),1073 symbols);1074 for (Operation *symbolOp : symbols)1075 symbolToUsers[symbolOp].insert(use.getUser());1076 }1077 }1078 };1079 // We just set `allSymUsesVisible` to false here because it isn't necessary1080 // for building the user map.1081 SymbolTable::walkSymbolTables(symbolTableOp, /*allSymUsesVisible=*/false,1082 walkFn);1083}1084 1085void SymbolUserMap::replaceAllUsesWith(Operation *symbol,1086 StringAttr newSymbolName) {1087 auto it = symbolToUsers.find(symbol);1088 if (it == symbolToUsers.end())1089 return;1090 1091 // Replace the uses within the users of `symbol`.1092 for (Operation *user : it->second)1093 (void)SymbolTable::replaceAllSymbolUses(symbol, newSymbolName, user);1094 1095 // Move the current users of `symbol` to the new symbol if it is in the1096 // symbol table.1097 Operation *newSymbol =1098 symbolTable.lookupSymbolIn(symbol->getParentOp(), newSymbolName);1099 if (newSymbol != symbol) {1100 // Transfer over the users to the new symbol. The reference to the old one1101 // is fetched again as the iterator is invalidated during the insertion.1102 auto newIt = symbolToUsers.try_emplace(newSymbol);1103 auto oldIt = symbolToUsers.find(symbol);1104 assert(oldIt != symbolToUsers.end() && "missing old users list");1105 if (newIt.second)1106 newIt.first->second = std::move(oldIt->second);1107 else1108 newIt.first->second.set_union(oldIt->second);1109 symbolToUsers.erase(oldIt);1110 }1111}1112 1113//===----------------------------------------------------------------------===//1114// Visibility parsing implementation.1115//===----------------------------------------------------------------------===//1116 1117ParseResult impl::parseOptionalVisibilityKeyword(OpAsmParser &parser,1118 NamedAttrList &attrs) {1119 StringRef visibility;1120 if (parser.parseOptionalKeyword(&visibility, {"public", "private", "nested"}))1121 return failure();1122 1123 StringAttr visibilityAttr = parser.getBuilder().getStringAttr(visibility);1124 attrs.push_back(parser.getBuilder().getNamedAttr(1125 SymbolTable::getVisibilityAttrName(), visibilityAttr));1126 return success();1127}1128 1129//===----------------------------------------------------------------------===//1130// Symbol Interfaces1131//===----------------------------------------------------------------------===//1132 1133/// Include the generated symbol interfaces.1134#include "mlir/IR/SymbolInterfaces.cpp.inc"1135