[mlir] Update the use-list algorithms in SymbolTable to support nested references.

Summary: This updates the use list algorithms to support querying from a specific symbol, allowing for the collection and detection of nested references. This works by walking the parent "symbol scopes" and applying the existing algorithm at each level.

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D72042

1 files changed, 60 insertions, 42 deletions

diff --git a/mlir/test/lib/IR/TestSymbolUses.cpp b/mlir/test/lib/IR/TestSymbolUses.cpp
index c8fb1d8eecf..dc0608b21c5 100644
--- a/mlir/test/lib/IR/TestSymbolUses.cpp
+++ b/mlir/test/lib/IR/TestSymbolUses.cpp
@@ -16,54 +16,70 @@ 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 ModulePass<SymbolUsesPass> {
-  void runOnModule() override {
-    auto module = getModule();
-    std::vector<FuncOp> ops_to_delete;
+  WalkResult operateOnSymbol(Operation *symbol, Operation *module,
+                             SmallVectorImpl<FuncOp> &deadFunctions) {
+    // Test computing uses on a non symboltable op.
+    Optional<SymbolTable::UseRange> symbolUses =
+        SymbolTable::getSymbolUses(symbol);
 
-    for (FuncOp func : module.getOps<FuncOp>()) {
-      // Test computing uses on a non symboltable op.
-      Optional<SymbolTable::UseRange> symbolUses =
-          SymbolTable::getSymbolUses(func);
+    // 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 conservative failure case.
-      if (!symbolUses) {
-        func.emitRemark() << "function contains an unknown nested operation "
-                             "that 'may' define a new symbol table";
-        return;
-      }
-      if (unsigned numUses = llvm::size(*symbolUses))
-        func.emitRemark() << "function contains " << numUses
-                          << " nested references";
+    // Test the functionality of symbolKnownUseEmpty.
+    if (SymbolTable::symbolKnownUseEmpty(symbol, module)) {
+      FuncOp funcSymbol = dyn_cast<FuncOp>(symbol);
+      if (funcSymbol && funcSymbol.isExternal())
+        deadFunctions.push_back(funcSymbol);
 
-      // Test the functionality of symbolKnownUseEmpty.
-      if (func.symbolKnownUseEmpty(module)) {
-        func.emitRemark() << "function has no uses";
-        if (func.getBody().empty())
-          ops_to_delete.push_back(func);
-        continue;
-      }
+      symbol->emitRemark() << "symbol has no uses";
+      return WalkResult::advance();
+    }
 
-      // Test the functionality of getSymbolUses.
-      symbolUses = func.getSymbolUses(module);
-      assert(symbolUses.hasValue() && "expected no unknown operations");
-      for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
+    // Test the functionality of getSymbolUses.
+    symbolUses = SymbolTable::getSymbolUses(symbol, module);
+    assert(symbolUses.hasValue() && "expected no unknown operations");
+    for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
+      // Check that we can resolve back to our symbol.
+      if (Operation *op = SymbolTable::lookupNearestSymbolFrom(
+              symbolUse.getUser()->getParentOp(), symbolUse.getSymbolRef())) {
         symbolUse.getUser()->emitRemark()
-            << "found use of function : " << symbolUse.getSymbolRef();
+            << "found use of symbol : " << symbolUse.getSymbolRef() << " : "
+            << symbol->getAttr(SymbolTable::getSymbolAttrName());
       }
-      func.emitRemark() << "function has " << llvm::size(*symbolUses)
-                        << " uses";
     }
+    symbol->emitRemark() << "symbol has " << llvm::size(*symbolUses) << " uses";
+    return WalkResult::advance();
+  }
+
+  void runOnModule() override {
+    auto module = getModule();
 
-    for (FuncOp func : ops_to_delete) {
+    // Walk nested symbols.
+    SmallVector<FuncOp, 4> deadFunctions;
+    module.getBodyRegion().walk([&](Operation *nestedOp) {
+      if (SymbolTable::isSymbol(nestedOp))
+        return operateOnSymbol(nestedOp, module, deadFunctions);
+      return WalkResult::advance();
+    });
+
+    for (Operation *op : deadFunctions) {
       // In order to test the SymbolTable::erase method, also erase completely
       // useless functions.
       SymbolTable table(module);
-      auto func_name = func.getName();
-      assert(table.lookup(func_name) && "expected no unknown operations");
-      table.erase(func);
-      assert(!table.lookup(func_name) &&
+      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() << func_name << " function successfully erased";
+      module.emitRemark() << name << " function successfully erased";
     }
   }
 };
@@ -74,13 +90,15 @@ struct SymbolReplacementPass : public ModulePass<SymbolReplacementPass> {
   void runOnModule() override {
     auto module = getModule();
 
-    for (FuncOp func : module.getOps<FuncOp>()) {
-      StringAttr newName = func.getAttrOfType<StringAttr>("sym.new_name");
+    // Walk nested functions and modules.
+    module.getBodyRegion().walk([&](Operation *nestedOp) {
+      StringAttr newName = nestedOp->getAttrOfType<StringAttr>("sym.new_name");
       if (!newName)
-        continue;
-      if (succeeded(func.replaceAllSymbolUses(newName.getValue(), module)))
-        func.setName(newName.getValue());
-    }
+        return;
+      if (succeeded(SymbolTable::replaceAllSymbolUses(
+              nestedOp, newName.getValue(), module)))
+        SymbolTable::setSymbolName(nestedOp, newName.getValue());
+    });
   }
 };
 } // end anonymous namespace