1 //===- TestSymbolUses.cpp - Pass to test symbol uselists ------------------===//
2 //
3 // Part of the MLIR 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 "TestDialect.h"
10 #include "mlir/IR/Function.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 : public ModulePass<SymbolUsesPass> {
19   WalkResult operateOnSymbol(Operation *symbol, Operation *module,
20                              SmallVectorImpl<FuncOp> &deadFunctions) {
21     // Test computing uses on a non symboltable op.
22     Optional<SymbolTable::UseRange> symbolUses =
23         SymbolTable::getSymbolUses(symbol);
24 
25     // Test the conservative failure case.
26     if (!symbolUses) {
27       symbol->emitRemark()
28           << "symbol contains an unknown nested operation that "
29              "'may' define a new symbol table";
30       return WalkResult::interrupt();
31     }
32     if (unsigned numUses = llvm::size(*symbolUses))
33       symbol->emitRemark() << "symbol contains " << numUses
34                            << " nested references";
35 
36     // Test the functionality of symbolKnownUseEmpty.
37     if (SymbolTable::symbolKnownUseEmpty(symbol, module)) {
38       FuncOp funcSymbol = dyn_cast<FuncOp>(symbol);
39       if (funcSymbol && funcSymbol.isExternal())
40         deadFunctions.push_back(funcSymbol);
41 
42       symbol->emitRemark() << "symbol has no uses";
43       return WalkResult::advance();
44     }
45 
46     // Test the functionality of getSymbolUses.
47     symbolUses = SymbolTable::getSymbolUses(symbol, module);
48     assert(symbolUses.hasValue() && "expected no unknown operations");
49     for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
50       // Check that we can resolve back to our symbol.
51       if (SymbolTable::lookupNearestSymbolFrom(
52               symbolUse.getUser()->getParentOp(), symbolUse.getSymbolRef())) {
53         symbolUse.getUser()->emitRemark()
54             << "found use of symbol : " << symbolUse.getSymbolRef() << " : "
55             << symbol->getAttr(SymbolTable::getSymbolAttrName());
56       }
57     }
58     symbol->emitRemark() << "symbol has " << llvm::size(*symbolUses) << " uses";
59     return WalkResult::advance();
60   }
61 
62   void runOnModule() override {
63     auto module = getModule();
64 
65     // Walk nested symbols.
66     SmallVector<FuncOp, 4> deadFunctions;
67     module.getBodyRegion().walk([&](Operation *nestedOp) {
68       if (SymbolTable::isSymbol(nestedOp))
69         return operateOnSymbol(nestedOp, module, deadFunctions);
70       return WalkResult::advance();
71     });
72 
73     for (Operation *op : deadFunctions) {
74       // In order to test the SymbolTable::erase method, also erase completely
75       // useless functions.
76       SymbolTable table(module);
77       auto name = SymbolTable::getSymbolName(op);
78       assert(table.lookup(name) && "expected no unknown operations");
79       table.erase(op);
80       assert(!table.lookup(name) &&
81              "expected erased operation to be unknown now");
82       module.emitRemark() << name << " function successfully erased";
83     }
84   }
85 };
86 
87 /// This is a symbol test pass that tests the symbol use replacement
88 /// functionality provided by the symbol table.
89 struct SymbolReplacementPass : public ModulePass<SymbolReplacementPass> {
90   void runOnModule() override {
91     auto module = getModule();
92 
93     // Walk nested functions and modules.
94     module.getBodyRegion().walk([&](Operation *nestedOp) {
95       StringAttr newName = nestedOp->getAttrOfType<StringAttr>("sym.new_name");
96       if (!newName)
97         return;
98       if (succeeded(SymbolTable::replaceAllSymbolUses(
99               nestedOp, newName.getValue(), module)))
100         SymbolTable::setSymbolName(nestedOp, newName.getValue());
101     });
102   }
103 };
104 } // end anonymous namespace
105 
106 static PassRegistration<SymbolUsesPass> pass("test-symbol-uses",
107                                              "Test detection of symbol uses");
108 
109 static PassRegistration<SymbolReplacementPass>
110     rauwPass("test-symbol-rauw", "Test replacement of symbol uses");
111