[ORC] Align weak definition handling in JITDylib::defineMaterializing() and JITDylib::defineImpl()

[sgraenitz]

Summary

In both functions handle all cases for symbol name collisions:

(1) The incoming symbol is weak: reject the incoming symbol definition
(2) The existing symbol is weak (and not materialized!): discard the existing symbol definition and insert the incoming one
(3) Both symbol definitions are strong: DuplicateDefinition error

This behavior would be very handy for ThinLtoJIT, because it (currently) uses weak definitions for fallback handlers. One example: let's say we have two modules, main.bc and rare.bc. The latter defines a function rarely() that is called from main.bc, so it will be resolved as a transitive dependency when linking main.bc even though it might never be called (or not even be reachable from our entry point!).

ThinLtoJIT will try to minimize the immediate/synchronous workload and bet on discovery to identify and emit all necessary code asynchronously. In many cases this may work, but it still needs a fallback in case rarely() gets called and discovery didn't manage to provide rare.bc in time. For this case it emits a synthetic call-through stub that does the work (load and emit the module). The symbol it defines for rarely() in the shadow dylib would then simply be weak, so that discovery can override it with the actual definition from rare.bc once it's ready.

[lhames]

ThinLtoJIT will try to minimize the immediate/synchronous workload and bet on discovery to identify and emit all necessary code asynchronously. In many cases this may work, but it still needs a fallback in case rarely() gets called and discovery didn't manage to provide rare.bc in time. For this case it emits a synthetic call-through stub that does the work (load and emit the module). The symbol it defines for rarely() in the shadow dylib would then simply be weak, so that discovery can override it with the actual definition from rare.bc once it's ready.

Is the idea here that the call-through stub is emitted as weak, and then replaced by discovery with a strong definition if one is found in time? Or have I misunderstood?

[sgraenitz]

In D75090#1904569, @lhames wrote:

Is the idea here that the call-through stub is emitted as weak, and then replaced by discovery with a strong definition if one is found in time? Or have I misunderstood?

Yes that's the idea. (Though the actual definition found by discovery might be weak itself -- but I don't handle this yet.) I think the cases where I ran into this are actually race conditions in my code, where multiple threads compete for symbol table access. While I will have to revisit this part, I wondered whether it may be worth having an aligned mechanism for both cases. Something like the (1), (2), (3) described above?

llvm/examples/ThinLtoJIT/OnDemandLayer.cpp

#ifndef NDEBUG
 
   JITDylib &JD = R.getTargetJITDylib();
   LLVM_DEBUG(dbgs() << "[" << ModulePath << "] Adding symbols ("
                     << JD.getName() << "): " << ModuleSymbols << "\n");
   LLVM_DEBUG(dbgs() << "[" << ModulePath << "] Existing symbols: "
                     << R.getSymbols() << "\n");
 #endif
 
