[mlir] Optimize symbol related checks in SymbolDCE

[rriddle]

Summary

This revision contains two optimizations related to symbol checking:

• Optimize SymbolOpInterface to only check for a name attribute if the operation is an optional symbol.

This removes an otherwise unnecessary attribute lookup from a majority of symbols.

• Add a new SymbolTableCollection class to represent a collection of SymbolTables.

This allows for perfoming non-flat symbol lookups in O(1) time by caching SymbolTables for symbol table operations. This class is very useful for algorithms that operate on multiple symbol tables, either recursively or not.

Depends On D89504

mlir/include/mlir/IR/OpDefinition.h

 public:
   using Base = OpInterface<ConcreteType, Traits>;
   using InterfaceBase = detail::Interface<ConcreteType, Operation *, Traits,
                                           Op<ConcreteType>, OpTrait::TraitBase>;
 
   /// Inherit the base class constructor.
   using InterfaceBase::InterfaceBase;
 
-private:
+protected:
   /// Returns the impl interface instance for the given operation.
   static typename InterfaceBase::Concept *getInterfaceFor(Operation *op) {
     // Access the raw interface from the abstract operation.
     auto *abstractOp = op->getAbstractOperation();
     return abstractOp ? abstractOp->getInterface<ConcreteType>() : nullptr;
   }
 
   /// Allow access to `getInterfaceFor`.

mlir/include/mlir/IR/SymbolInterfaces.td

if (concreteOp.isOptionalSymbol()) {
         return success();
     }
     return ::mlir::detail::verifySymbol($_op);
   }];
 
   let extraClassDeclaration = [{
     /// Custom classof that handles the case where the symbol is optional.
     static bool classof(Operation *op) {
-      return Base::classof(op)
-        && op->getAttr(::mlir::SymbolTable::getSymbolAttrName());
+      auto *concept = getInterfaceFor(op);
+      if (!concept)
+        return false;
+      return !concept->isOptionalSymbol(op) ||
+             op->getAttr(::mlir::SymbolTable::getSymbolAttrName());
     }
   }];
 
   let extraTraitClassDeclaration = [{
     using Visibility = mlir::SymbolTable::Visibility;
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Symbol Traits
 //===----------------------------------------------------------------------===//
 
 // Op defines a symbol table.
 def SymbolTable : NativeOpTrait<"SymbolTable">;
 
 #endif // MLIR_IR_SYMBOLINTERFACES

mlir/include/mlir/IR/SymbolTable.h

private:
   /// This is a mapping from a name to the symbol with that name.
   llvm::StringMap<Operation *> symbolTable;
 
   /// This is used when name conflicts are detected.
   unsigned uniquingCounter = 0;
 };
 
 //===----------------------------------------------------------------------===//
+// SymbolTableCollection
+//===----------------------------------------------------------------------===//
+
+/// This class represents a collection of `SymbolTable`s. This simplifies
+/// certain algorithms that run recursively on nested symbol tables. Symbol
+/// tables are constructed lazily to reduce the upfront cost of constructing
+/// unnecessary tables.
+class SymbolTableCollection {
+public:
+  /// 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, StringRef symbol);
+  Operation *lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr name);
+  template <typename T, typename NameT>
+  T lookupSymbolIn(Operation *symbolTableOp, NameT &&name) const {
+    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);
+
+  /// Lookup, or create, a symbol table for an operation.
+  SymbolTable &getSymbolTable(Operation *op);
+
+private:
+  /// The constructed symbol tables nested within this table.
+  DenseMap<Operation *, std::unique_ptr<SymbolTable>> symbolTables;
+};
+
+//===----------------------------------------------------------------------===//
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 LogicalResult verifySymbolTable(Operation *op);
 LogicalResult verifySymbol(Operation *op);
 } // namespace detail
 

mlir/lib/IR/SymbolTable.cpp

 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
                                        SymbolRefAttr symbol) {
   SmallVector<Operation *, 4> resolvedSymbols;
   if (failed(lookupSymbolIn(symbolTableOp, symbol, resolvedSymbols)))
     return nullptr;
   return resolvedSymbols.back();
 }
 
