[ORC] Add absoluteSymbolsLinkGraph to expose absolute symbols to platform

Adds a function to create a LinkGraph of absolute symbols, and a
callback in dynamic library search generators to enable using it to
expose its symbols to the platform/orc runtime. This allows e.g. using
__orc_rt_run_program to run a precompiled function that was found via
dlsym. Ideally we would use this in llvm-jitlink's own search generator,
but it will require more work to align with the Process/Platform
JITDylib split, so not handled here.

As part of this change we need to handle LinkGraphs that only have
absolute symbols.

Author: benlangmuir

bolt/lib/Rewrite/JITLinkLinker.cpp

@@ -173,7 +173,8 @@ struct JITLinkLinker::Context : jitlink::JITLinkContext {
 
   void notifyFinalized(
       jitlink::JITLinkMemoryManager::FinalizedAlloc Alloc) override {
-    Linker.Allocs.push_back(std::move(Alloc));
+    if (Alloc)
+      Linker.Allocs.push_back(std::move(Alloc));
     ++Linker.MM->ObjectsLoaded;
   }
 };

llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h

@@ -19,6 +19,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h"
@@ -1923,6 +1924,10 @@ void visitExistingEdges(LinkGraph &G, VisitorTs &&...Vs) {
 Expected<std::unique_ptr<LinkGraph>>
 createLinkGraphFromObject(MemoryBufferRef ObjectBuffer);
 
+/// Create a \c LinkGraph defining the given absolute symbols.
+std::unique_ptr<LinkGraph> absoluteSymbolsLinkGraph(const Triple &TT,
+                                                    orc::SymbolMap Symbols);
+
 /// Link the given graph.
 void link(std::unique_ptr<LinkGraph> G, std::unique_ptr<JITLinkContext> Ctx);
 

llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h

@@ -25,31 +25,39 @@ class ExecutorProcessControl;
 class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
 public:
   using SymbolPredicate = unique_function<bool(const SymbolStringPtr &)>;
+  using AddAbsoluteSymbolsFn = unique_function<Error(JITDylib &, SymbolMap)>;
 
   /// Create a DynamicLibrarySearchGenerator that searches for symbols in the
   /// library with the given handle.
   ///
   /// If the Allow predicate is given then only symbols matching the predicate
   /// will be searched for. If the predicate is not given then all symbols will
   /// be searched for.
-  EPCDynamicLibrarySearchGenerator(ExecutionSession &ES,
-                                   tpctypes::DylibHandle H,
-                                   SymbolPredicate Allow = SymbolPredicate())
-      : EPC(ES.getExecutorProcessControl()), H(H), Allow(std::move(Allow)) {}
+  ///
+  /// If \p AddAbsoluteSymbols is provided, it is used to add the symbols to the
+  /// \c JITDylib; otherwise it uses JD.define(absoluteSymbols(...)).
+  EPCDynamicLibrarySearchGenerator(
+      ExecutionSession &ES, tpctypes::DylibHandle H,
+      SymbolPredicate Allow = SymbolPredicate(),
+      AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr)
+      : EPC(ES.getExecutorProcessControl()), H(H), Allow(std::move(Allow)),
+        AddAbsoluteSymbols(std::move(AddAbsoluteSymbols)) {}
 
   /// Permanently loads the library at the given path and, on success, returns
   /// a DynamicLibrarySearchGenerator that will search it for symbol definitions
   /// in the library. On failure returns the reason the library failed to load.
   static Expected<std::unique_ptr<EPCDynamicLibrarySearchGenerator>>
   Load(ExecutionSession &ES, const char *LibraryPath,
-       SymbolPredicate Allow = SymbolPredicate());
+       SymbolPredicate Allow = SymbolPredicate(),
+       AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr);
 
   /// Creates a EPCDynamicLibrarySearchGenerator that searches for symbols in
   /// the target process.
   static Expected<std::unique_ptr<EPCDynamicLibrarySearchGenerator>>
   GetForTargetProcess(ExecutionSession &ES,
-                      SymbolPredicate Allow = SymbolPredicate()) {
-    return Load(ES, nullptr, std::move(Allow));
+                      SymbolPredicate Allow = SymbolPredicate(),
+                      AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr) {
+    return Load(ES, nullptr, std::move(Allow), std::move(AddAbsoluteSymbols));
   }
 
   Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
@@ -60,6 +68,7 @@ class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
   ExecutorProcessControl &EPC;
   tpctypes::DylibHandle H;
   SymbolPredicate Allow;
+  AddAbsoluteSymbolsFn AddAbsoluteSymbols;
 };
 
 } // end namespace orc

llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h

