[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

3 files changed, 129 insertions, 50 deletions

diff --git a/mlir/test/IR/test-symbol-rauw.mlir b/mlir/test/IR/test-symbol-rauw.mlir
index 963875e762e..25d22e43225 100644
--- a/mlir/test/IR/test-symbol-rauw.mlir
+++ b/mlir/test/IR/test-symbol-rauw.mlir
@@ -32,6 +32,39 @@ module attributes {sym.outside_use = @symbol_foo } {
 
 // -----
 
+// Check the support for nested references.
+
+// CHECK: module
+module {
+  // CHECK: module @module_a
+  module @module_a {
+    // CHECK: func @replaced_foo
+    func @foo() attributes {sym.new_name = "replaced_foo" }
+  }
+
+  // CHECK: module @replaced_module_b
+  module @module_b attributes {sym.new_name = "replaced_module_b"} {
+    // CHECK: module @replaced_module_c
+    module @module_c attributes {sym.new_name = "replaced_module_c"} {
+      // CHECK: func @replaced_foo
+      func @foo() attributes {sym.new_name = "replaced_foo" }
+    }
+  }
+
+  // CHECK: func @symbol_bar
+  func @symbol_bar() {
+    // CHECK: foo.op
+    // CHECK-SAME: use_1 = @module_a::@replaced_foo
+    // CHECK-SAME: use_2 = @replaced_module_b::@replaced_module_c::@replaced_foo
+    "foo.op"() {
+      use_1 = @module_a::@foo,
+      use_2 = @module_b::@module_c::@foo
+    } : () -> ()
+  }
+}
+
+// -----
+
 // Check that the replacement fails for potentially unknown symbol tables.
 module {
   // CHECK: func @failed_repl
diff --git a/mlir/test/IR/test-symbol-uses.mlir b/mlir/test/IR/test-symbol-uses.mlir
index 4d6555aceed..f95ac6b2d70 100644
--- a/mlir/test/IR/test-symbol-uses.mlir
+++ b/mlir/test/IR/test-symbol-uses.mlir
@@ -4,14 +4,14 @@
 // its table.
 // expected-remark@below {{symbol_removable function successfully erased}}
 module attributes {sym.outside_use = @symbol_foo } {
-  // expected-remark@+1 {{function has 2 uses}}
+  // expected-remark@+1 {{symbol has 2 uses}}
   func @symbol_foo()
 
-  // expected-remark@below {{function has no uses}}
-  // expected-remark@below {{found use of function : @symbol_foo}}
-  // expected-remark@below {{function contains 2 nested references}}
+  // expected-remark@below {{symbol has no uses}}
+  // expected-remark@below {{found use of symbol : @symbol_foo}}
+  // expected-remark@below {{symbol contains 2 nested references}}
   func @symbol_bar() attributes {sym.use = @symbol_foo} {
-    // expected-remark@+1 {{found use of function : @symbol_foo}}
+    // expected-remark@+1 {{found use of symbol : @symbol_foo}}
     "foo.op"() {
       non_symbol_attr,
       use = [{ nested_symbol = [@symbol_foo]}],
@@ -19,13 +19,13 @@ module attributes {sym.outside_use = @symbol_foo } {
     } : () -> ()
   }
 
-  // expected-remark@below {{function has no uses}}
+  // expected-remark@below {{symbol has no uses}}
   func @symbol_removable()
 
-  // expected-remark@+1 {{function has 1 use}}
+  // expected-remark@+1 {{symbol has 1 use}}
   func @symbol_baz()
 
-  // expected-remark@+1 {{found use of function : @symbol_baz}}
+  // expected-remark@+1 {{found use of symbol : @symbol_baz}}
   module attributes {test.reference = @symbol_baz} {
     "foo.op"() {test.nested_reference = @symbol_baz} : () -> ()
   }
@@ -33,6 +33,34 @@ module attributes {sym.outside_use = @symbol_foo } {
 
 // -----
 
+// Test nested attribute support
+module {
+  // expected-remark@+1 {{symbol has 2 uses}}
+  module @module_b {
+    // expected-remark@+1 {{symbol has 1 uses}}
+    module @module_c {
+      // expected-remark@+1 {{symbol has 1 uses}}
+      func @foo()
+    }
+  }
+
+  // expected-remark@below {{symbol has no uses}}
+  // expected-remark@below {{symbol contains 2 nested references}}
+  func @symbol_bar() {
+    // expected-remark@below {{found use of symbol : @module_b::@module_c::@foo : "foo"}}
+    // expected-remark@below {{found use of symbol : @module_b::@module_c::@foo : "module_c"}}
+    // expected-remark@below {{found use of symbol : @module_b::@module_c::@foo : "module_b"}}
+    // expected-remark@below {{found use of symbol : @module_b : "module_b"}}
+    "foo.op"() {
+      use_1 = [{ nested_symbol = [@module_b::@module_c::@foo]}],
+      use_2 = @module_b
+    } : () -> ()
+  }
+}
+
+
+// -----
+
 // expected-remark@+1 {{contains an unknown nested operation that 'may' define a new symbol table}}
 func @symbol_bar() {
   "foo.possibly_unknown_symbol_table"() ({
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