[ORC] Merge MaterializationResponsibility notifyEmitted and addDependencies

Removes the MaterializationResponsibility::addDependencies and
addDependenciesForAll methods, and transfers dependency registration to
the notifyEmitted operation. The new dependency registration allows
dependencies to be specified for arbitrary subsets of the
MaterializationResponsibility's symbols (rather than just single symbols
or all symbols) via an array of SymbolDependenceGroups (pairs of symbol
sets and corresponding dependencies for that set).

This patch aims to both improve emission performance and simplify
dependence tracking. By eliminating some states (e.g. symbols having
registered dependencies but not yet being resolved or emitted) we make
some errors impossible by construction, and reduce the number of error
cases that we need to check. NonOwningSymbolStringPtrs are used for
dependence tracking under the session lock, which should reduce
ref-counting operations, and intra-emit dependencies are resolved
outside the session lock, which should provide better performance when
JITing concurrently (since some dependence tracking can happen in
parallel).

The Orc C API is updated to account for this change, with the
LLVMOrcMaterializationResponsibilityNotifyEmitted API being modified and
the LLVMOrcMaterializationResponsibilityAddDependencies and
LLVMOrcMaterializationResponsibilityAddDependenciesForAll operations
being removed.

1 files changed, 73 insertions, 61 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index b828294..fffa95e 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ExecutionEngine/JITLink/EHFrameSupport.h"
 #include "llvm/ExecutionEngine/JITLink/aarch32.h"
 #include "llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h"
+#include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h"
 #include "llvm/Support/MemoryBuffer.h"
@@ -216,16 +217,13 @@ public:
       }
     };
 
-    for (auto &KV : InternalNamedSymbolDeps) {
-      SymbolDependenceMap InternalDeps;
-      InternalDeps[&MR->getTargetJITDylib()] = std::move(KV.second);
-      MR->addDependencies(KV.first, InternalDeps);
-    }
-
     ES.lookup(LookupKind::Static, LinkOrder, std::move(LookupSet),
               SymbolState::Resolved, std::move(OnResolve),
               [this](const SymbolDependenceMap &Deps) {
-                registerDependencies(Deps);
+                // Translate LookupDeps map to SymbolSourceJD.
+                for (auto &[DepJD, Deps] : Deps)
+                  for (auto &DepSym : Deps)
+                    SymbolSourceJDs[NonOwningSymbolStringPtr(DepSym)] = DepJD;
               });
   }
 
@@ -329,7 +327,7 @@ public:
       MR->failMaterialization();
       return;
     }
-    if (auto Err = MR->notifyEmitted()) {
+    if (auto Err = MR->notifyEmitted(SymbolDepGroups)) {
       Layer.getExecutionSession().reportError(std::move(Err));
       MR->failMaterialization();
     }
@@ -348,8 +346,8 @@ public:
 
     Layer.modifyPassConfig(*MR, LG, Config);
 
-    Config.PostPrunePasses.push_back(
-        [this](LinkGraph &G) { return computeNamedSymbolDependencies(G); });
+    Config.PreFixupPasses.push_back(
+        [this](LinkGraph &G) { return registerDependencies(G); });
 
     return Error::success();
   }
@@ -413,9 +411,10 @@ private:
       for (auto &E : B.edges()) {
         auto &Tgt = E.getTarget();
         if (Tgt.getScope() != Scope::Local) {
-          if (Tgt.isExternal())
-            BIDCacheVal.External.insert(getInternedName(Tgt));
-          else
+          if (Tgt.isExternal()) {
+            if (Tgt.getAddress() || !Tgt.isWeaklyReferenced())
+              BIDCacheVal.External.insert(getInternedName(Tgt));
+          } else
             BIDCacheVal.Internal.insert(getInternedName(Tgt));
         }
       }
@@ -481,58 +480,92 @@ private:
     return Error::success();
   }
 
