[ORC] Add LookupKind::TransitiveStatic for linker-resolved transitive dependencies

[sgraenitz]

Summary

This lookup kind tells symbol generators that the lookup is being performed for transitive module dependencies. ThinLtoJIT wants to avoid loading and emitting modules for transitive dependencies where possible (if all requested symbols for one module are callable) and instead emit synthetic call-through stubs, which will load and emit the module when reached.

[sgraenitz]

llvm/include/llvm/ExecutionEngine/Orc/Core.h

89

This would need a description in the comment. Not sure though, if it's the best way to achieve my goal in the first place.

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

54

Alternatively, the generator could make the decision dependent on the required symbol state. However, currently this information isn't passed to the generators.

[lhames]

So the aim here is to enable generation of lazy stubs without having to materialize any IR, right?

I think the right way to do this: They’re supposed to have just enough information to decide whether the materialization unit being added should be part of the JIT’d program or not. They’re not intended for making smart decisions about how to build the materialization units.

I think a better approach would be to define a ModuleSummaryMaterializationUnit (or something similar) and have that be the thing that you add to your JIT layer. The layer would build stubs based on the interface defined in the module summary and defer compilation of the IR.

Let me know if I’m missing problems with that approach though: I’m afraid I haven’t had time to think about this deeply.

btw — I’m on #llvm on discord a lot these days. This might be a handy one to chat about in real time.

[sgraenitz]

In D75087#1904552, @lhames wrote:

So the aim here is to enable generation of lazy stubs without having to materialize any IR, right?

Right, AND without having to parse the respective bitcode files. In many cases the main module has transitive dependencies to a huge number of modules and materializing it causes parsing and emitting reexports for all of them. I do see that the case, which I refer to as TransitiveStatic here, could be the default for my generator. Entry points would then get a synthetic call-through too and load&emit the module on call (that would be acceptable). The problem, however, would then be the lookups issued from discovery. These should cause materialization of the requested modules. Unfortunately, there isn't a good way for the generator to distinguish them either (see inline comment).

I think the right way to do this: They’re supposed to have just enough information to decide whether the materialization unit being added should be part of the JIT’d program or not. They’re not intended for making smart decisions about how to build the materialization units.

I think a better approach would be to define a ModuleSummaryMaterializationUnit (or something similar) and have that be the thing that you add to your JIT layer. The layer would build stubs based on the interface defined in the module summary and defer compilation of the IR.

Not sure if that goes too much into the details, but one issue here is that module summaries only provide GUIDs. I don't have the full interface before parsing bitcode. The only symbol names I know are the requested ones.

The underlying question for me was, whether -- in general -- it could make sense for generators to distinguish between lookups coming from the linker and lookups from the "user".

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

263

This is the lookup issued from discovery. It's intended to force materialization of the modules that define LookaheadSymbols. As described above, I could assume TransitiveStatic the default. Then I'd probably pass LookupKind::DLSym here for the generator to indicate forced materialization, but that would be an abuse too.

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

Error ThinLtoJIT::runDiscovery(FunctionSummary *Origin, StringRef MangledName,
 
   // For each module, get one symbol that we can lookup in order to force
   // materialization.
   SymbolLookupSet LookaheadSymbols =
       GlobalIndex->createLookupSet(NonReadyModules, *Mangle);
 
   // Blocking lookup. We are done after this.
   auto SymMap = ES.lookup(SearchOrder, LookaheadSymbols,
                           LookupKind::Static, SymbolState::Ready);
 
   return SymMap ? Error::success() : SymMap.takeError();
 }
 
 Error ThinLtoJIT::submitWholeModuleBlocking(ModuleIndex::FetchModuleResult R,
                                             bool &NotReady) {
   ModuleNames ModulesToMaterialize;
 
Error ThinLtoJIT::generateMissingSymbolsBlocking(
   // For transitive dependencies, we emit synthetic call-through stubs if all
   // requested symbols for one module are callable. It avoids the blocking load
   // of all their modules here in the generator. Those modules that we actually
   // need, will hopefully be discovered and materialized in parallel before
   // calling through the stubs, but we cannot be sure.
   ModuleNames RemainingPaths;
   for (auto &KV : ModuleInfoMap) {
     StringRef Path = KV.first();
-    // TODO: Allow to distinguish regular static lookups from lookups for
-    // transitive dependencies issued by the linker.
+    ModuleInfo &MI = KV.second;
+    if (K == LookupKind::TransitiveStatic && MI.AllCallable) {
+      if (Error Err = submitCallThroughStubs(Path, std::move(MI)))
+        return Err;
+    } else {
       RemainingPaths.push_back(Path);
     }
+  }
 
   // For direct requests and transitive data dependencies, emit reexports.
   if (!RemainingPaths.empty()) {
     PSTATS("Generator requesting reexports (", K, "): ", RemainingPaths);
 
     for (FetchModuleResult &R : GlobalIndex->fetchAllModules(RemainingPaths)) {
       if (Error Err = submitReexportsBlocking(std::move(R)))
         return Err;

llvm/include/llvm/ExecutionEngine/Orc/Core.h

 /// definitions to generate.
 ///
 /// Static -- Lookup is being performed as-if at static link time (e.g.
 ///           generators representing static archives should pull in new
 ///           definitions).
 ///
 /// DLSym -- Lookup is being performed as-if at runtime (e.g. generators
 ///          representing static archives should not pull in new definitions).
-enum class LookupKind { Static, DLSym };
+enum class LookupKind { Static, DLSym, TransitiveStatic };
 
 /// A list of (JITDylib*, JITDylibLookupFlags) pairs to be used as a search
 /// order during symbol lookup.
 using JITDylibSearchOrder =
     std::vector<std::pair<JITDylib *, JITDylibLookupFlags>>;
 
 /// Convenience function for creating a search order from an ArrayRef of
 /// JITDylib*, all with the same flags.

llvm/lib/ExecutionEngine/Orc/Core.cpp

 }
 
 raw_ostream &operator<<(raw_ostream &OS, const LookupKind &K) {
   switch (K) {
   case LookupKind::Static:
     return OS << "Static";
   case LookupKind::DLSym:
     return OS << "DLSym";
+  case LookupKind::TransitiveStatic:
+    return OS << "TransitiveStatic";
   }
   llvm_unreachable("Invalid lookup kind");
 }
 
 raw_ostream &operator<<(raw_ostream &OS,
                         const JITDylibLookupFlags &JDLookupFlags) {
   switch (JDLookupFlags) {
   case JITDylibLookupFlags::MatchExportedSymbolsOnly:

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

void lookup(const LookupSet &Symbols, OnResolvedFunction OnResolved) {
     // Register dependencies for all symbols contained in this set.
     auto RegisterDependencies = [&](const SymbolDependenceMap &Deps) {
       MR.addDependenciesForAll(Deps);
     };
 
     JITDylibSearchOrder SearchOrder;
     MR.getTargetJITDylib().withSearchOrderDo(
         [&](const JITDylibSearchOrder &JDs) { SearchOrder = JDs; });
-    ES.lookup(LookupKind::Static, SearchOrder, InternedSymbols,
+    ES.lookup(LookupKind::TransitiveStatic, SearchOrder, InternedSymbols,
               SymbolState::Resolved, std::move(OnResolvedWithUnwrap),
               RegisterDependencies);
   }
 
   Expected<LookupSet> getResponsibilitySet(const LookupSet &Symbols) {
     LookupSet Result;
 
     for (auto &KV : MR.getSymbols()) {
       if (Symbols.count(*KV.first))