-  // TODO: Take responsibility for all symbols in the module. They should
-  // replace all synthetic weak definitions inherited from the
-  // EmitCallThroughStubs unit.
+  // Take responsibility for all symbols in the module. In the JITDylib, this
+  // will replace the synthetic weak definitions inherited from the
+  // EmitCallThroughStubs unit with the definitions parsed from the module.
+  //
+  // Note: The MaterializationResponsibility will keep the old entries, because
+  //       it implicitly drops duplicates upon insert().
+  //
+  // TODO: What happens if the module naturally contains weak definitions?
+  //       So far defineMaterializing() would reject the incoming definition
+  //       and incorrectly keep the existing one. It seems acceptable for now as
+  //       we only replace their flags values under the hood. If we started
+  //       recording other properties for symbols, this would probably break.
+  //       Introducing a 'synthetic' flag may help to fix the ambiguity.
+  //
+  if (Error Err = R.defineMaterializing(std::move(ModuleSymbols))) {
+    Parent.getExecutionSession().reportError(std::move(Err));
+    return;
+  }
 
   static constexpr bool SubmittedAsynchronously = false;
   Parent.emitWholeModule(std::move(R), std::move(TSM), SubmittedAsynchronously);
 }
 
 class EmitCallThroughStubs : public NamedMaterializationUnit {
 public:
   EmitCallThroughStubs(SymbolFlagsMap SymbolFlags, StringRef ModulePath,
void CompileWholeModuleOnDemandLayer::emitStubsOnly(
   SymbolAliasMap Callables;
   for (auto &KV : R.getSymbols()) {
     SymbolStringPtr Name = KV.first;
     JITSymbolFlags Flags = KV.second;
     assert(Flags.isCallable() && Flags.isWeak());
     Callables[Name] = SymbolAliasMapEntry(Name, Flags);
   }
 
-  // TODO: Create a materialization unit that will load and emit the module. The
-  // reexports that we generate in the end will issue a query on call-through,
-  // which should trigger the actual materialization.
+  // Create a materialization unit that will load and emit the whole module and
+  // lodge it with the shadow dylib. The reexports that we generate in the end
+  // will issue a query into the shadow dylib on call-through. This will trigger
+  // the actual materialization.
+  JITDylibResources &JDR = getJITDylibResources(R.getTargetJITDylib());
+  if (auto Err = JDR.ShadowJD->define(
+          std::make_unique<LoadAndEmitWholeModule>(
+              R.getSymbols(), ModulePath, R.getVModuleKey(), *this))) {
+    getExecutionSession().reportError(std::move(Err));
+    R.failMaterialization();
+    return;
+  }
 
   // No need to track function aliasees.
   constexpr ImplSymbolMap *NoAliaseeImpls = nullptr;
 
   // Emit call-through symbols to main dylib.
-  JITDylibResources &JDR = getJITDylibResources(R.getTargetJITDylib());
   R.replace(lazyReexports(*CallThroughManager, *JDR.ISManager, *JDR.ShadowJD,
                           std::move(Callables), NoAliaseeImpls));
 
   PSTATS("[", ModulePath , "] Done emitting stubs-only");
 }
 
 void CompileWholeModuleOnDemandLayer::emitReexports(
     MaterializationResponsibility R, ThreadSafeModule TSM) {
for (auto &KV : R.getSymbols()) {
     JITSymbolFlags Flags = KV.second;
 
     if (Flags.isCallable())
       Callables[Name] = SymbolAliasMapEntry(Name, Flags | JITSymbolFlags::Weak);
     else
       NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
   }
 
-  // TODO: Create a materialization unit that will emit the whole module. The
-  // reexports that we generate in the end will issue a query on call-through,
-  // which should trigger the actual materialization.
+  // Create a materialization unit that will emit the whole module and lodge it
+  // with the shadow dylib. The reexports that we generate in the end will issue
+  // a query into the shadow dylib on call-through. This will trigger the actual
+  // materialization.
+  JITDylib &JD = R.getTargetJITDylib();
+  static constexpr bool SubmittedFromDiscovery = false;
+  if (auto Err = addWholeModule(JD, std::move(TSM), R.getVModuleKey(),
+                                SubmittedFromDiscovery)) {
+    getExecutionSession().reportError(std::move(Err));
+    R.failMaterialization();
+    return;
+  }
 
   PSTATS("[", ModulePath , "] Start emitting reexports");
 
   // No need to track function aliasees.
   constexpr ImplSymbolMap *NoAliaseeImpls = nullptr;
-  JITDylibResources &JDR = getJITDylibResources(R.getTargetJITDylib());
+  JITDylibResources &JDR = getJITDylibResources(JD);
 
   R.replace(reexports(*JDR.ShadowJD, std::move(NonCallables),
                       JITDylibLookupFlags::MatchAllSymbols));
   R.replace(lazyReexports(*CallThroughManager, *JDR.ISManager, *JDR.ShadowJD,
                           std::move(Callables), NoAliaseeImpls));
 
   PSTATS("[", ModulePath , "] Done emitting reexports");
 }

llvm/lib/ExecutionEngine/Orc/Core.cpp

for (auto SFItr = SymbolFlags.begin(), SFEnd = SymbolFlags.end();
 
       auto &Name = SFItr->first;
       auto &Flags = SFItr->second;
 
       auto EntryItr = Symbols.find(Name);
 
       // If the entry already exists...
       if (EntryItr != Symbols.end()) {
+        if (Flags.isWeak()) {
+          // The incoming symbol definition is weak.
+          RejectedWeakDefs.push_back(SFItr);
+          continue;
+        }
 
-        // If this is a strong definition then error out.
-        if (!Flags.isWeak()) {
-          // Remove any symbols already added.
+        if (EntryItr->second.getFlags().isWeak()) {
+          // The existing symbol definition is weak.
+          if (!EntryItr->second.isInMaterializationPhase()) {
+            auto UMII = UnmaterializedInfos.find(Name);
+            assert(UMII != UnmaterializedInfos.end() &&
+                  "Overridden existing def should have an UnmaterializedInfo");
+            UMII->second->MU->doDiscard(*this, Name);
+          }
+          Symbols.erase(EntryItr);
+        } else {
+          // Both definitions are strong. Remove any symbols already added.
           for (auto &SI : AddedSyms)
             Symbols.erase(SI);
 
           // FIXME: Return all duplicates.
           return make_error<DuplicateDefinition>(std::string(*Name));
         }
+      }
 
-        // Otherwise just make a note to discard this symbol after the loop.
-        RejectedWeakDefs.push_back(SFItr);
-        continue;
-      } else
-        EntryItr =
-          Symbols.insert(std::make_pair(Name, SymbolTableEntry(Flags))).first;
 
       AddedSyms.push_back(EntryItr);
       EntryItr->second.setState(SymbolState::Materializing);
     }
 
     // Remove any rejected weak definitions from the SymbolFlags map.
     while (!RejectedWeakDefs.empty()) {
       SymbolFlags.erase(RejectedWeakDefs.back());

 JITDylib::JITDylib(ExecutionSession &ES, std::string Name)
     : ES(ES), JITDylibName(std::move(Name)) {
   SearchOrder.push_back({this, JITDylibLookupFlags::MatchAllSymbols});
 }
 
 Error JITDylib::defineImpl(MaterializationUnit &MU) {
   SymbolNameSet Duplicates;
   std::vector<SymbolStringPtr> ExistingDefsOverridden;
-  std::vector<SymbolStringPtr> MUDefsOverridden;
+  std::vector<SymbolStringPtr> RejectedWeakDefs;
 
   for (const auto &KV : MU.getSymbols()) {
     auto I = Symbols.find(KV.first);
 
     if (I != Symbols.end()) {
-      if (KV.second.isStrong()) {
-        if (I->second.getFlags().isStrong() ||
-            I->second.getState() > SymbolState::NeverSearched)
-          Duplicates.insert(KV.first);
-        else {
-          assert(I->second.getState() == SymbolState::NeverSearched &&
-                 "Overridden existing def should be in the never-searched "
-                 "state");
-          ExistingDefsOverridden.push_back(KV.first);
-        }
-      } else
-        MUDefsOverridden.push_back(KV.first);
+      if (KV.second.isWeak()) {
+        // The incoming symbol definition is weak.
+        RejectedWeakDefs.push_back(KV.first);
+        continue;
+      } else if (I->second.getFlags().isWeak() &&
+                 I->second.getState() == SymbolState::NeverSearched) {
+        // The existing symbol definition is weak and can be overridden.
+      } else {
+        // Both definitions are strong.
+        Duplicates.insert(KV.first);
     }
   }
 
   // If there were any duplicate definitions then bail out.
   if (!Duplicates.empty())
     return make_error<DuplicateDefinition>(std::string(**Duplicates.begin()));
 
   // Discard any overridden defs in this MU.
-  for (auto &S : MUDefsOverridden)
+  for (auto &S : RejectedWeakDefs)
     MU.doDiscard(*this, S);
 
   // Discard existing overridden defs.
   for (auto &S : ExistingDefsOverridden) {
 
     auto UMII = UnmaterializedInfos.find(S);
     assert(UMII != UnmaterializedInfos.end() &&
            "Overridden existing def should have an UnmaterializedInfo");