[mlir] Update SCCP and the Inliner to use SymbolTableCollection for symbol lookups

[rriddle]

Summary

This transforms the symbol lookups to O(1) from O(NM), greatly speeding up both passes. For a large MLIR module this shaved seconds off of the compilation time.

Depends On D89512

mlir/include/mlir/Analysis/CallGraph.h

 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class CallOpInterface;
 struct CallInterfaceCallable;
 class Operation;
 class Region;
+class SymbolTableCollection;
 
 //===----------------------------------------------------------------------===//
 // CallGraphNode
 //===----------------------------------------------------------------------===//
 
 /// This class represents a single callable in the callgraph. Aside from the
 /// external node, each node represents a callable node in the graph and
 /// contains a valid corresponding Region. The external node is a virtual node
public:
   CallGraphNode *lookupNode(Region *region) const;
 
   /// Return the callgraph node representing the indirect-external callee.
   CallGraphNode *getExternalNode() const {
     return const_cast<CallGraphNode *>(&externalNode);
   }
 
   /// Resolve the callable for given callee to a node in the callgraph, or the
-  /// external node if a valid node was not resolved.
-  CallGraphNode *resolveCallable(CallOpInterface call) const;
+  /// external node if a valid node was not resolved. The provided symbol table
+  /// is used when resolving calls that reference callables via a symbol
+  /// reference.
+  CallGraphNode *resolveCallable(CallOpInterface call,
+                                 SymbolTableCollection &symbolTable) const;
 
   /// Erase the given node from the callgraph.
   void eraseNode(CallGraphNode *node);
 
   /// An iterator over the nodes of the graph.
   using iterator = NodeIterator;
   iterator begin() const { return nodes.begin(); }
   iterator end() const { return nodes.end(); }

mlir/include/mlir/Interfaces/CallInterfaces.td

