[mlir][IR] Add a new SymbolUserMap class

[rriddle]

Summary

This class provides efficient implementations of symbol queries related to uses, such as collecting the users of a symbol, replacing all uses, etc. This provides similar benefits to use related queries, as SymbolTableCollection did for lookup queries.

[jpienaar]

Thanks this will help a lot :)

mlir/lib/IR/SymbolTable.cpp

1027

Comment for allSymUsesVisible logic here?

[mehdi_amini]

mlir/include/mlir/IR/SymbolTable.h

273

The sentence is ambiguous: will you get symbol uses outside of limit for symbols defined under limit?

[rriddle]

update

mlir/include/mlir/IR/SymbolTable.h

 //===- SymbolTable.h - MLIR Symbol Table Class ------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_SYMBOLTABLE_H
 #define MLIR_IR_SYMBOLTABLE_H
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace mlir {
 class Identifier;
 class Operation;
 
 /// This class allows for representing and managing the symbol table used by
 /// operations with the 'SymbolTable' trait. Inserting into and erasing from
public:
   SymbolTable &getSymbolTable(Operation *op);
 
 private:
   /// The constructed symbol tables nested within this table.
   DenseMap<Operation *, std::unique_ptr<SymbolTable>> symbolTables;
 };
 
 //===----------------------------------------------------------------------===//
+// SymbolUserMap
+//===----------------------------------------------------------------------===//
+
+/// This class represents a map of symbols to users, and provides efficient
+/// implementations of symbol queries related to users; such as collecting the
+/// users of a symbol, replacing all uses, etc.
+class SymbolUserMap {
+public:
+  /// Build a user map for all of the symbols defined in regions nested under
+  /// 'symbolTableOp'. A reference to the provided symbol table collection is
+  /// kept by the user map to ensure efficient lookups, thus the lifetime should
+  /// extend beyond that of this map.
+  SymbolUserMap(SymbolTableCollection &symbolTable, Operation *symbolTableOp);
+
+  /// Return the users of the provided symbol operation.
+  ArrayRef<Operation *> getUsers(Operation *symbol) const {
+    auto it = symbolToUsers.find(symbol);
+    return it != symbolToUsers.end() ? it->second.getArrayRef() : llvm::None;
+  }
+
+  /// Return true if the given symbol has no uses.
+  bool use_empty(Operation *symbol) const {
+    return !symbolToUsers.count(symbol);
+  }
+
+  /// Replace all of the uses of the given symbol with `newSymbolName`.
+  void replaceAllUsesWith(Operation *symbol, StringRef newSymbolName);
+
+private:
+  /// A reference to the symbol table used to construct this map.
+  SymbolTableCollection &symbolTable;
+
+  /// A map of symbol operations to symbol users.
+  DenseMap<Operation *, llvm::SetVector<Operation *>> symbolToUsers;
+};
+
+//===----------------------------------------------------------------------===//
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 LogicalResult verifySymbolTable(Operation *op);
 LogicalResult verifySymbol(Operation *op);
 } // namespace detail
 

mlir/lib/IR/SymbolTable.cpp

 SymbolTable &SymbolTableCollection::getSymbolTable(Operation *op) {
   auto it = symbolTables.try_emplace(op, nullptr);
   if (it.second)
     it.first->second = std::make_unique<SymbolTable>(op);
   return *it.first->second;
 }
 
 //===----------------------------------------------------------------------===//
+// SymbolUserMap
+//===----------------------------------------------------------------------===//
+
+SymbolUserMap::SymbolUserMap(SymbolTableCollection &symbolTable,
+                             Operation *symbolTableOp)
+    : symbolTable(symbolTable) {
+  // Walk each of the symbol tables looking for discardable callgraph nodes.
+  SmallVector<Operation *> symbols;
+  auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {
+    for (Operation &nestedOp : symbolTableOp->getRegion(0).getOps()) {
+      auto symbolUses = SymbolTable::getSymbolUses(&nestedOp);
+      assert(symbolUses && "expected uses to be valid");
+
+      for (const SymbolTable::SymbolUse &use : *symbolUses) {
+        symbols.clear();
+        (void)symbolTable.lookupSymbolIn(symbolTableOp, use.getSymbolRef(),
+                                         symbols);
+        for (Operation *symbolOp : symbols)
+          symbolToUsers[symbolOp].insert(use.getUser());
+      }
+    }
+  };
+  // We just set `allSymUsesVisible` to false here because it isn't necessary
+  // for building the user map.
+  SymbolTable::walkSymbolTables(symbolTableOp, /*allSymUsesVisible=*/false,
+                                walkFn);
+}
+
+void SymbolUserMap::replaceAllUsesWith(Operation *symbol,
+                                       StringRef newSymbolName) {
+  auto it = symbolToUsers.find(symbol);
+  if (it == symbolToUsers.end())
+    return;
+  llvm::SetVector<Operation *> &users = it->second;
+
+  // Replace the uses within the users of `symbol`.
+  for (Operation *user : users)
+    (void)SymbolTable::replaceAllSymbolUses(symbol, newSymbolName, user);
+
+  // Move the current users of `symbol` to the new symbol if it is in the
+  // symbol table.
+  Operation *newSymbol =
+      symbolTable.lookupSymbolIn(symbol->getParentOp(), newSymbolName);
+  if (newSymbol != symbol) {
+    // Transfer over the users to the new symbol.
+    auto newIt = symbolToUsers.find(newSymbol);
+    if (newIt == symbolToUsers.end())
+      symbolToUsers.try_emplace(newSymbol, std::move(users));
+    else
+      newIt->second.set_union(users);
+    symbolToUsers.erase(symbol);
+  }
+}
+
+//===----------------------------------------------------------------------===//
 // Visibility parsing implementation.
 //===----------------------------------------------------------------------===//
 
 ParseResult impl::parseOptionalVisibilityKeyword(OpAsmParser &parser,
                                                  NamedAttrList &attrs) {
   StringRef visibility;
   if (parser.parseOptionalKeyword(&visibility, {"public", "private", "nested"}))
     return failure();

mlir/test/lib/IR/TestSymbolUses.cpp

struct SymbolUsesPass
   }
 };
 
 /// 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>> {
   void runOnOperation() override {
-    auto module = getOperation();
+    ModuleOp module = getOperation();
+
+    // Don't try to replace if we can't collect symbol uses.
+    if (!SymbolTable::getSymbolUses(&module.getBodyRegion()))
+      return;
 
-    // Walk nested functions and modules.
+    SymbolTableCollection symbolTable;
+    SymbolUserMap symbolUsers(symbolTable, module);
     module.getBodyRegion().walk([&](Operation *nestedOp) {
       StringAttr newName = nestedOp->getAttrOfType<StringAttr>("sym.new_name");
       if (!newName)
         return;
-      if (succeeded(SymbolTable::replaceAllSymbolUses(
-              nestedOp, newName.getValue(), &module.getBodyRegion())))
-        SymbolTable::setSymbolName(nestedOp, newName.getValue());
+      symbolUsers.replaceAllUsesWith(nestedOp, newName.getValue());
     });
   }
 };
 } // end anonymous namespace
 
 namespace mlir {
 void registerSymbolTestPasses() {
   PassRegistration<SymbolUsesPass>("test-symbol-uses",
                                    "Test detection of symbol uses");
 
   PassRegistration<SymbolReplacementPass>("test-symbol-rauw",
                                           "Test replacement of symbol uses");
 }
 } // namespace mlir