-  Error computeNamedSymbolDependencies(LinkGraph &G) {
-    auto &ES = MR->getTargetJITDylib().getExecutionSession();
+  Error registerDependencies(LinkGraph &G) {
+    auto &TargetJD = MR->getTargetJITDylib();
+    auto &ES = TargetJD.getExecutionSession();
     auto BlockDeps = computeBlockNonLocalDeps(G);
 
+    DenseSet<Block *> BlockDepsProcessed;
+    DenseMap<Block *, SymbolDependenceGroup> DepGroupForBlock;
+
     // Compute dependencies for symbols defined in the JITLink graph.
     for (auto *Sym : G.defined_symbols()) {
 
-      // Skip local symbols: we do not track dependencies for these.
+      // Skip local symbols.
       if (Sym->getScope() == Scope::Local)
         continue;
       assert(Sym->hasName() &&
              "Defined non-local jitlink::Symbol should have a name");
 
-      auto &SymDeps = BlockDeps[Sym->getBlock()];
-      if (SymDeps.External.empty() && SymDeps.Internal.empty())
+      auto &BDeps = BlockDeps[Sym->getBlock()];
+
+      // Skip symbols in blocks that don't depend on anything.
+      if (BDeps.Internal.empty() && BDeps.External.empty())
         continue;
 
-      auto SymName = ES.intern(Sym->getName());
-      if (!SymDeps.External.empty())
-        ExternalNamedSymbolDeps[SymName] = SymDeps.External;
-      if (!SymDeps.Internal.empty())
-        InternalNamedSymbolDeps[SymName] = SymDeps.Internal;
+      SymbolDependenceGroup &SDG = DepGroupForBlock[&Sym->getBlock()];
+      SDG.Symbols.insert(ES.intern(Sym->getName()));
+
+      if (!BlockDepsProcessed.count(&Sym->getBlock())) {
+        BlockDepsProcessed.insert(&Sym->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);
+        }
+      }
     }
 
+    SymbolDependenceGroup SynthSDG;
+
     for (auto &P : Layer.Plugins) {
       auto SynthDeps = P->getSyntheticSymbolDependencies(*MR);
       if (SynthDeps.empty())
         continue;
 
       DenseSet<Block *> BlockVisited;
-      for (auto &KV : SynthDeps) {
-        auto &Name = KV.first;
-        auto &DepsForName = KV.second;
-        for (auto *Sym : DepsForName) {
+      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)
-              InternalNamedSymbolDeps[Name].insert(S);
-            for (auto &S : BDeps.External)
-              ExternalNamedSymbolDeps[Name].insert(S);
+              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 {
-            if (Sym->isExternal())
-              ExternalNamedSymbolDeps[Name].insert(
-                  BlockDeps.getInternedName(*Sym));
-            else
-              InternalNamedSymbolDeps[Name].insert(
-                  BlockDeps.getInternedName(*Sym));
+            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);
           }
         }
       }
     }
 
+    // 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));
+
     return Error::success();
   }
 
@@ -601,34 +634,13 @@ private:
                                 std::move(BlockDeps));
   }
 
-  void registerDependencies(const SymbolDependenceMap &QueryDeps) {
-    for (auto &NamedDepsEntry : ExternalNamedSymbolDeps) {
-      auto &Name = NamedDepsEntry.first;
-      auto &NameDeps = NamedDepsEntry.second;
-      SymbolDependenceMap SymbolDeps;
-
-      for (const auto &QueryDepsEntry : QueryDeps) {
-        JITDylib &SourceJD = *QueryDepsEntry.first;
-        const SymbolNameSet &Symbols = QueryDepsEntry.second;
-        auto &DepsForJD = SymbolDeps[&SourceJD];
-
-        for (const auto &S : Symbols)
-          if (NameDeps.count(S))
-            DepsForJD.insert(S);
-
-        if (DepsForJD.empty())
-          SymbolDeps.erase(&SourceJD);
-      }
-
-      MR->addDependencies(Name, SymbolDeps);
-    }
-  }
-
   ObjectLinkingLayer &Layer;
   std::unique_ptr<MaterializationResponsibility> MR;
   std::unique_ptr<MemoryBuffer> ObjBuffer;
-  DenseMap<SymbolStringPtr, SymbolNameSet> ExternalNamedSymbolDeps;
-  DenseMap<SymbolStringPtr, SymbolNameSet> InternalNamedSymbolDeps;
+  DenseMap<Block *, SymbolNameSet> ExternalBlockDeps;
+  DenseMap<Block *, SymbolNameSet> InternalBlockDeps;
+  DenseMap<NonOwningSymbolStringPtr, JITDylib *> SymbolSourceJDs;
+  std::vector<SymbolDependenceGroup> SymbolDepGroups;
 };
 
 ObjectLinkingLayer::Plugin::~Plugin() = default;