@@ -216,29 +216,38 @@ class ItaniumCXAAtExitSupport {
 class DynamicLibrarySearchGenerator : public DefinitionGenerator {
 public:
   using SymbolPredicate = std::function<bool(const SymbolStringPtr &)>;
+  using AddAbsoluteSymbolsFn = unique_function<Error(JITDylib &, SymbolMap)>;
 
   /// Create a DynamicLibrarySearchGenerator that searches for symbols in the
   /// given sys::DynamicLibrary.
   ///
   /// If the Allow predicate is given then only symbols matching the predicate
   /// will be searched for. If the predicate is not given then all symbols will
   /// be searched for.
-  DynamicLibrarySearchGenerator(sys::DynamicLibrary Dylib, char GlobalPrefix,
-                                SymbolPredicate Allow = SymbolPredicate());
+  ///
+  /// If \p AddAbsoluteSymbols is provided, it is used to add the symbols to the
+  /// \c JITDylib; otherwise it uses JD.define(absoluteSymbols(...)).
+  DynamicLibrarySearchGenerator(
+      sys::DynamicLibrary Dylib, char GlobalPrefix,
+      SymbolPredicate Allow = SymbolPredicate(),
+      AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr);
 
   /// Permanently loads the library at the given path and, on success, returns
   /// a DynamicLibrarySearchGenerator that will search it for symbol definitions
   /// in the library. On failure returns the reason the library failed to load.
   static Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
   Load(const char *FileName, char GlobalPrefix,
-       SymbolPredicate Allow = SymbolPredicate());
+       SymbolPredicate Allow = SymbolPredicate(),
+       AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr);
 
   /// Creates a DynamicLibrarySearchGenerator that searches for symbols in
   /// the current process.
   static Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
   GetForCurrentProcess(char GlobalPrefix,
-                       SymbolPredicate Allow = SymbolPredicate()) {
-    return Load(nullptr, GlobalPrefix, std::move(Allow));
+                       SymbolPredicate Allow = SymbolPredicate(),
+                       AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr) {
+    return Load(nullptr, GlobalPrefix, std::move(Allow),
+                std::move(AddAbsoluteSymbols));
   }
 
   Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
@@ -248,6 +257,7 @@ class DynamicLibrarySearchGenerator : public DefinitionGenerator {
 private:
   sys::DynamicLibrary Dylib;
   SymbolPredicate Allow;
+  AddAbsoluteSymbolsFn AddAbsoluteSymbols;
   char GlobalPrefix;
 };
 

llvm/lib/ExecutionEngine/JITLink/JITLink.cpp

@@ -468,6 +468,41 @@ createLinkGraphFromObject(MemoryBufferRef ObjectBuffer) {
   };
 }
 
