1 | //===- TestSymbolUses.cpp - Pass to test symbol uselists ------------------===// |
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-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | |
9 | #include "TestOps.h" |
10 | #include "mlir/IR/BuiltinOps.h" |
11 | #include "mlir/Pass/Pass.h" |
12 | |
13 | using namespace mlir; |
14 | |
15 | namespace { |
16 | /// This is a symbol test pass that tests the symbol uselist functionality |
17 | /// provided by the symbol table along with erasing from the symbol table. |
18 | struct SymbolUsesPass |
19 | : public PassWrapper<SymbolUsesPass, OperationPass<ModuleOp>> { |
20 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SymbolUsesPass) |
21 | |
22 | StringRef getArgument() const final { return "test-symbol-uses" ; } |
23 | StringRef getDescription() const final { |
24 | return "Test detection of symbol uses" ; |
25 | } |
26 | WalkResult operateOnSymbol(Operation *symbol, ModuleOp module, |
27 | SmallVectorImpl<func::FuncOp> &deadFunctions) { |
28 | // Test computing uses on a non symboltable op. |
29 | std::optional<SymbolTable::UseRange> symbolUses = |
30 | SymbolTable::getSymbolUses(from: symbol); |
31 | |
32 | // Test the conservative failure case. |
33 | if (!symbolUses) { |
34 | symbol->emitRemark() |
35 | << "symbol contains an unknown nested operation that " |
36 | "'may' define a new symbol table" ; |
37 | return WalkResult::interrupt(); |
38 | } |
39 | if (unsigned numUses = llvm::size(Range&: *symbolUses)) |
40 | symbol->emitRemark() << "symbol contains " << numUses |
41 | << " nested references" ; |
42 | |
43 | // Test the functionality of symbolKnownUseEmpty. |
44 | if (SymbolTable::symbolKnownUseEmpty(symbol, &module.getBodyRegion())) { |
45 | func::FuncOp funcSymbol = dyn_cast<func::FuncOp>(symbol); |
46 | if (funcSymbol && funcSymbol.isExternal()) |
47 | deadFunctions.push_back(funcSymbol); |
48 | |
49 | symbol->emitRemark() << "symbol has no uses" ; |
50 | return WalkResult::advance(); |
51 | } |
52 | |
53 | // Test the functionality of getSymbolUses. |
54 | symbolUses = SymbolTable::getSymbolUses(symbol, &module.getBodyRegion()); |
55 | assert(symbolUses && "expected no unknown operations" ); |
56 | for (SymbolTable::SymbolUse symbolUse : *symbolUses) { |
57 | // Check that we can resolve back to our symbol. |
58 | if (SymbolTable::lookupNearestSymbolFrom( |
59 | symbolUse.getUser()->getParentOp(), symbolUse.getSymbolRef())) { |
60 | symbolUse.getUser()->emitRemark() |
61 | << "found use of symbol : " << symbolUse.getSymbolRef() << " : " |
62 | << *symbol->getInherentAttr(name: SymbolTable::getSymbolAttrName()); |
63 | } |
64 | } |
65 | symbol->emitRemark() << "symbol has " << llvm::size(Range&: *symbolUses) << " uses" ; |
66 | return WalkResult::advance(); |
67 | } |
68 | |
69 | void runOnOperation() override { |
70 | auto module = getOperation(); |
71 | |
72 | // Walk nested symbols. |
73 | SmallVector<func::FuncOp, 4> deadFunctions; |
74 | module.getBodyRegion().walk([&](Operation *nestedOp) { |
75 | if (isa<SymbolOpInterface>(nestedOp)) |
76 | return operateOnSymbol(nestedOp, module, deadFunctions); |
77 | return WalkResult::advance(); |
78 | }); |
79 | |
80 | SymbolTable table(module); |
81 | for (Operation *op : deadFunctions) { |
82 | // In order to test the SymbolTable::erase method, also erase completely |
83 | // useless functions. |
84 | auto name = SymbolTable::getSymbolName(op); |
85 | assert(table.lookup(name) && "expected no unknown operations" ); |
86 | table.erase(op); |
87 | assert(!table.lookup(name) && |
88 | "expected erased operation to be unknown now" ); |
89 | module.emitRemark() << name.getValue() << " function successfully erased" ; |
90 | } |
91 | } |
92 | }; |
93 | |
94 | /// This is a symbol test pass that tests the symbol use replacement |
95 | /// functionality provided by the symbol table. |
96 | struct SymbolReplacementPass |
97 | : public PassWrapper<SymbolReplacementPass, OperationPass<ModuleOp>> { |
98 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SymbolReplacementPass) |
99 | |
100 | StringRef getArgument() const final { return "test-symbol-rauw" ; } |
101 | StringRef getDescription() const final { |
102 | return "Test replacement of symbol uses" ; |
103 | } |
104 | void runOnOperation() override { |
105 | ModuleOp module = getOperation(); |
106 | |
107 | // Don't try to replace if we can't collect symbol uses. |
108 | if (!SymbolTable::getSymbolUses(&module.getBodyRegion())) |
109 | return; |
110 | |
111 | SymbolTableCollection symbolTable; |
112 | SymbolUserMap symbolUsers(symbolTable, module); |
113 | module.getBodyRegion().walk([&](Operation *nestedOp) { |
114 | StringAttr newName = nestedOp->getAttrOfType<StringAttr>("sym.new_name" ); |
115 | if (!newName) |
116 | return; |
117 | symbolUsers.replaceAllUsesWith(symbol: nestedOp, newSymbolName: newName); |
118 | SymbolTable::setSymbolName(nestedOp, newName); |
119 | }); |
120 | } |
121 | }; |
122 | } // namespace |
123 | |
124 | namespace mlir { |
125 | void registerSymbolTestPasses() { |
126 | PassRegistration<SymbolUsesPass>(); |
127 | |
128 | PassRegistration<SymbolReplacementPass>(); |
129 | } |
130 | } // namespace mlir |
131 | |