-LogicalResult
-SymbolTable::lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr symbol,
-                            SmallVectorImpl<Operation *> &symbols) {
+/// Internal implementation of `lookupSymbolIn` that allows for specialized
+/// implementations of the lookup function.
+static LogicalResult lookupSymbolInImpl(
+    Operation *symbolTableOp, SymbolRefAttr symbol,
+    SmallVectorImpl<Operation *> &symbols,
+    function_ref<Operation *(Operation *, StringRef)> lookupSymbolFn) {
   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
 
   // Lookup the root reference for this symbol.
-  symbolTableOp = lookupSymbolIn(symbolTableOp, symbol.getRootReference());
+  symbolTableOp = lookupSymbolFn(symbolTableOp, symbol.getRootReference());
   if (!symbolTableOp)
     return failure();
   symbols.push_back(symbolTableOp);
 
   // If there are no nested references, just return the root symbol directly.
   ArrayRef<FlatSymbolRefAttr> nestedRefs = symbol.getNestedReferences();
   if (nestedRefs.empty())
     return success();
 
   // Verify that the root is also a symbol table.
   if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())
     return failure();
 
   // Otherwise, lookup each of the nested non-leaf references and ensure that
   // each corresponds to a valid symbol table.
   for (FlatSymbolRefAttr ref : nestedRefs.drop_back()) {
-    symbolTableOp = lookupSymbolIn(symbolTableOp, ref.getValue());
+    symbolTableOp = lookupSymbolFn(symbolTableOp, ref.getValue());
     if (!symbolTableOp || !symbolTableOp->hasTrait<OpTrait::SymbolTable>())
       return failure();
     symbols.push_back(symbolTableOp);
   }
-  symbols.push_back(lookupSymbolIn(symbolTableOp, symbol.getLeafReference()));
+  symbols.push_back(lookupSymbolFn(symbolTableOp, symbol.getLeafReference()));
   return success(symbols.back());
 }
 
+LogicalResult
+SymbolTable::lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr symbol,
+                            SmallVectorImpl<Operation *> &symbols) {
+  auto lookupFn = [](Operation *symbolTableOp, StringRef symbol) {
+    return lookupSymbolIn(symbolTableOp, symbol);
+  };
+  return lookupSymbolInImpl(symbolTableOp, symbol, symbols, lookupFn);
+}
+
 /// Returns the operation registered with the given symbol name within the
 /// closes parent operation with the 'OpTrait::SymbolTable' trait. Returns
 /// nullptr if no valid symbol was found.
 Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,
                                                 StringRef symbol) {
   Operation *symbolTableOp = getNearestSymbolTable(from);
   return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;
 }

 }
 LogicalResult SymbolTable::replaceAllSymbolUses(Operation *oldSymbol,
                                                 StringRef newSymbol,
                                                 Region *from) {
   return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);
 }
 
 //===----------------------------------------------------------------------===//
+// SymbolTableCollection
+//===----------------------------------------------------------------------===//
+
+Operation *SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
+                                                 StringRef symbol) {
+  return getSymbolTable(symbolTableOp).lookup(symbol);
+}
+Operation *SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
+                                                 SymbolRefAttr name) {
+  SmallVector<Operation *, 4> symbols;
+  if (failed(lookupSymbolIn(symbolTableOp, name, symbols)))
+    return nullptr;
+  return symbols.back();
+}
+/// A variant of 'lookupSymbolIn' that returns all of the symbols referenced by
+/// a given SymbolRefAttr. Returns failure if any of the nested references could
+/// not be resolved.
+LogicalResult
+SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
+                                      SymbolRefAttr name,
+                                      SmallVectorImpl<Operation *> &symbols) {
+  auto lookupFn = [this](Operation *symbolTableOp, StringRef symbol) {
+    return lookupSymbolIn(symbolTableOp, symbol);
+  };
+  return lookupSymbolInImpl(symbolTableOp, name, symbols, lookupFn);
+}
+
+/// Lookup, or create, a symbol table for an operation.
+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;
+}
+
+//===----------------------------------------------------------------------===//
 // Symbol Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the generated symbol interfaces.
 #include "mlir/IR/SymbolInterfaces.cpp.inc"

