[mlir] Optimize symbol related checks in SymbolDCE

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.

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

1 files changed, 7 insertions, 3 deletions

diff --git a/mlir/lib/Transforms/SymbolDCE.cpp b/mlir/lib/Transforms/SymbolDCE.cpp
index 56997b6d2af7..5d94245489c0 100644
--- a/mlir/lib/Transforms/SymbolDCE.cpp
+++ b/mlir/lib/Transforms/SymbolDCE.cpp
@@ -24,6 +24,7 @@ struct SymbolDCE : public SymbolDCEBase<SymbolDCE> {
   /// `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);
 };
@@ -49,7 +50,9 @@ void SymbolDCE::runOnOperation() {
 
   // 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
@@ -71,6 +74,7 @@ void SymbolDCE::runOnOperation() {
 /// `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.
@@ -104,7 +108,7 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
       // 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();
     }
 
@@ -120,7 +124,7 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
     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 "