+std::unique_ptr<LinkGraph> absoluteSymbolsLinkGraph(const Triple &TT,
+                                                    orc::SymbolMap Symbols) {
+  unsigned PointerSize;
+  endianness Endianness =
+      TT.isLittleEndian() ? endianness::little : endianness::big;
+  switch (TT.getArch()) {
+  case Triple::aarch64:
+  case llvm::Triple::riscv64:
+  case Triple::x86_64:
+    PointerSize = 8;
+    break;
+  case llvm::Triple::arm:
+  case llvm::Triple::riscv32:
+  case llvm::Triple::x86:
+    PointerSize = 4;
+    break;
+  default:
+    llvm::report_fatal_error("unhandled target architecture");
+  }
+
+  static std::atomic<uint64_t> Counter = {0};
+  auto Index = Counter.fetch_add(1, std::memory_order_relaxed);
+  auto G = std::make_unique<LinkGraph>(
+      "<Absolute Symbols " + std::to_string(Index) + ">", TT, PointerSize,
+      Endianness, /*GetEdgeKindName=*/nullptr);
+  for (auto &[Name, Def] : Symbols) {
+    auto &Sym =
+        G->addAbsoluteSymbol(*Name, Def.getAddress(), /*Size=*/0,
+                             Linkage::Strong, Scope::Default, /*IsLive=*/true);
+    Sym.setCallable(Def.getFlags().isCallable());
+  }
+
+  return G;
+}
+
 void link(std::unique_ptr<LinkGraph> G, std::unique_ptr<JITLinkContext> Ctx) {
   switch (G->getTargetTriple().getObjectFormat()) {
   case Triple::MachO:

llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp

@@ -48,6 +48,14 @@ void JITLinkerBase::linkPhase1(std::unique_ptr<JITLinkerBase> Self) {
   if (auto Err = runPasses(Passes.PostPrunePasses))
     return Ctx->notifyFailed(std::move(Err));
 
+  // Skip straight to phase 2 if the graph is empty with no associated actions.
+  if (G->allocActions().empty() && llvm::all_of(G->sections(), [](Section &S) {
+        return S.getMemLifetime() == orc::MemLifetime::NoAlloc;
+      })) {
+    linkPhase2(std::move(Self), nullptr);
+    return;
+  }
+
   Ctx->getMemoryManager().allocate(
       Ctx->getJITLinkDylib(), *G,
       [S = std::move(Self)](AllocResult AR) mutable {
@@ -163,6 +171,12 @@ void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self,
   if (auto Err = runPasses(Passes.PostFixupPasses))
     return abandonAllocAndBailOut(std::move(Self), std::move(Err));
 
+  // Skip straight to phase 4 if the graph has no allocation.
+  if (!Alloc) {
+    linkPhase4(std::move(Self), JITLinkMemoryManager::FinalizedAlloc{});
+    return;
+  }
+
   Alloc->finalize([S = std::move(Self)](FinalizeResult FR) mutable {
     // FIXME: Once MSVC implements c++17 order of evaluation rules for calls
     // this can be simplified to

llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp

@@ -12,15 +12,15 @@ namespace llvm {
 namespace orc {
 
 Expected<std::unique_ptr<EPCDynamicLibrarySearchGenerator>>
-EPCDynamicLibrarySearchGenerator::Load(ExecutionSession &ES,
-                                       const char *LibraryPath,
-                                       SymbolPredicate Allow) {
+EPCDynamicLibrarySearchGenerator::Load(
+    ExecutionSession &ES, const char *LibraryPath, SymbolPredicate Allow,
+    AddAbsoluteSymbolsFn AddAbsoluteSymbols) {
   auto Handle = ES.getExecutorProcessControl().loadDylib(LibraryPath);
   if (!Handle)
     return Handle.takeError();
 
-  return std::make_unique<EPCDynamicLibrarySearchGenerator>(ES, *Handle,
-                                                            std::move(Allow));
+  return std::make_unique<EPCDynamicLibrarySearchGenerator>(
+      ES, *Handle, std::move(Allow), std::move(AddAbsoluteSymbols));
 }
 
 Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
@@ -62,6 +62,8 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
     return Error::success();
 
   // Define resolved symbols.
+  if (AddAbsoluteSymbols)
+    return AddAbsoluteSymbols(JD, std::move(NewSymbols));
   return JD.define(absoluteSymbols(std::move(NewSymbols)));
 }
 

llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp

@@ -218,19 +218,23 @@ void ItaniumCXAAtExitSupport::runAtExits(void *DSOHandle) {
 }
 
 DynamicLibrarySearchGenerator::DynamicLibrarySearchGenerator(
-    sys::DynamicLibrary Dylib, char GlobalPrefix, SymbolPredicate Allow)
+    sys::DynamicLibrary Dylib, char GlobalPrefix, SymbolPredicate Allow,
+    AddAbsoluteSymbolsFn AddAbsoluteSymbols)
     : Dylib(std::move(Dylib)), Allow(std::move(Allow)),
+      AddAbsoluteSymbols(std::move(AddAbsoluteSymbols)),
       GlobalPrefix(GlobalPrefix) {}
 
 Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
 DynamicLibrarySearchGenerator::Load(const char *FileName, char GlobalPrefix,
-                                    SymbolPredicate Allow) {
+                                    SymbolPredicate Allow,
+                                    AddAbsoluteSymbolsFn AddAbsoluteSymbols) {
   std::string ErrMsg;
   auto Lib = sys::DynamicLibrary::getPermanentLibrary(FileName, &ErrMsg);
   if (!Lib.isValid())
     return make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode());
   return std::make_unique<DynamicLibrarySearchGenerator>(
-      std::move(Lib), GlobalPrefix, std::move(Allow));
+      std::move(Lib), GlobalPrefix, std::move(Allow),
+      std::move(AddAbsoluteSymbols));
 }
 
 Error DynamicLibrarySearchGenerator::tryToGenerate(
@@ -261,6 +265,8 @@ Error DynamicLibrarySearchGenerator::tryToGenerate(
   if (NewSymbols.empty())
     return Error::success();
 
+  if (AddAbsoluteSymbols)
+    return AddAbsoluteSymbols(JD, std::move(NewSymbols));
   return JD.define(absoluteSymbols(std::move(NewSymbols)));
 }
 

llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp

@@ -1608,6 +1608,8 @@ Error MachOPlatform::MachOPlatformPlugin::prepareSymbolTableRegistration(
     SmallVector<jitlink::Symbol *> SymsToProcess;
     for (auto *Sym : G.defined_symbols())
       SymsToProcess.push_back(Sym);
+    for (auto *Sym : G.absolute_symbols())
+      SymsToProcess.push_back(Sym);
 
     for (auto *Sym : SymsToProcess) {
       if (!Sym->hasName())

llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp

@@ -93,15 +93,20 @@ class LinkGraphMaterializationUnit : public MaterializationUnit {
 
     Interface LGI;
 
-    for (auto *Sym : G.defined_symbols()) {
+    auto AddSymbol = [&](Symbol *Sym) {
       // Skip local symbols.
       if (Sym->getScope() == Scope::Local)
-        continue;
+        return;
       assert(Sym->hasName() && "Anonymous non-local symbol?");
 
       LGI.SymbolFlags[ES.intern(Sym->getName())] =
           getJITSymbolFlagsForSymbol(*Sym);
-    }
+    };
+
+    for (auto *Sym : G.defined_symbols())
+      AddSymbol(Sym);
+    for (auto *Sym : G.absolute_symbols())
+      AddSymbol(Sym);
 
     if (hasInitializerSection(G))
       LGI.InitSymbol = makeInitSymbol(ES, G);
@@ -705,6 +710,9 @@ Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
   if (Err)
     return Err;
 
+  if (!FA)
+    return Error::success();
+
   return MR.withResourceKeyDo(
       [&](ResourceKey K) { Allocs[K].push_back(std::move(FA)); });
 }