mlir/lib/Transforms/SymbolDCE.cpp

 namespace {
 struct SymbolDCE : public SymbolDCEBase<SymbolDCE> {
   void runOnOperation() override;
 
   /// Compute the liveness of the symbols within the given symbol table.
   /// `symbolTableIsHidden` is true if this symbol table is known to be
   /// unaccessible from operations in its parent regions.
   LogicalResult computeLiveness(Operation *symbolTableOp,
+                                SymbolTableCollection &symbolTable,
                                 bool symbolTableIsHidden,
                                 DenseSet<Operation *> &liveSymbols);
 };
 } // end anonymous namespace
 
 void SymbolDCE::runOnOperation() {
   Operation *symbolTableOp = getOperation();
 
void SymbolDCE::runOnOperation() {
   // accessible from parent scopes.
   bool symbolTableIsHidden = true;
   SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(symbolTableOp);
   if (symbolTableOp->getParentOp() && symbol)
     symbolTableIsHidden = symbol.isPrivate();
 
   // Compute the set of live symbols within the symbol table.
   DenseSet<Operation *> liveSymbols;
-  if (failed(computeLiveness(symbolTableOp, symbolTableIsHidden, liveSymbols)))
+  SymbolTableCollection symbolTable;
+  if (failed(computeLiveness(symbolTableOp, symbolTable, symbolTableIsHidden,
+                             liveSymbols)))
     return signalPassFailure();
 
   // After computing the liveness, delete all of the symbols that were found to
   // be dead.
   symbolTableOp->walk([&](Operation *nestedSymbolTable) {
     if (!nestedSymbolTable->hasTrait<OpTrait::SymbolTable>())
       return;
     for (auto &block : nestedSymbolTable->getRegion(0)) {
       for (Operation &op :
            llvm::make_early_inc_range(block.without_terminator())) {
         if (isa<SymbolOpInterface>(&op) && !liveSymbols.count(&op))
           op.erase();
       }
     }
   });
 }
 
 /// Compute the liveness of the symbols within the given symbol table.
 /// `symbolTableIsHidden` is true if this symbol table is known to be
 /// unaccessible from operations in its parent regions.
 LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
+                                         SymbolTableCollection &symbolTable,
                                          bool symbolTableIsHidden,
                                          DenseSet<Operation *> &liveSymbols) {
   // A worklist of live operations to propagate uses from.
   SmallVector<Operation *, 16> worklist;
 
   // Walk the symbols within the current symbol table, marking the symbols that
   // are known to be live.
   for (auto &block : symbolTableOp->getRegion(0)) {
while (!worklist.empty()) {
     Operation *op = worklist.pop_back_val();
 
     // If this is a symbol table, recursively compute its liveness.
     if (op->hasTrait<OpTrait::SymbolTable>()) {
       // The internal symbol table is hidden if the parent is, if its not a
       // symbol, or if it is a private symbol.
       SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);
       bool symIsHidden = symbolTableIsHidden || !symbol || symbol.isPrivate();
-      if (failed(computeLiveness(op, symIsHidden, liveSymbols)))
+      if (failed(computeLiveness(op, symbolTable, symIsHidden, liveSymbols)))
         return failure();
     }
 
     // Collect the uses held by this operation.
     Optional<SymbolTable::UseRange> uses = SymbolTable::getSymbolUses(op);
     if (!uses) {
       return op->emitError()
              << "operation contains potentially unknown symbol table, "
                 "meaning that we can't reliable compute symbol uses";
     }
 
     SmallVector<Operation *, 4> resolvedSymbols;
     for (const SymbolTable::SymbolUse &use : *uses) {
       // Lookup the symbols referenced by this use.
       resolvedSymbols.clear();
-      if (failed(SymbolTable::lookupSymbolIn(
+      if (failed(symbolTable.lookupSymbolIn(
               op->getParentOp(), use.getSymbolRef(), resolvedSymbols))) {
         return use.getUser()->emitError()
                << "unable to resolve reference to symbol "
                << use.getSymbolRef();
       }
 
       // Mark each of the resolved symbols as live.
       for (Operation *resolvedSymbol : resolvedSymbols)