| //===- TestSymbolUses.cpp - Pass to test symbol uselists ------------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "TestDialect.h" |
| #include "mlir/IR/BuiltinOps.h" |
| #include "mlir/Pass/Pass.h" |
| |
| using namespace mlir; |
| |
| namespace { |
| /// This is a symbol test pass that tests the symbol uselist functionality |
| /// provided by the symbol table along with erasing from the symbol table. |
| struct SymbolUsesPass |
| : public PassWrapper<SymbolUsesPass, OperationPass<ModuleOp>> { |
| StringRef getArgument() const final { return "test-symbol-uses"; } |
| StringRef getDescription() const final { |
| return "Test detection of symbol uses"; |
| } |
| WalkResult operateOnSymbol(Operation *symbol, ModuleOp module, |
| SmallVectorImpl<FuncOp> &deadFunctions) { |
| // Test computing uses on a non symboltable op. |
| Optional<SymbolTable::UseRange> symbolUses = |
| SymbolTable::getSymbolUses(symbol); |
| |
| // Test the conservative failure case. |
| if (!symbolUses) { |
| symbol->emitRemark() |
| << "symbol contains an unknown nested operation that " |
| "'may' define a new symbol table"; |
| return WalkResult::interrupt(); |
| } |
| if (unsigned numUses = llvm::size(*symbolUses)) |
| symbol->emitRemark() << "symbol contains " << numUses |
| << " nested references"; |
| |
| // Test the functionality of symbolKnownUseEmpty. |
| if (SymbolTable::symbolKnownUseEmpty(symbol, &module.getBodyRegion())) { |
| FuncOp funcSymbol = dyn_cast<FuncOp>(symbol); |
| if (funcSymbol && funcSymbol.isExternal()) |
| deadFunctions.push_back(funcSymbol); |
| |
| symbol->emitRemark() << "symbol has no uses"; |
| return WalkResult::advance(); |
| } |
| |
| // Test the functionality of getSymbolUses. |
| symbolUses = SymbolTable::getSymbolUses(symbol, &module.getBodyRegion()); |
| assert(symbolUses.hasValue() && "expected no unknown operations"); |
| for (SymbolTable::SymbolUse symbolUse : *symbolUses) { |
| // Check that we can resolve back to our symbol. |
| if (SymbolTable::lookupNearestSymbolFrom( |
| symbolUse.getUser()->getParentOp(), symbolUse.getSymbolRef())) { |
| symbolUse.getUser()->emitRemark() |
| << "found use of symbol : " << symbolUse.getSymbolRef() << " : " |
| << symbol->getAttr(SymbolTable::getSymbolAttrName()); |
| } |
| } |
| symbol->emitRemark() << "symbol has " << llvm::size(*symbolUses) << " uses"; |
| return WalkResult::advance(); |
| } |
| |
| void runOnOperation() override { |
| auto module = getOperation(); |
| |
| // Walk nested symbols. |
| SmallVector<FuncOp, 4> deadFunctions; |
| module.getBodyRegion().walk([&](Operation *nestedOp) { |
| if (isa<SymbolOpInterface>(nestedOp)) |
| return operateOnSymbol(nestedOp, module, deadFunctions); |
| return WalkResult::advance(); |
| }); |
| |
| SymbolTable table(module); |
| for (Operation *op : deadFunctions) { |
| // In order to test the SymbolTable::erase method, also erase completely |
| // useless functions. |
| auto name = SymbolTable::getSymbolName(op); |
| assert(table.lookup(name) && "expected no unknown operations"); |
| table.erase(op); |
| assert(!table.lookup(name) && |
| "expected erased operation to be unknown now"); |
| module.emitRemark() << name.getValue() << " function successfully erased"; |
| } |
| } |
| }; |
| |
| /// This is a symbol test pass that tests the symbol use replacement |
| /// functionality provided by the symbol table. |
| struct SymbolReplacementPass |
| : public PassWrapper<SymbolReplacementPass, OperationPass<ModuleOp>> { |
| StringRef getArgument() const final { return "test-symbol-rauw"; } |
| StringRef getDescription() const final { |
| return "Test replacement of symbol uses"; |
| } |
| void runOnOperation() override { |
| ModuleOp module = getOperation(); |
| |
| // Don't try to replace if we can't collect symbol uses. |
| if (!SymbolTable::getSymbolUses(&module.getBodyRegion())) |
| return; |
| |
| SymbolTableCollection symbolTable; |
| SymbolUserMap symbolUsers(symbolTable, module); |
| module.getBodyRegion().walk([&](Operation *nestedOp) { |
| StringAttr newName = nestedOp->getAttrOfType<StringAttr>("sym.new_name"); |
| if (!newName) |
| return; |
| symbolUsers.replaceAllUsesWith(nestedOp, newName); |
| SymbolTable::setSymbolName(nestedOp, newName); |
| }); |
| } |
| }; |
| } // end anonymous namespace |
| |
| namespace mlir { |
| void registerSymbolTestPasses() { |
| PassRegistration<SymbolUsesPass>(); |
| |
| PassRegistration<SymbolReplacementPass>(); |
| } |
| } // namespace mlir |