[ORC] LLJIT updates: ExecutorNativePlatform, default link order, Process JD.

This commit includes several related ergonomic improvements to LLJIT.

(1) Adds a default JITDylibSearchOrder to be appended to the initial link order
of JITDylibs created via LLJIT::createJITDylib (dropping any duplicate entries).

This was introduced to support automatic reflection of process symbols (see
(2) below), but has been made visible to clients as it's generically useful,
e.g. if clients have some extra set of libraries that they want to be visible
to JIT'd code by default.

The default JITDylibSearchOrder is only appended to the link order of JITDylibs
created via LLJIT::createJITDylib, and will not be apply to JITDylibs created by
directly calling the underlying ExecutionSession -- in that case clients can set
up the link order manually.

(2) Makes process symbols visible to JIT'd code by default via the new "Process"
JITDylib, which is added to the default link order.

LLJIT clients usually want symbols in the executor process to be accessible to
JIT'd code. Until now clients have been left to set this up themselves by adding
a DynamicLibrarySearchGenerator to the Main JITDylib. This patch adds a new
process symbols JITDylib that will be created by default (with an
EPCDynamicLibrarySearchGenerator attached) and added to the default link order,
making process symbols available to JIT'd code.

Clients who do not want process symbols to be visible to JIT'd code by default
can call setLinkProcessSymbolsByDefault(false) on their LLJITBuilder to disable
this:

LLJITBuilder()
  ...
  .setLinkProcessSymbolsByDefault(false)
  ...
  .create();

Clients can also call setProcessSymbolsJITDylibSetup to take over responsibility
for configuring the process symbols JITDylib (the callback that the client
supplies will be called on the bare process symbols JITDylib immediately after
it is created).

If setLinkProcessSymbolsByDefault(false) is called and no JITDylib setup
callback has been set then the process symbols JITDylib will not be created and
LLJIT::getProcessSymbolsJITDylib will return null.

(3) Adds an ExecutorNativePlatform utility that makes it easier to enable
native platform features.

Some object format features (e.g. native static initializers and thread locals)
require runtime support in the executing process. Support for these features in
ORC is implemented cooperatively between the ORC runtime and the LLVM Platform
subclasses (COFFPlatform, ELFNixPlatform, and MachOPlatform).
ExecutorNativePlatfrom simplifies the process of loading the ORC runtime and
creating the appropriate platform class for the executor process.

ExecutorNativePlatform takes a path to the ORC runtime (or a MemoryBuffer
containing the runtime) and other required runtimes for the executor platform
(e.g. MSVC on Windows) and then configures LLJIT with an appropriate platform
class based on the executor's target triple:

LLJITBuilder()
  .setPlatformSetUp(ExecutorNativePlatform("/path/to/orc-runtime.a"));

(The ORC runtime is built as part of compiler-rt, and the exact name of the
archive is platform dependent).

The ORC runtime and platform symbols will be added to a new "Platform" JITDylib,
which will be added to the *front* of the default link order (so JIT'd code will
prefer symbol definitions in the platform/runtime to definitions in the executor
process).

ExecutorNativePlatform assumes that the Process JITDylib is available, as
the ORC runtime may depend on symbols provided by the executor process.

Differential Revision: https://reviews.llvm.org/D144276

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h

@@ -39,6 +39,14 @@ class COFFPlatform : public Platform {
 
   /// Try to create a COFFPlatform instance, adding the ORC runtime to the
   /// given JITDylib.
+  static Expected<std::unique_ptr<COFFPlatform>>
+  Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+         JITDylib &PlatformJD,
+         std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
+         LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime = false,
+         const char *VCRuntimePath = nullptr,
+         std::optional<SymbolAliasMap> RuntimeAliases = std::nullopt);
+
   static Expected<std::unique_ptr<COFFPlatform>>
   Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
          JITDylib &PlatformJD, const char *OrcRuntimePath,
