[ORC] Simplify intra-graph dependence tracking in ObjectLinkingLayer.

ObjectLinkingLayer::registerDependencies used to propagate external symbol
dependencies (dependencies on symbols outside the current graph) to all nodes.
Since ebe8733a11e, which merged addDependencies into notifyEmitted, the
notifyEmitted function will propagate intra-graph dependencies, so
registerDependencies no longer needs to do this.

This patch updates ObjectLinkingLayer::registerDependencies to just propagate
named dependencies (on both internal and external symbols) through anonymous
blocks, leaving the rest of the work to ExecutionSession::notifyEmitted.
It also choses a key symbol to use for blocks containing multiple symbols. The
result is both easier to read and faster.

1 files changed, 133 insertions, 203 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 073c25e..19d8307 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -14,8 +14,8 @@
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h"
 #include "llvm/Support/MemoryBuffer.h"
+
 #include <string>
-#include <vector>
 
 #define DEBUG_TYPE "orc"
 
@@ -330,6 +330,7 @@ public:
       MR->failMaterialization();
       return;
     }
+
     if (auto Err = MR->notifyEmitted(SymbolDepGroups)) {
       Layer.getExecutionSession().reportError(std::move(Err));
       MR->failMaterialization();
@@ -380,84 +381,6 @@ public:
   }
 
 private:
-  // Symbol name dependencies:
-  // Internal: Defined in this graph.
-  // External: Defined externally.
-  struct BlockSymbolDependencies {
-    SymbolNameSet Internal, External;
-  };
-
-  // Lazily populated map of blocks to BlockSymbolDependencies values.
-  class BlockDependenciesMap {
-  public:
-    BlockDependenciesMap(ExecutionSession &ES,
-                         DenseMap<const Block *, DenseSet<Block *>> BlockDeps)
-        : ES(ES), BlockDeps(std::move(BlockDeps)) {}
-
-    const BlockSymbolDependencies &operator[](const Block &B) {
-      // Check the cache first.
-      auto I = BlockTransitiveDepsCache.find(&B);
-      if (I != BlockTransitiveDepsCache.end())
-        return I->second;
-
-      // No value. Populate the cache.
-      BlockSymbolDependencies BTDCacheVal;
-      auto BDI = BlockDeps.find(&B);
-      assert(BDI != BlockDeps.end() && "No block dependencies");
-
-      for (auto *BDep : BDI->second) {
-        auto &BID = getBlockImmediateDeps(*BDep);
-        for (auto &ExternalDep : BID.External)
-          BTDCacheVal.External.insert(ExternalDep);
-        for (auto &InternalDep : BID.Internal)
-          BTDCacheVal.Internal.insert(InternalDep);
-      }
-
-      return BlockTransitiveDepsCache
-          .insert(std::make_pair(&B, std::move(BTDCacheVal)))
-          .first->second;
-    }
-
-    SymbolStringPtr &getInternedName(Symbol &Sym) {
-      auto I = NameCache.find(&Sym);
-      if (I != NameCache.end())
-        return I->second;
-
-      return NameCache.insert(std::make_pair(&Sym, ES.intern(Sym.getName())))
-          .first->second;
-    }
-
-  private:
-    BlockSymbolDependencies &getBlockImmediateDeps(Block &B) {
-      // Check the cache first.
-      auto I = BlockImmediateDepsCache.find(&B);
-      if (I != BlockImmediateDepsCache.end())
-        return I->second;
-
-      BlockSymbolDependencies BIDCacheVal;
-      for (auto &E : B.edges()) {
-        auto &Tgt = E.getTarget();
-        if (Tgt.getScope() != Scope::Local) {
-          if (Tgt.isExternal()) {
-            if (Tgt.getAddress() || !Tgt.isWeaklyReferenced())
-              BIDCacheVal.External.insert(getInternedName(Tgt));
-          } else
-            BIDCacheVal.Internal.insert(getInternedName(Tgt));
-        }
-      }
-
-      return BlockImmediateDepsCache
-          .insert(std::make_pair(&B, std::move(BIDCacheVal)))
-          .first->second;
-    }
-
-    ExecutionSession &ES;
-    DenseMap<const Block *, DenseSet<Block *>> BlockDeps;
-    DenseMap<const Symbol *, SymbolStringPtr> NameCache;
-    DenseMap<const Block *, BlockSymbolDependencies> BlockImmediateDepsCache;
-    DenseMap<const Block *, BlockSymbolDependencies> BlockTransitiveDepsCache;
-  };
-
   Error claimOrExternalizeWeakAndCommonSymbols(LinkGraph &G) {
     auto &ES = Layer.getExecutionSession();
 
@@ -509,168 +432,174 @@ private:
   }
 
   Error registerDependencies(LinkGraph &G) {
-    auto &TargetJD = MR->getTargetJITDylib();
-    auto &ES = TargetJD.getExecutionSession();
-    auto BlockDeps = computeBlockNonLocalDeps(G);
 
-    DenseSet<Block *> BlockDepsProcessed;
-    DenseMap<Block *, SymbolDependenceGroup> DepGroupForBlock;
+    struct BlockInfo {
+      bool InWorklist = false;
+      DenseSet<Symbol *> Defs;
+      DenseSet<Symbol *> SymbolDeps;
+      DenseSet<Block *> AnonEdges, AnonBackEdges;
+    };
 
-    // Compute dependencies for symbols defined in the JITLink graph.
+    DenseMap<Block *, BlockInfo> BlockInfos;
+
+    // Reserve space so that BlockInfos doesn't need to resize. This is
+    // essential to avoid invalidating pointers to entries below.
+    {
+      size_t NumBlocks = 0;
+      for (auto &Sec : G.sections())
+        NumBlocks += Sec.blocks_size();
+      BlockInfos.reserve(NumBlocks);
+    }
+
+    // Identify non-locally-scoped symbols defined by each block.
     for (auto *Sym : G.defined_symbols()) {
+      if (Sym->getScope() != Scope::Local)
+        BlockInfos[&Sym->getBlock()].Defs.insert(Sym);
+    }
 
-      // Skip local symbols.
-      if (Sym->getScope() == Scope::Local)
-        continue;
-      assert(Sym->hasName() &&
-             "Defined non-local jitlink::Symbol should have a name");
+    // Identify the symbolic and anonymous-block dependencies for each block.
+    for (auto *B : G.blocks()) {
+      auto &BI = BlockInfos[B];
 
-      auto &BDeps = BlockDeps[Sym->getBlock()];
+      for (auto &E : B->edges()) {
 
-      // Skip symbols in blocks that don't depend on anything.
-      if (BDeps.Internal.empty() && BDeps.External.empty())
-        continue;
+        // External symbols are trivially depended on.
+        if (E.getTarget().isExternal()) {
+          BI.SymbolDeps.insert(&E.getTarget());
+          continue;
+        }
 
-      SymbolDependenceGroup &SDG = DepGroupForBlock[&Sym->getBlock()];
-      SDG.Symbols.insert(ES.intern(Sym->getName()));
+        // Anonymous symbols aren't depended on at all (they're assumed to be
+        // already available).
+        if (E.getTarget().isAbsolute())
+          continue;
 
-      if (!BlockDepsProcessed.count(&Sym->getBlock())) {
-        BlockDepsProcessed.insert(&Sym->getBlock());
+        // If we get here then we depend on a symbol defined by some other
+        // block.
+        auto &TgtBI = BlockInfos[&E.getTarget().getBlock()];
 
-        if (!BDeps.Internal.empty())
-          SDG.Dependencies[&TargetJD] = BDeps.Internal;
-        for (auto &Dep : BDeps.External) {
-          auto DepSrcItr = SymbolSourceJDs.find(NonOwningSymbolStringPtr(Dep));
-          if (DepSrcItr != SymbolSourceJDs.end())
-            SDG.Dependencies[DepSrcItr->second].insert(Dep);
+        // If that block has any definitions then use the first one as the
+        // "effective" dependence here (all symbols in TgtBI will become
+        // ready at the same time, and chosing a single symbol to represent
+        // the block keeps the SymbolDepGroup size small).
+        if (!TgtBI.Defs.empty()) {
+          BI.SymbolDeps.insert(*TgtBI.Defs.begin());
+          continue;
         }
-      }
-    }
 
-    SymbolDependenceGroup SynthSDG;
-
-    for (auto &P : Plugins) {
-      auto SynthDeps = P->getSyntheticSymbolDependencies(*MR);
-      if (SynthDeps.empty())
-        continue;
-
-      DenseSet<Block *> BlockVisited;
-      for (auto &[Name, DepSyms] : SynthDeps) {
-        SynthSDG.Symbols.insert(Name);
-        for (auto *Sym : DepSyms) {
-          if (Sym->getScope() == Scope::Local) {
-            auto &BDeps = BlockDeps[Sym->getBlock()];
-            for (auto &S : BDeps.Internal)
-              SynthSDG.Dependencies[&TargetJD].insert(S);
-            for (auto &S : BDeps.External) {
-              auto DepSrcItr =
-                  SymbolSourceJDs.find(NonOwningSymbolStringPtr(S));
-              if (DepSrcItr != SymbolSourceJDs.end())
-                SynthSDG.Dependencies[DepSrcItr->second].insert(S);
-            }
-          } else {
-            auto SymName = ES.intern(Sym->getName());
-            if (Sym->isExternal()) {
-              assert(SymbolSourceJDs.count(NonOwningSymbolStringPtr(SymName)) &&
-                     "External symbol source entry missing");
-              SynthSDG
-                  .Dependencies[SymbolSourceJDs[NonOwningSymbolStringPtr(
-                      SymName)]]
-                  .insert(SymName);
-            } else
-              SynthSDG.Dependencies[&TargetJD].insert(SymName);
-          }
-        }
+        // Otherwise we've got a dependence on an anonymous block. Record it
+        // here for back-propagating symbol dependencies below.
+        BI.AnonEdges.insert(&E.getTarget().getBlock());
+        TgtBI.AnonBackEdges.insert(B);
       }
     }
 
-    // Transfer SDGs to SymbolDepGroups.
-    DepGroupForBlock.reserve(DepGroupForBlock.size() + 1);
-    for (auto &[B, SDG] : DepGroupForBlock) {
-      assert(!SDG.Symbols.empty() && "SymbolDependenceGroup covers no symbols");
-      if (!SDG.Dependencies.empty())
-        SymbolDepGroups.push_back(std::move(SDG));
-    }
-    if (!SynthSDG.Symbols.empty() && !SynthSDG.Dependencies.empty())
-      SymbolDepGroups.push_back(std::move(SynthSDG));
+    // Prune anonymous blocks.
+    {
+      std::vector<Block *> BlocksToRemove;
+      for (auto &[B, BI] : BlockInfos) {
+        // Skip blocks with defs. We only care about anonyous blocks.
+        if (!BI.Defs.empty())
+          continue;
 
-    return Error::success();
-  }
+        BlocksToRemove.push_back(B);
 
-  BlockDependenciesMap computeBlockNonLocalDeps(LinkGraph &G) {
-    // First calculate the reachable-via-non-local-symbol blocks for each block.
-    struct BlockInfo {
-      DenseSet<Block *> Dependencies;
-      DenseSet<Block *> Dependants;
-      bool DependenciesChanged = true;
-    };
-    DenseMap<Block *, BlockInfo> BlockInfos;
-    std::deque<Block *> WorkList;
+        for (auto *FB : BI.AnonEdges)
+          BlockInfos[FB].AnonBackEdges.erase(B);
 
-    // Pre-allocate map entries. This prevents any iterator/reference
-    // invalidation in the next loop.
-    for (auto *B : G.blocks())
-      (void)BlockInfos[B];
+        for (auto *BB : BI.AnonBackEdges)
+          BlockInfos[BB].AnonEdges.erase(B);
 
-    // Build initial worklist, record block dependencies/dependants and
-    // non-local symbol dependencies.
-    for (auto *B : G.blocks()) {
-      auto &BI = BlockInfos[B];
-      for (auto &E : B->edges()) {
-        if (E.getTarget().getScope() == Scope::Local &&
-            !E.getTarget().isAbsolute()) {
-          auto &TgtB = E.getTarget().getBlock();
-          if (&TgtB != B) {
-            BI.Dependencies.insert(&TgtB);
-            BlockInfos[&TgtB].Dependants.insert(B);
-          }
+        for (auto *FB : BI.AnonEdges) {
+          auto &FBI = BlockInfos[FB];
+          for (auto *BB : BI.AnonBackEdges)
+            FBI.AnonBackEdges.insert(BB);
+        }
+
+        for (auto *BB : BI.AnonBackEdges) {
+          auto &BBI = BlockInfos[BB];
+          for (auto *SD : BI.SymbolDeps)
+            BBI.SymbolDeps.insert(SD);
+          for (auto *FB : BI.AnonEdges)
+            BBI.AnonEdges.insert(FB);
         }
       }
+
+      for (auto *B : BlocksToRemove)
+        BlockInfos.erase(B);
     }
 
-    // Add blocks with both dependants and dependencies to the worklist to
-    // propagate dependencies to dependants.
+    // Build the initial dependence propagation worklist.
+    std::deque<Block *> Worklist;
     for (auto &[B, BI] : BlockInfos) {
-      if (!BI.Dependants.empty() && !BI.Dependencies.empty())
-        WorkList.push_back(B);
+      if (!BI.SymbolDeps.empty() && !BI.AnonBackEdges.empty()) {
+        Worklist.push_back(B);
+        BI.InWorklist = true;
+      }
     }
 
-    // Propagate block-level dependencies through the block-dependence graph.
-    while (!WorkList.empty()) {
-      auto *B = WorkList.back();
-      WorkList.pop_back();
+    // Propagate symbol deps through the graph.
+    while (!Worklist.empty()) {
+      auto *B = Worklist.front();
+      Worklist.pop_front();
 
       auto &BI = BlockInfos[B];
-      assert(BI.DependenciesChanged &&
-             "Block in worklist has unchanged dependencies");
-      BI.DependenciesChanged = false;
-      for (auto *Dependant : BI.Dependants) {
-        auto &DependantBI = BlockInfos[Dependant];
-        for (auto *Dependency : BI.Dependencies) {
-          if (Dependant != Dependency &&
-              DependantBI.Dependencies.insert(Dependency).second)
-            if (!DependantBI.DependenciesChanged) {
-              DependantBI.DependenciesChanged = true;
-              WorkList.push_front(Dependant);
-            }
+      BI.InWorklist = false;
+
+      for (auto *DB : BI.AnonBackEdges) {
+        auto &DBI = BlockInfos[DB];
+        for (auto *Sym : BI.SymbolDeps) {
+          if (DBI.SymbolDeps.insert(Sym).second && !DBI.InWorklist) {
+            Worklist.push_back(DB);
+            DBI.InWorklist = true;
+          }
         }
       }
     }
 
-    DenseMap<const Block *, DenseSet<Block *>> BlockDeps;
-    for (auto &KV : BlockInfos)
-      BlockDeps[KV.first] = std::move(KV.second.Dependencies);
+    // Transform our local dependence information into a list of
+    // SymbolDependenceGroups (in the SymbolDepGroups member), ready for use in
+    // the upcoming notifyFinalized call.
+    auto &TargetJD = MR->getTargetJITDylib();
+    auto &ES = TargetJD.getExecutionSession();
+
+    DenseMap<Symbol *, SymbolStringPtr> InternedNames;
+    auto GetInternedName = [&](Symbol *S) {
+      auto &Name = InternedNames[S];
+      if (!Name)
+        Name = ES.intern(S->getName());
+      return Name;
+    };
+
+    for (auto &[B, BI] : BlockInfos) {
+      if (!BI.Defs.empty()) {
+        SymbolDepGroups.push_back(SymbolDependenceGroup());
+        auto &SDG = SymbolDepGroups.back();
+
+        for (auto *Def : BI.Defs)
+          SDG.Symbols.insert(GetInternedName(Def));
+
+        for (auto *Dep : BI.SymbolDeps) {
+          auto DepName = GetInternedName(Dep);
+          if (Dep->isDefined())
+            SDG.Dependencies[&TargetJD].insert(std::move(DepName));
+          else {
+            auto SourceJDItr =
+                SymbolSourceJDs.find(NonOwningSymbolStringPtr(DepName));
+            if (SourceJDItr != SymbolSourceJDs.end())
+              SDG.Dependencies[SourceJDItr->second].insert(std::move(DepName));
+          }
+        }
+      }
+    }
 
-    return BlockDependenciesMap(Layer.getExecutionSession(),
-                                std::move(BlockDeps));
+    return Error::success();
   }
 
   ObjectLinkingLayer &Layer;
   std::vector<std::shared_ptr<ObjectLinkingLayer::Plugin>> Plugins;
   std::unique_ptr<MaterializationResponsibility> MR;
   std::unique_ptr<MemoryBuffer> ObjBuffer;
-  DenseMap<Block *, SymbolNameSet> ExternalBlockDeps;
-  DenseMap<Block *, SymbolNameSet> InternalBlockDeps;
   DenseMap<NonOwningSymbolStringPtr, JITDylib *> SymbolSourceJDs;
   std::vector<SymbolDependenceGroup> SymbolDepGroups;
 };
@@ -717,6 +646,7 @@ void ObjectLinkingLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
 
   auto Ctx = std::make_unique<ObjectLinkingLayerJITLinkContext>(
       *this, std::move(R), std::move(O));
+
   if (auto G = createLinkGraphFromObject(ObjBuffer)) {
     Ctx->notifyMaterializing(**G);
     link(std::move(*G), std::move(Ctx));