[mlir] Add a SharedSymbolTableCollection class

[Mogball]

Summary

This class wraps a SymbolTableCollection to allow shared access to the
collection of symbol tables (but not the individual symbol tables). This
can be used, for example, in a pass that shards work among symbols that
requires symbol lookups.

[rriddle]

This seems like a fine start. It'd be nice if the name was a bit closer to "ThreadSafe" aspect, like "LockedSymbolTableCollection"(or something), to make it more apparent what the purpose is on first glance.

[rriddle]

mlir/include/mlir/IR/SymbolTable.h

296–297

Can you detail the thread safe aspects in the docs? I don't think "shared" accurately conveys that on its own.

[Mogball]

comments

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"
+#include "llvm/Support/RWMutex.h"
 
 namespace mlir {
 
 /// This class allows for representing and managing the symbol table used by
 /// operations with the 'SymbolTable' trait. Inserting into and erasing from
 /// this SymbolTable will also insert and erase from the Operation given to it
 /// at construction.
 class SymbolTable {
public:
   T lookupNearestSymbolFrom(Operation *from, SymbolRefAttr symbol) {
     return dyn_cast_or_null<T>(lookupNearestSymbolFrom(from, symbol));
   }
 
   /// Lookup, or create, a symbol table for an operation.
   SymbolTable &getSymbolTable(Operation *op);
 
 private:
+  friend class LockedSymbolTableCollection;
+
   /// The constructed symbol tables nested within this table.
   DenseMap<Operation *, std::unique_ptr<SymbolTable>> symbolTables;
 };
 
 //===----------------------------------------------------------------------===//
+// LockedSymbolTableCollection
+//===----------------------------------------------------------------------===//
+
+/// This class implements a lock-based shared wrapper around a symbol table
+/// collection that allows shared access to the collection of symbol tables.
+/// This class does not protect shared access to individual symbol tables.
+/// `SymbolTableCollection` lazily instantiates `SymbolTable` instances for
+/// symbol table operations, making read operations not thread-safe. This class
+/// provides a thread-safe `lookupSymbolIn` implementation by synchronizing the
+/// lazy `SymbolTable` lookup.
+class LockedSymbolTableCollection : public SymbolTableCollection {
+public:
+  explicit LockedSymbolTableCollection(SymbolTableCollection &collection)
+      : collection(collection) {}
+
+  /// Look up a symbol with the specified name within the specified symbol table
+  /// operation, returning null if no such name exists.
+  Operation *lookupSymbolIn(Operation *symbolTableOp, StringAttr symbol);
+  /// Look up a symbol with the specified name within the specified symbol table
+  /// operation, returning null if no such name exists.
+  Operation *lookupSymbolIn(Operation *symbolTableOp, FlatSymbolRefAttr symbol);
+  /// Look up a potentially nested symbol within the specified symbol table
+  /// operation, returning null if no such symbol exists.
+  Operation *lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr name);
+
+  /// Lookup a symbol of a particular kind within the specified symbol table,
+  /// returning null if the symbol was not found.
+  template <typename T, typename NameT>
+  T lookupSymbolIn(Operation *symbolTableOp, NameT &&name) {
+    return dyn_cast_or_null<T>(
+        lookupSymbolIn(symbolTableOp, std::forward<NameT>(name)));
+  }
+
+  /// A variant of 'lookupSymbolIn' that returns all of the symbols referenced
+  /// by a given SymbolRefAttr when resolved within the provided symbol table
+  /// operation. Returns failure if any of the nested references could not be
+  /// resolved.
+  LogicalResult lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr name,
+                               SmallVectorImpl<Operation *> &symbols);
+
+private:
+  /// Get the symbol table for the symbol table operation, constructing if it
+  /// does not exist. This function provides thread safety over `collection`
+  /// by locking when performing the lookup and when inserting
+  /// lazily-constructed symbol tables.
+  SymbolTable &getSymbolTable(Operation *symbolTableOp);
+
+  /// The symbol tables to manage.
+  SymbolTableCollection &collection;
+  /// The mutex protecting access to the symbol table collection.
+  llvm::sys::SmartRWMutex<true> mutex;
+};
+
+//===----------------------------------------------------------------------===//
 // 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:

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;
 }
 
 //===----------------------------------------------------------------------===//
+// LockedSymbolTableCollection
+//===----------------------------------------------------------------------===//
+
+Operation *LockedSymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
+                                                       StringAttr symbol) {
+  return getSymbolTable(symbolTableOp).lookup(symbol);
+}
+
+Operation *
+LockedSymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
+                                            FlatSymbolRefAttr symbol) {
+  return lookupSymbolIn(symbolTableOp, symbol.getAttr());
+}
+
+Operation *LockedSymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
+                                                       SymbolRefAttr name) {
+  SmallVector<Operation *> symbols;
+  if (failed(lookupSymbolIn(symbolTableOp, name, symbols)))
+    return nullptr;
+  return symbols.back();
+}
+
+LogicalResult LockedSymbolTableCollection::lookupSymbolIn(
+    Operation *symbolTableOp, SymbolRefAttr name,
+    SmallVectorImpl<Operation *> &symbols) {
+  auto lookupFn = [this](Operation *symbolTableOp, StringAttr symbol) {
+    return lookupSymbolIn(symbolTableOp, symbol);
+  };
+  return lookupSymbolInImpl(symbolTableOp, name, symbols, lookupFn);
+}
+
+SymbolTable &
+LockedSymbolTableCollectio::getSymbolTable(Operation *symbolTableOp) {
+  assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
+  // Try to find an existing symbol table.
+  {
+    llvm::sys::SmartScopedReader<true> lock(mutex);
+    auto it = collection.symbolTables.find(symbolTableOp);
+    if (it != collection.symbolTables.end())
+      return *it->second;
+  }
+  // Create a symbol table for the operation. Perform construction outside of
+  // the critical section.
+  auto symbolTable = std::make_unique<SymbolTable>(symbolTableOp);
+  // Insert the constructed symbol table.
+  llvm::sys::SmartScopedWriter<true> lock(mutex);
+  return *collection.symbolTables
+              .insert({symbolTableOp, std::move(symbolTable)})
+              .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;