[ORC] Introduce SetUpExecutorNativePlatform utility.

Simplifies the process of building an LLJIT instance that supports the native
platform features (initializers, TLV, etc.).

SetUpExecutorNativePlatform can be passed to LLJITBuilder::setPlatformSetUp
method. It takes a reference to the ORC runtime (as a path or an in-memory
archive) and automatically sets the platform for LLJIT's ExecutionSession based
on the executor process's triple.

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/Core.h

@@ -1054,6 +1054,10 @@ class JITDylib : public ThreadSafeRefCountedBase<JITDylib>,
   void setLinkOrder(JITDylibSearchOrder NewSearchOrder,
                     bool LinkAgainstThisJITDylibFirst = true);
 
+  /// Append the given JITDylibSearchOrder to the link order for this
+  /// JITDylib.
+  void addToLinkOrder(const JITDylibSearchOrder &NewLinks);
+
   /// Add the given JITDylib to the link order for definitions in this
   /// JITDylib.
   ///

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) {
@@ -82,9 +96,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);
@@ -203,6 +220,10 @@ class LLJIT {
   std::unique_ptr<PlatformSupport> PS;
 
   JITDylib *Main = nullptr;
+  JITDylib *ProcessSymbols = nullptr;
+  JITDylib *Platform = nullptr;
+
+  JITDylibSearchOrder DefaultLinks;
 
   DataLayout DL;
   Triple TT;
@@ -258,12 +279,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 LinkProcessSymbolsJITDylibByDefault = true;
+  ProcessSymbolsJITDylibSetupFunction SetupProcessSymbolsJITDylib;
   ObjectLinkingLayerCreator CreateObjectLinkingLayer;
   CompileFunctionCreator CreateCompileFunction;
   PlatformSetupFunction SetUpPlatform;
@@ -316,6 +342,29 @@ class LLJITBuilderSetters {
     return impl();
   }
 
+  /// The LinkProcessSymbolsJITDylibDyDefault 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 &
+  setLinkProcessSymbolsJITDylibByDefault(bool LinkProcessSymsByDefault) {
+    impl().LinkProcessSymbolsJITDylibByDefault = LinkProcessSymsByDefault;
+    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
@@ -447,20 +496,52 @@ 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 SetUpExecutorNativePlatform {
+public:
+  /// Set up using path to Orc runtime. CreatePlatformJD will be run before
+  /// attempting to construct the platform instance. It should be used (if
+  /// needed) to provide the offers an opportunity
+  /// to load process symbols.
+  SetUpExecutorNativePlatform(std::string OrcRuntimePath)
+      : OrcRuntime(std::move(OrcRuntimePath)) {}
+
+  /// Set up using the given memory buffer.
+  SetUpExecutorNativePlatform(std::unique_ptr<MemoryBuffer> OrcRuntimeMB)
+      : OrcRuntime(std::move(OrcRuntimeMB)) {}
+
+  // TODO: add compiler-rt.
+
+  /// Add a path to the VC runtime.
+  SetUpExecutorNativePlatform &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)) {

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -1329,6 +1329,13 @@ void JITDylib::setLinkOrder(JITDylibSearchOrder NewLinkOrder,
   });
 }
 
+void JITDylib::addToLinkOrder(const JITDylibSearchOrder &NewLinks) {
+  ES.runSessionLocked([&]() {
+    LinkOrder.reserve(LinkOrder.size() + NewLinks.size());
+    llvm::append_range(LinkOrder, NewLinks);
+  });
+}
+
 void JITDylib::addToLinkOrder(JITDylib &JD, JITDylibLookupFlags JDLookupFlags) {
   ES.runSessionLocked([&]() { LinkOrder.push_back({&JD, JDLookupFlags}); });
 }