InterfaceMethod<[{
       "CallInterfaceCallable", "getCallableForCallee"
     >,
     InterfaceMethod<[{
         Returns the operands within this call that are used as arguments to the
         callee.
       }],
       "Operation::operand_range", "getArgOperands"
     >,
-    InterfaceMethod<[{
-        Resolve the callable operation for given callee to a
-        CallableOpInterface, or nullptr if a valid callable was not resolved.
-      }],
-      "Operation *", "resolveCallable", (ins), [{
-        // If the callable isn't a value, lookup the symbol reference.
-        CallInterfaceCallable callable = $_op.getCallableForCallee();
-        if (auto symbolRef = callable.dyn_cast<SymbolRefAttr>())
-          return SymbolTable::lookupNearestSymbolFrom($_op, symbolRef);
-        return callable.get<Value>().getDefiningOp();
-      }]
-    >,
   ];
+
+  let extraClassDeclaration = [{
+    /// Resolve the callable operation for given callee to a
+    /// CallableOpInterface, or nullptr if a valid callable was not resolved.
+    /// `symbolTable` is an optional parameter that will allow for using a
+    /// cached symbol table for symbol lookups instead of performing an O(N)
+    /// scan.
+    Operation *resolveCallable(SymbolTableCollection *symbolTable = nullptr);
+  }];
 }
 
 /// Interface for callable operations.
 def CallableOpInterface : OpInterface<"CallableOpInterface"> {
   let description = [{
     A callable operation is one who represents a potential sub-routine, and may
     be a target for a call-like operation (those providing the CallOpInterface
     above). These operations may be traditional functional operation

mlir/lib/Analysis/CallGraph.cpp

 
 //===----------------------------------------------------------------------===//
 // CallGraph
 //===----------------------------------------------------------------------===//
 
 /// Recursively compute the callgraph edges for the given operation. Computed
 /// edges are placed into the given callgraph object.
 static void computeCallGraph(Operation *op, CallGraph &cg,
+                             SymbolTableCollection &symbolTable,
                              CallGraphNode *parentNode, bool resolveCalls) {
   if (CallOpInterface call = dyn_cast<CallOpInterface>(op)) {
     // If there is no parent node, we ignore this operation. Even if this
     // operation was a call, there would be no callgraph node to attribute it
     // to.
     if (resolveCalls && parentNode)
-      parentNode->addCallEdge(cg.resolveCallable(call));
+      parentNode->addCallEdge(cg.resolveCallable(call, symbolTable));
     return;
   }
 
   // Compute the callgraph nodes and edges for each of the nested operations.
   if (CallableOpInterface callable = dyn_cast<CallableOpInterface>(op)) {
     if (auto *callableRegion = callable.getCallableRegion())
       parentNode = cg.getOrAddNode(callableRegion, parentNode);
     else
       return;
   }
 
   for (Region &region : op->getRegions())
     for (Operation &nested : region.getOps())
-      computeCallGraph(&nested, cg, parentNode, resolveCalls);
+      computeCallGraph(&nested, cg, symbolTable, parentNode, resolveCalls);
 }
 
 CallGraph::CallGraph(Operation *op) : externalNode(/*callableRegion=*/nullptr) {
   // Make two passes over the graph, one to compute the callables and one to
   // resolve the calls. We split these up as we may have nested callable objects
   // that need to be reserved before the calls.
-  computeCallGraph(op, *this, /*parentNode=*/nullptr, /*resolveCalls=*/false);
-  computeCallGraph(op, *this, /*parentNode=*/nullptr, /*resolveCalls=*/true);
+  SymbolTableCollection symbolTable;
+  computeCallGraph(op, *this, symbolTable, /*parentNode=*/nullptr,
+                   /*resolveCalls=*/false);
+  computeCallGraph(op, *this, symbolTable, /*parentNode=*/nullptr,
+                   /*resolveCalls=*/true);
 }
 
 /// Get or add a call graph node for the given region.
 CallGraphNode *CallGraph::getOrAddNode(Region *region,
                                        CallGraphNode *parentNode) {
   assert(region && isa<CallableOpInterface>(region->getParentOp()) &&
          "expected parent operation to be callable");
   std::unique_ptr<CallGraphNode> &node = nodes[region];
   if (!node) {
     node.reset(new CallGraphNode(region));
 
     // Add this node to the given parent node if necessary.
-    if (parentNode)
+    if (parentNode) {
       parentNode->addChildEdge(node.get());
-    else
+    } else {
       // Otherwise, connect all callable nodes to the external node, this allows
       // for conservatively including all callable nodes within the graph.
-      // FIXME(riverriddle) This isn't correct, this is only necessary for
-      // callable nodes that *could* be called from external sources. This
-      // requires extending the interface for callables to check if they may be
-      // referenced externally.
+      // FIXME This isn't correct, this is only necessary for callable nodes
+      // that *could* be called from external sources. This requires extending
+      // the interface for callables to check if they may be referenced
+      // externally.
       externalNode.addAbstractEdge(node.get());
     }
+  }
   return node.get();
 }
 
 /// Lookup a call graph node for the given region, or nullptr if none is
 /// registered.
 CallGraphNode *CallGraph::lookupNode(Region *region) const {
   auto it = nodes.find(region);
   return it == nodes.end() ? nullptr : it->second.get();
 }
 
 /// Resolve the callable for given callee to a node in the callgraph, or the
 /// external node if a valid node was not resolved.
-CallGraphNode *CallGraph::resolveCallable(CallOpInterface call) const {
-  Operation *callable = call.resolveCallable();
+CallGraphNode *
+CallGraph::resolveCallable(CallOpInterface call,
+                           SymbolTableCollection &symbolTable) const {
+  Operation *callable = call.resolveCallable(&symbolTable);
   if (auto callableOp = dyn_cast_or_null<CallableOpInterface>(callable))
     if (auto *node = lookupNode(callableOp.getCallableRegion()))
       return node;
 
   // If we don't have a valid direct region, this is an external call.
   return getExternalNode();
 }
 

mlir/lib/Interfaces/CallInterfaces.cpp

 //===- CallInterfaces.cpp - ControlFlow Interfaces ------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/CallInterfaces.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
+// CallOpInterface
+//===----------------------------------------------------------------------===//
+
+/// Resolve the callable operation for given callee to a CallableOpInterface, or
+/// nullptr if a valid callable was not resolved. `symbolTable` is an optional
+/// parameter that will allow for using a cached symbol table for symbol lookups
+/// instead of performing an O(N) scan.
+Operation *
+CallOpInterface::resolveCallable(SymbolTableCollection *symbolTable) {
+  CallInterfaceCallable callable = getCallableForCallee();
+  if (auto symbolVal = callable.dyn_cast<Value>())
+    return symbolVal.getDefiningOp();
+
+  // If the callable isn't a value, lookup the symbol reference.
+  auto symbolRef = callable.get<SymbolRefAttr>();
+  if (symbolTable)
+    return symbolTable->lookupNearestSymbolFrom(getOperation(), symbolRef);
+  return SymbolTable::lookupNearestSymbolFrom(getOperation(), symbolRef);
+}
+
+//===----------------------------------------------------------------------===//
 // CallInterfaces
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/CallInterfaces.cpp.inc"

mlir/lib/Transforms/Inliner.cpp

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Tracking
 //===----------------------------------------------------------------------===//
 
 /// Walk all of the used symbol callgraph nodes referenced with the given op.
 static void walkReferencedSymbolNodes(
-    Operation *op, CallGraph &cg,
+    Operation *op, CallGraph &cg, SymbolTableCollection &symbolTable,
     DenseMap<Attribute, CallGraphNode *> &resolvedRefs,
     function_ref<void(CallGraphNode *, Operation *)> callback) {
   auto symbolUses = SymbolTable::getSymbolUses(op);
   assert(symbolUses && "expected uses to be valid");
 
   Operation *symbolTableOp = op->getParentOp();
   for (const SymbolTable::SymbolUse &use : *symbolUses) {
     auto refIt = resolvedRefs.insert({use.getSymbolRef(), nullptr});
     CallGraphNode *&node = refIt.first->second;
 
     // If this is the first instance of this reference, try to resolve a
     // callgraph node for it.
     if (refIt.second) {
-      auto *symbolOp = SymbolTable::lookupNearestSymbolFrom(symbolTableOp,
-                                                            use.getSymbolRef());
+      auto *symbolOp = symbolTable.lookupNearestSymbolFrom(symbolTableOp,
       auto callableOp = dyn_cast_or_null<CallableOpInterface>(symbolOp);
       if (!callableOp)
         continue;
       node = cg.lookupNode(callableOp.getCallableRegion());
     }
     if (node)
       callback(node, use.getUser());
   }
struct CGUser {
     /// Any nodes referenced in the top-level attribute list of this user. We
     /// use a set here because the number of references does not matter.
     DenseSet<CallGraphNode *> topLevelUses;
 
     /// Uses of nodes referenced by nested operations.
     DenseMap<CallGraphNode *, int> innerUses;
   };
 
-  CGUseList(Operation *op, CallGraph &cg);
+  CGUseList(Operation *op, CallGraph &cg, SymbolTableCollection &symbolTable);
 
   /// Drop uses of nodes referred to by the given call operation that resides
   /// within 'userNode'.
   void dropCallUses(CallGraphNode *userNode, Operation *callOp, CallGraph &cg);
 
   /// Remove the given node from the use list.
   void eraseNode(CallGraphNode *node);
 

 
 private:
   /// Decrement the uses of discardable nodes referenced by the given user.
   void decrementDiscardableUses(CGUser &uses);
 
   /// A mapping between a discardable callgraph node (that is a symbol) and the
   /// number of uses for this node.
   DenseMap<CallGraphNode *, int> discardableSymNodeUses;
+
   /// A mapping between a callgraph node and the symbol callgraph nodes that it
   /// uses.
   DenseMap<CallGraphNode *, CGUser> nodeUses;
+
+  /// A symbol table to use when resolving call lookups.
+  SymbolTableCollection &symbolTable;
 };
 } // end anonymous namespace
 
-CGUseList::CGUseList(Operation *op, CallGraph &cg) {
+CGUseList::CGUseList(Operation *op, CallGraph &cg,
+                     SymbolTableCollection &symbolTable)
+    : symbolTable(symbolTable) {
   /// A set of callgraph nodes that are always known to be live during inlining.
   DenseMap<Attribute, CallGraphNode *> alwaysLiveNodes;
 
   // Walk each of the symbol tables looking for discardable callgraph nodes.
   auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {
     for (Operation &op : symbolTableOp->getRegion(0).getOps()) {
       // If this is a callgraph operation, check to see if it is discardable.
       if (auto callable = dyn_cast<CallableOpInterface>(&op)) {
         if (auto *node = cg.lookupNode(callable.getCallableRegion())) {
           SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(&op);
           if (symbol && (allUsesVisible || symbol.isPrivate()) &&
               symbol.canDiscardOnUseEmpty()) {
             discardableSymNodeUses.try_emplace(node, 0);
           }
           continue;
         }
       }
       // Otherwise, check for any referenced nodes. These will be always-live.
-      walkReferencedSymbolNodes(&op, cg, alwaysLiveNodes,
+      walkReferencedSymbolNodes(&op, cg, symbolTable, alwaysLiveNodes,
                                 [](CallGraphNode *, Operation *) {});
     }
   };
   SymbolTable::walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(),
                                 walkFn);
 
   // Drop the use information for any discardable nodes that are always live.
   for (auto &it : alwaysLiveNodes)
void CGUseList::dropCallUses(CallGraphNode *userNode, Operation *callOp,
   auto walkFn = [&](CallGraphNode *node, Operation *user) {
     auto parentIt = userRefs.find(node);
     if (parentIt == userRefs.end())
       return;
     --parentIt->second;
     --discardableSymNodeUses[node];
   };
   DenseMap<Attribute, CallGraphNode *> resolvedRefs;
-  walkReferencedSymbolNodes(callOp, cg, resolvedRefs, walkFn);
+  walkReferencedSymbolNodes(callOp, cg, symbolTable, resolvedRefs, walkFn);
 }
 
 void CGUseList::eraseNode(CallGraphNode *node) {
   // Drop all child nodes.
   for (auto &edge : *node)
     if (edge.isChild())
       eraseNode(edge.getTarget());
 
if (discardSymIt == discardableSymNodeUses.end())
       return;
 
     if (user != parentOp)
       ++uses.innerUses[refNode];
     else if (!uses.topLevelUses.insert(refNode).second)
       return;
     ++discardSymIt->second;
   };
-  walkReferencedSymbolNodes(parentOp, cg, resolvedRefs, walkFn);
+  walkReferencedSymbolNodes(parentOp, cg, symbolTable, resolvedRefs, walkFn);
 }
 
 void CGUseList::mergeUsesAfterInlining(CallGraphNode *lhs, CallGraphNode *rhs) {
   auto &lhsUses = nodeUses[lhs], &rhsUses = nodeUses[rhs];
   for (auto &useIt : lhsUses.innerUses) {
     rhsUses.innerUses[useIt.first] += useIt.second;
     discardableSymNodeUses[useIt.first] += useIt.second;
   }

 };
 } // end anonymous namespace
 
 /// Collect all of the callable operations within the given range of blocks. If
 /// `traverseNestedCGNodes` is true, this will also collect call operations
 /// inside of nested callgraph nodes.
 static void collectCallOps(iterator_range<Region::iterator> blocks,
                            CallGraphNode *sourceNode, CallGraph &cg,
+                           SymbolTableCollection &symbolTable,
                            SmallVectorImpl<ResolvedCall> &calls,
                            bool traverseNestedCGNodes) {
   SmallVector<std::pair<Block *, CallGraphNode *>, 8> worklist;
   auto addToWorklist = [&](CallGraphNode *node,
                            iterator_range<Region::iterator> blocks) {
     for (Block &block : blocks)
       worklist.emplace_back(&block, node);
   };
 
   addToWorklist(sourceNode, blocks);
   while (!worklist.empty()) {
     Block *block;
     std::tie(block, sourceNode) = worklist.pop_back_val();
 
     for (Operation &op : *block) {
       if (auto call = dyn_cast<CallOpInterface>(op)) {
         // TODO: Support inlining nested call references.
         CallInterfaceCallable callable = call.getCallableForCallee();
         if (SymbolRefAttr symRef = callable.dyn_cast<SymbolRefAttr>()) {
           if (!symRef.isa<FlatSymbolRefAttr>())
             continue;
         }
 
-        CallGraphNode *targetNode = cg.resolveCallable(call);
+        CallGraphNode *targetNode = cg.resolveCallable(call, symbolTable);
         if (!targetNode->isExternal())
           calls.emplace_back(call, sourceNode, targetNode);
         continue;
       }
 
       // If this is not a call, traverse the nested regions. If
       // `traverseNestedCGNodes` is false, then don't traverse nested call graph
       // regions.
       for (auto &nestedRegion : op.getRegions()) {
         CallGraphNode *nestedNode = cg.lookupNode(&nestedRegion);
         if (traverseNestedCGNodes || !nestedNode)
           addToWorklist(nestedNode ? nestedNode : sourceNode, nestedRegion);
       }
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Inliner
 //===----------------------------------------------------------------------===//
 namespace {
 /// This class provides a specialization of the main inlining interface.
 struct Inliner : public InlinerInterface {
-  Inliner(MLIRContext *context, CallGraph &cg)
-      : InlinerInterface(context), cg(cg) {}
+  Inliner(MLIRContext *context, CallGraph &cg,
+          SymbolTableCollection &symbolTable)
+      : InlinerInterface(context), cg(cg), symbolTable(symbolTable) {}
 
   /// Process a set of blocks that have been inlined. This callback is invoked
   /// *before* inlined terminator operations have been processed.
   void
   processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks) final {
     // Find the closest callgraph node from the first block.
     CallGraphNode *node;
     Region *region = inlinedBlocks.begin()->getParent();
     while (!(node = cg.lookupNode(region))) {
       region = region->getParentRegion();
       assert(region && "expected valid parent node");
     }
 
-    collectCallOps(inlinedBlocks, node, cg, calls,
+    collectCallOps(inlinedBlocks, node, cg, symbolTable, calls,
                    /*traverseNestedCGNodes=*/true);
   }
 
   /// Mark the given callgraph node for deletion.
   void markForDeletion(CallGraphNode *node) { deadNodes.insert(node); }
 
   /// This method properly disposes of callables that became dead during
   /// inlining. This should not be called while iterating over the SCCs.
   void eraseDeadCallables() {
     for (CallGraphNode *node : deadNodes)
       node->getCallableRegion()->getParentOp()->erase();
   }
 
   /// The set of callables known to be dead.
   SmallPtrSet<CallGraphNode *, 8> deadNodes;
 
   /// The current set of call instructions to consider for inlining.
   SmallVector<ResolvedCall, 8> calls;
 
   /// The callgraph being operated on.
   CallGraph &cg;
+
+  /// A symbol table to use when resolving call lookups.
+  SymbolTableCollection &symbolTable;
 };
 } // namespace
 
 /// Returns true if the given call should be inlined.
 static bool shouldInline(ResolvedCall &resolvedCall) {
   // Don't allow inlining terminator calls. We currently don't support this
   // case.
   if (resolvedCall.call.getOperation()->isKnownTerminator())
static LogicalResult inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
   // Collect all of the direct calls within the nodes of the current SCC. We
   // don't traverse nested callgraph nodes, because they are handled separately
   // likely within a different SCC.
   for (CallGraphNode *node : currentSCC) {
     if (node->isExternal())
       continue;
 
     // Don't collect calls if the node is already dead.
-    if (useList.isDead(node))
+    if (useList.isDead(node)) {
       deadNodes.push_back(node);
-    else
-      collectCallOps(*node->getCallableRegion(), node, cg, calls,
-                     /*traverseNestedCGNodes=*/false);
+    } else {
+      collectCallOps(*node->getCallableRegion(), node, cg, inliner.symbolTable,
+                     calls, /*traverseNestedCGNodes=*/false);
+    }
   }
 
   // Try to inline each of the call operations. Don't cache the end iterator
   // here as more calls may be added during inlining.
   bool inlinedAnyCalls = false;
   for (unsigned i = 0; i != calls.size(); ++i) {
     ResolvedCall it = calls[i];
     bool doInline = shouldInline(it);
void InlinerPass::runOnOperation() {
 
   // Collect a set of canonicalization patterns to use when simplifying
   // callable regions within an SCC.
   OwningRewritePatternList canonPatterns;
   for (auto *op : context->getRegisteredOperations())
     op->getCanonicalizationPatterns(canonPatterns, context);
 
   // Run the inline transform in post-order over the SCCs in the callgraph.
-  Inliner inliner(context, cg);
-  CGUseList useList(getOperation(), cg);
+  SymbolTableCollection symbolTable;
+  Inliner inliner(context, cg, symbolTable);
+  CGUseList useList(getOperation(), cg, symbolTable);
   runTransformOnCGSCCs(cg, [&](CallGraphSCC &scc) {
     inlineSCC(inliner, useList, scc, context, canonPatterns);
   });
 
   // After inlining, make sure to erase any callables proven to be dead.
   inliner.eraseDeadCallables();
 }
 

mlir/lib/Transforms/SCCP.cpp

private:
 
   /// The callable operations that have their argument/result state tracked.
   DenseMap<Operation *, CallableLatticeState> callableLatticeState;
 
   /// A map between a call operation and the resolved symbol callable. This
   /// avoids re-resolving symbol references during propagation. Value based
   /// callables are trivial to resolve, so they can be done in-place.
   DenseMap<Operation *, Operation *> callToSymbolCallable;
+
+  /// A symbol table used for O(1) symbol lookups during simplification.
+  SymbolTableCollection symbolTable;
 };
 } // end anonymous namespace
 
 SCCPSolver::SCCPSolver(Operation *op) {
   /// Initialize the solver with the regions within this operation.
   for (Region &region : op->getRegions()) {
     if (region.empty())
       continue;
if (!uses) {
       op->walk([&](CallableOpInterface op) { callableLatticeState.erase(op); });
       return;
     }
 
     for (const SymbolTable::SymbolUse &use : *uses) {
       // If the use is a call, track it to avoid the need to recompute the
       // reference later.
       if (auto callOp = dyn_cast<CallOpInterface>(use.getUser())) {
-        Operation *symCallable = callOp.resolveCallable();
+        Operation *symCallable = callOp.resolveCallable(&symbolTable);
         auto callableLatticeIt = callableLatticeState.find(symCallable);
         if (callableLatticeIt != callableLatticeState.end()) {
           callToSymbolCallable.try_emplace(callOp, symCallable);
 
           // We only need to record the call in the lattice if it produces any
           // values.
           if (callOp.getOperation()->getNumResults())
             callableLatticeIt->second.addSymbolCall(callOp);
         }
         continue;
       }
       // This use isn't a call, so don't we know all of the callers.
-      auto *symbol = SymbolTable::lookupSymbolIn(op, use.getSymbolRef());
+      auto *symbol = symbolTable.lookupSymbolIn(op, use.getSymbolRef());
       auto it = callableLatticeState.find(symbol);
       if (it != callableLatticeState.end())
         markAllOverdefined(it->second.getCallableArguments());
     }
   };
   SymbolTable::walkSymbolTables(op, /*allSymUsesVisible=*/!op->getBlock(),
                                 walkFn);
 }