@@ -136,10 +144,14 @@ class COFFPlatform : public Platform {
 
   static bool supportedTarget(const Triple &TT);
 
-  COFFPlatform(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
-               JITDylib &PlatformJD, const char *OrcRuntimePath,
-               LoadDynamicLibrary LoadDynamicLibrary, bool StaticVCRuntime,
-               const char *VCRuntimePath, Error &Err);
+  COFFPlatform(
+      ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+      JITDylib &PlatformJD,
+      std::unique_ptr<StaticLibraryDefinitionGenerator> OrcRuntimeGenerator,
+      std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
+      std::unique_ptr<object::Archive> OrcRuntimeArchive,
+      LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
+      const char *VCRuntimePath, Error &Err);
 
   // Associate COFFPlatform JIT-side runtime support functions with handlers.
   Error associateRuntimeSupportFunctions(JITDylib &PlatformJD);

llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h

@@ -37,7 +37,7 @@ class ExecutorProcessControl;
 class LLJIT {
   template <typename, typename, typename> friend class LLJITBuilderSetters;
 
-  friend void setUpGenericLLVMIRPlatform(LLJIT &J);
+  friend Expected<JITDylibSP> setUpGenericLLVMIRPlatform(LLJIT &J);
 
 public:
   /// Initializer support for LLJIT.
@@ -70,6 +70,20 @@ class LLJIT {
   /// Returns a reference to the JITDylib representing the JIT'd main program.
   JITDylib &getMainJITDylib() { return *Main; }
 
+  /// Returns the ProcessSymbols JITDylib, which by default reflects non-JIT'd
+  /// symbols in the host process.
+  ///
+  /// Note: JIT'd code should not be added to the ProcessSymbols JITDylib. Use
+  /// the main JITDylib or a custom JITDylib instead.
+  JITDylibSP getProcessSymbolsJITDylib();
+
+  /// Returns the Platform JITDylib, which will contain the ORC runtime (if
+  /// given) and any platform symbols.
+  ///
+  /// Note: JIT'd code should not be added to the Platform JITDylib. Use the
+  /// main JITDylib or a custom JITDylib instead.
+  JITDylibSP getPlatformJITDylib();
+
   /// Returns the JITDylib with the given name, or nullptr if no JITDylib with
   /// that name exists.
   JITDylib *getJITDylibByName(StringRef Name) {
@@ -108,9 +122,12 @@ class LLJIT {
   /// input or elsewhere in the environment then the client should check
   /// (e.g. by calling getJITDylibByName) that the given name is not already in
   /// use.
-  Expected<JITDylib &> createJITDylib(std::string Name) {
-    return ES->createJITDylib(std::move(Name));
-  }
+  Expected<JITDylib &> createJITDylib(std::string Name);
+
+  /// Returns the default link order for this LLJIT instance. This link order
+  /// will be appended to the link order of JITDylibs created by LLJIT's
+  /// createJITDylib method.
+  JITDylibSearchOrder defaultLinkOrder() { return DefaultLinks; }
 
   /// Adds an IR module with the given ResourceTracker.
   Error addIRModule(ResourceTrackerSP RT, ThreadSafeModule TSM);
@@ -229,6 +246,10 @@ class LLJIT {
   std::unique_ptr<PlatformSupport> PS;
 
   JITDylib *Main = nullptr;
+  JITDylib *ProcessSymbols = nullptr;
+  JITDylib *Platform = nullptr;
+
+  JITDylibSearchOrder DefaultLinks;
 
   DataLayout DL;
   Triple TT;
@@ -284,12 +305,17 @@ class LLJITBuilderState {
       std::function<Expected<std::unique_ptr<IRCompileLayer::IRCompiler>>(
           JITTargetMachineBuilder JTMB)>;
 
-  using PlatformSetupFunction = std::function<Error(LLJIT &J)>;
+  using ProcessSymbolsJITDylibSetupFunction =
+      std::function<Error(JITDylib &JD)>;
+
+  using PlatformSetupFunction = unique_function<Expected<JITDylibSP>(LLJIT &J)>;
 
   std::unique_ptr<ExecutorProcessControl> EPC;
   std::unique_ptr<ExecutionSession> ES;
   std::optional<JITTargetMachineBuilder> JTMB;
   std::optional<DataLayout> DL;
+  bool LinkProcessSymbolsByDefault = true;
+  ProcessSymbolsJITDylibSetupFunction SetupProcessSymbolsJITDylib;
   ObjectLinkingLayerCreator CreateObjectLinkingLayer;
   CompileFunctionCreator CreateCompileFunction;
   PlatformSetupFunction SetUpPlatform;
@@ -342,6 +368,28 @@ class LLJITBuilderSetters {
     return impl();
   }
 
+  /// The LinkProcessSymbolsDyDefault flag determines whether the "Process"
+  /// JITDylib will be added to the default link order at LLJIT construction
+  /// time. If true, the Process JITDylib will be added as the last item in the
+  /// default link order. If false (or if the Process JITDylib is disabled via
+  /// setProcessSymbolsJITDylibSetup) then the Process JITDylib will not appear
+  /// in the default link order.
+  SetterImpl &setLinkProcessSymbolsByDefault(bool LinkProcessSymbolsByDefault) {
+    impl().LinkProcessSymbolsByDefault = LinkProcessSymbolsByDefault;
+    return impl();
+  }
+
+  /// Set a setup function for the process symbols dylib. If not provided,
+  /// but LinkProcessSymbolsJITDylibByDefault is true, then the process-symbols
+  /// JITDylib will be configured with a DynamicLibrarySearchGenerator with a
+  /// default symbol filter.
+  SetterImpl &setProcessSymbolsJITDylibSetup(
+      LLJITBuilderState::ProcessSymbolsJITDylibSetupFunction
+          SetupProcessSymbolsJITDylib) {
+    impl().SetupProcessSymbolsJITDylib = std::move(SetupProcessSymbolsJITDylib);
+    return impl();
+  }
+
   /// Set an ObjectLinkingLayer creation function.
   ///
   /// If this method is not called, a default creation function will be used
@@ -473,20 +521,49 @@ class LLLazyJITBuilder
       public LLLazyJITBuilderSetters<LLLazyJIT, LLLazyJITBuilder,
                                      LLLazyJITBuilderState> {};
 
-/// Configure the LLJIT instance to use orc runtime support. 
-Error setUpOrcPlatform(LLJIT& J);
+/// Configure the LLJIT instance to use orc runtime support. This overload
+/// assumes that the client has manually configured a Platform object.
+Error setUpOrcPlatformManually(LLJIT &J);
+
+/// Configure the LLJIT instance to use the ORC runtime and the detected
+/// native target for the executor.
+class ExecutorNativePlatform {
+public:
+  /// Set up using path to Orc runtime.
+  ExecutorNativePlatform(std::string OrcRuntimePath)
+      : OrcRuntime(std::move(OrcRuntimePath)) {}
+
+  /// Set up using the given memory buffer.
+  ExecutorNativePlatform(std::unique_ptr<MemoryBuffer> OrcRuntimeMB)
+      : OrcRuntime(std::move(OrcRuntimeMB)) {}
+
+  // TODO: add compiler-rt.
+
+  /// Add a path to the VC runtime.
+  ExecutorNativePlatform &addVCRuntime(std::string VCRuntimePath,
+                                       bool StaticVCRuntime) {
+    VCRuntime = {std::move(VCRuntimePath), StaticVCRuntime};
+    return *this;
+  }
+
+  Expected<JITDylibSP> operator()(LLJIT &J);
+
+private:
+  std::variant<std::string, std::unique_ptr<MemoryBuffer>> OrcRuntime;
+  std::optional<std::pair<std::string, bool>> VCRuntime;
+};
 
 /// Configure the LLJIT instance to scrape modules for llvm.global_ctors and
 /// llvm.global_dtors variables and (if present) build initialization and
 /// deinitialization functions. Platform specific initialization configurations
 /// should be preferred where available.
-void setUpGenericLLVMIRPlatform(LLJIT &J);
+Expected<JITDylibSP> setUpGenericLLVMIRPlatform(LLJIT &J);
 
 /// Configure the LLJIT instance to disable platform support explicitly. This is
 /// useful in two cases: for platforms that don't have such requirements and for
 /// platforms, that we have no explicit support yet and that don't work well
 /// with the generic IR platform.
-Error setUpInactivePlatform(LLJIT &J);
+Expected<JITDylibSP> setUpInactivePlatform(LLJIT &J);
 
 } // End namespace orc
 } // End namespace llvm

llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp

@@ -159,12 +159,11 @@ class COFFHeaderMaterializationUnit : public MaterializationUnit {
 namespace llvm {
 namespace orc {
 
-Expected<std::unique_ptr<COFFPlatform>>
-COFFPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
-                     JITDylib &PlatformJD, const char *OrcRuntimePath,
-                     LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
-                     const char *VCRuntimePath,
-                     std::optional<SymbolAliasMap> RuntimeAliases) {
+Expected<std::unique_ptr<COFFPlatform>> COFFPlatform::Create(
+    ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+    JITDylib &PlatformJD, std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
+    LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
+    const char *VCRuntimePath, std::optional<SymbolAliasMap> RuntimeAliases) {
 
   // If the target is not supported then bail out immediately.
   if (!supportedTarget(ES.getTargetTriple()))
@@ -174,6 +173,22 @@ COFFPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
 
   auto &EPC = ES.getExecutorProcessControl();
 
+  auto GeneratorArchive =
+      object::Archive::create(OrcRuntimeArchiveBuffer->getMemBufferRef());
+  if (!GeneratorArchive)
+    return GeneratorArchive.takeError();
+
+  auto OrcRuntimeArchiveGenerator = StaticLibraryDefinitionGenerator::Create(
+      ObjLinkingLayer, nullptr, std::move(*GeneratorArchive));
+  if (!OrcRuntimeArchiveGenerator)
+    return OrcRuntimeArchiveGenerator.takeError();
+
+  // We need a second instance of the archive (for now) for the Platform. We
+  // can `cantFail` this call, since if it were going to fail it would have
+  // failed above.
+  auto RuntimeArchive = cantFail(
+      object::Archive::create(OrcRuntimeArchiveBuffer->getMemBufferRef()));
+
   // Create default aliases if the caller didn't supply any.
   if (!RuntimeAliases)
     RuntimeAliases = standardPlatformAliases(ES);
@@ -199,13 +214,30 @@ COFFPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
   // Create the instance.
   Error Err = Error::success();
   auto P = std::unique_ptr<COFFPlatform>(new COFFPlatform(
-      ES, ObjLinkingLayer, PlatformJD, OrcRuntimePath,
+      ES, ObjLinkingLayer, PlatformJD, std::move(*OrcRuntimeArchiveGenerator),
+      std::move(OrcRuntimeArchiveBuffer), std::move(RuntimeArchive),
       std::move(LoadDynLibrary), StaticVCRuntime, VCRuntimePath, Err));
   if (Err)
     return std::move(Err);
   return std::move(P);
 }
 
+Expected<std::unique_ptr<COFFPlatform>>
+COFFPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+                     JITDylib &PlatformJD, const char *OrcRuntimePath,
+                     LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
+                     const char *VCRuntimePath,
+                     std::optional<SymbolAliasMap> RuntimeAliases) {
+
+  auto ArchiveBuffer = MemoryBuffer::getFile(OrcRuntimePath);
+  if (!ArchiveBuffer)
+    return createFileError(OrcRuntimePath, ArchiveBuffer.getError());
+
+  return Create(ES, ObjLinkingLayer, PlatformJD, std::move(*ArchiveBuffer),
+                std::move(LoadDynLibrary), StaticVCRuntime, VCRuntimePath,
+                std::move(RuntimeAliases));
+}
+
 Expected<MemoryBufferRef> COFFPlatform::getPerJDObjectFile() {
   auto PerJDObj = OrcRuntimeArchive->findSym("__orc_rt_coff_per_jd_marker");
   if (!PerJDObj)
@@ -349,37 +381,22 @@ bool COFFPlatform::supportedTarget(const Triple &TT) {
   }
 }
 
-COFFPlatform::COFFPlatform(ExecutionSession &ES,
-                           ObjectLinkingLayer &ObjLinkingLayer,
-                           JITDylib &PlatformJD, const char *OrcRuntimePath,
-                           LoadDynamicLibrary LoadDynamicLibrary,
-                           bool StaticVCRuntime, const char *VCRuntimePath,
-                           Error &Err)
+COFFPlatform::COFFPlatform(
+    ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+    JITDylib &PlatformJD,
+    std::unique_ptr<StaticLibraryDefinitionGenerator> OrcRuntimeGenerator,
+    std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
+    std::unique_ptr<object::Archive> OrcRuntimeArchive,
+    LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
+    const char *VCRuntimePath, Error &Err)
     : ES(ES), ObjLinkingLayer(ObjLinkingLayer),
-      LoadDynLibrary(std::move(LoadDynamicLibrary)),
+      LoadDynLibrary(std::move(LoadDynLibrary)),
+      OrcRuntimeArchiveBuffer(std::move(OrcRuntimeArchiveBuffer)),
+      OrcRuntimeArchive(std::move(OrcRuntimeArchive)),
       StaticVCRuntime(StaticVCRuntime),
       COFFHeaderStartSymbol(ES.intern("__ImageBase")) {
   ErrorAsOutParameter _(&Err);
 
-  // Create a generator for the ORC runtime archive.
-  auto OrcRuntimeArchiveGenerator =
-      StaticLibraryDefinitionGenerator::Load(ObjLinkingLayer, OrcRuntimePath);
-  if (!OrcRuntimeArchiveGenerator) {
-    Err = OrcRuntimeArchiveGenerator.takeError();
-    return;
-  }
-
-  auto ArchiveBuffer = MemoryBuffer::getFile(OrcRuntimePath);
-  if (!ArchiveBuffer) {
-    Err = createFileError(OrcRuntimePath, ArchiveBuffer.getError());
-    return;
-  }
-  OrcRuntimeArchiveBuffer = std::move(*ArchiveBuffer);
-  OrcRuntimeArchive =
-      std::make_unique<object::Archive>(*OrcRuntimeArchiveBuffer, Err);
-  if (Err)
-    return;
-
   Bootstrapping.store(true);
   ObjLinkingLayer.addPlugin(std::make_unique<COFFPlatformPlugin>(*this));
 
@@ -392,7 +409,7 @@ COFFPlatform::COFFPlatform(ExecutionSession &ES,
   }
   VCRuntimeBootstrap = std::move(*VCRT);
 
-  for (auto &Lib : (*OrcRuntimeArchiveGenerator)->getImportedDynamicLibraries())
+  for (auto &Lib : OrcRuntimeGenerator->getImportedDynamicLibraries())
     DylibsToPreload.insert(Lib);
 
   auto ImportedLibs =
@@ -406,7 +423,7 @@ COFFPlatform::COFFPlatform(ExecutionSession &ES,
   for (auto &Lib : *ImportedLibs)
     DylibsToPreload.insert(Lib);
 
-  PlatformJD.addGenerator(std::move(*OrcRuntimeArchiveGenerator));
+  PlatformJD.addGenerator(std::move(OrcRuntimeGenerator));
 
   // PlatformJD hasn't been set up by the platform yet (since we're creating
   // the platform now), so set it up.
@@ -416,10 +433,10 @@ COFFPlatform::COFFPlatform(ExecutionSession &ES,
   }
 
   for (auto& Lib : DylibsToPreload)
-      if (auto E2 = LoadDynLibrary(PlatformJD, Lib)) {
-          Err = std::move(E2);
-          return;
-      }
+    if (auto E2 = this->LoadDynLibrary(PlatformJD, Lib)) {
+      Err = std::move(E2);
+      return;
+    }
 
   if (StaticVCRuntime)
       if (auto E2 = VCRuntimeBootstrap->initializeStaticVCRuntime(PlatformJD)) {