[ORC] Refactor trampoline pool management out of JITCompileCallbackManager.

This will allow trampoline pools to be re-used for a new lazy-reexport utility
that generates looks up function bodies using the standard symbol lookup process
(rather than using a user provided compile function). This new utility provides
the same capabilities (since MaterializationUnits already allow user supplied
compile functions to be run) as JITCompileCallbackManager, but can use the new
asynchronous lookup functions to avoid blocking a compile thread.

This patch also updates createLocalCompileCallbackManager to return an error if
a callback manager can not be created, and updates clients of that API to
account for the change. Finally, the OrcCBindingsStack is updates so that if
a callback manager is not available for the target platform a valid stack
(without support for lazy compilation) can still be constructed.

llvm-svn: 343059

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h

@@ -47,92 +47,101 @@ class Value;
 
 namespace orc {
 
-/// Target-independent base class for compile callback management.
-class JITCompileCallbackManager {
+/// Base class for pools of compiler re-entry trampolines.
+/// These trampolines are callable addresses that save all register state
+/// before calling a supplied function to return the trampoline landing
+/// address, then restore all state before jumping to that address. They
+/// are used by various ORC APIs to support lazy compilation
+class TrampolinePool {
 public:
-  using CompileFunction = std::function<JITTargetAddress()>;
-
-  /// Construct a JITCompileCallbackManager.
-  /// @param ErrorHandlerAddress The address of an error handler in the target
-  ///                            process to be used if a compile callback fails.
-  JITCompileCallbackManager(ExecutionSession &ES,
-                            JITTargetAddress ErrorHandlerAddress)
-      : ES(ES), CallbacksJD(ES.createJITDylib("<Callbacks>")),
-        ErrorHandlerAddress(ErrorHandlerAddress) {}
+  virtual ~TrampolinePool() {}
 
-  virtual ~JITCompileCallbackManager() = default;
+  /// Get an available trampoline address.
+  /// Returns an error if no trampoline can be created.
+  virtual Expected<JITTargetAddress> getTrampoline() = 0;
 
-  /// Reserve a compile callback.
-  Expected<JITTargetAddress> getCompileCallback(CompileFunction Compile);
-
-  /// Execute the callback for the given trampoline id. Called by the JIT
-  ///        to compile functions on demand.
-  JITTargetAddress executeCompileCallback(JITTargetAddress TrampolineAddr);
+private:
+  virtual void anchor();
+};
 
-protected:
-  std::vector<JITTargetAddress> AvailableTrampolines;
+/// A trampoline pool for trampolines within the current process.
+template <typename ORCABI> class LocalTrampolinePool : public TrampolinePool {
+public:
+  using GetTrampolineLandingFunction =
+      std::function<JITTargetAddress(JITTargetAddress TrampolineAddr)>;
+
+  /// Creates a LocalTrampolinePool with the given RunCallback function.
+  /// Returns an error if this function is unable to correctly allocate, write
+  /// and protect the resolver code block.
+  static Expected<std::unique_ptr<LocalTrampolinePool>>
+  Create(GetTrampolineLandingFunction GetTrampolineLanding) {
+    Error Err = Error::success();
+
+    auto LTP = std::unique_ptr<LocalTrampolinePool>(
+        new LocalTrampolinePool(std::move(GetTrampolineLanding), Err));
+
+    if (Err)
+      return std::move(Err);
+    return std::move(LTP);
+  }
 
-private:
-  Expected<JITTargetAddress> getAvailableTrampolineAddr() {
-    if (this->AvailableTrampolines.empty())
+  /// Get a free trampoline. Returns an error if one can not be provide (e.g.
+  /// because the pool is empty and can not be grown).
+  Expected<JITTargetAddress> getTrampoline() override {
+    std::lock_guard<std::mutex> Lock(LTPMutex);
+    if (AvailableTrampolines.empty()) {
       if (auto Err = grow())
         return std::move(Err);
-    assert(!this->AvailableTrampolines.empty() &&
-           "Failed to grow available trampolines.");
-    JITTargetAddress TrampolineAddr = this->AvailableTrampolines.back();
-    this->AvailableTrampolines.pop_back();
+    }
+    assert(!AvailableTrampolines.empty() && "Failed to grow trampoline pool");
+    auto TrampolineAddr = AvailableTrampolines.back();
+    AvailableTrampolines.pop_back();
     return TrampolineAddr;
   }
 
-  // Create new trampolines - to be implemented in subclasses.
-  virtual Error grow() = 0;
+  /// Returns the given trampoline to the pool for re-use.
+  void releaseTrampoline(JITTargetAddress TrampolineAddr) {
+    std::lock_guard<std::mutex> Lock(LTPMutex);
+    AvailableTrampolines.push_back(TrampolineAddr);
+  }
 
-  virtual void anchor();
+private:
+  static JITTargetAddress reenter(void *TrampolinePoolPtr, void *TrampolineId) {
+    LocalTrampolinePool<ORCABI> *TrampolinePool =
+        static_cast<LocalTrampolinePool *>(TrampolinePoolPtr);
+    return TrampolinePool->GetTrampolineLanding(static_cast<JITTargetAddress>(
+        reinterpret_cast<uintptr_t>(TrampolineId)));
+  }
 
-  std::mutex CCMgrMutex;
-  ExecutionSession &ES;
-  JITDylib &CallbacksJD;
-  JITTargetAddress ErrorHandlerAddress;
-  std::map<JITTargetAddress, SymbolStringPtr> AddrToSymbol;
-  size_t NextCallbackId = 0;
-};
+  LocalTrampolinePool(GetTrampolineLandingFunction GetTrampolineLanding,
+                      Error &Err)
+      : GetTrampolineLanding(std::move(GetTrampolineLanding)) {
 
-/// Manage compile callbacks for in-process JITs.
-template <typename TargetT>
-class LocalJITCompileCallbackManager : public JITCompileCallbackManager {
-public:
-  /// Construct a InProcessJITCompileCallbackManager.
-  /// @param ErrorHandlerAddress The address of an error handler in the target
-  ///                            process to be used if a compile callback fails.
-  LocalJITCompileCallbackManager(ExecutionSession &ES,
-                                 JITTargetAddress ErrorHandlerAddress)
-      : JITCompileCallbackManager(ES, ErrorHandlerAddress) {
-    /// Set up the resolver block.
+    ErrorAsOutParameter _(&Err);
+
+    /// Try to set up the resolver block.
     std::error_code EC;
     ResolverBlock = sys::OwningMemoryBlock(sys::Memory::allocateMappedMemory(
-        TargetT::ResolverCodeSize, nullptr,
+        ORCABI::ResolverCodeSize, nullptr,
         sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC));
-    assert(!EC && "Failed to allocate resolver block");
+    if (EC) {
+      Err = errorCodeToError(EC);
+      return;
+    }
 
-    TargetT::writeResolverCode(static_cast<uint8_t *>(ResolverBlock.base()),
-                               &reenter, this);
+    ORCABI::writeResolverCode(static_cast<uint8_t *>(ResolverBlock.base()),
+                              &reenter, this);
 
     EC = sys::Memory::protectMappedMemory(ResolverBlock.getMemoryBlock(),
                                           sys::Memory::MF_READ |
                                               sys::Memory::MF_EXEC);
-    assert(!EC && "Failed to mprotect resolver block");
-  }
-
-private:
-  static JITTargetAddress reenter(void *CCMgr, void *TrampolineId) {
-    JITCompileCallbackManager *Mgr =
-        static_cast<JITCompileCallbackManager *>(CCMgr);
-    return Mgr->executeCompileCallback(
-        static_cast<JITTargetAddress>(
-            reinterpret_cast<uintptr_t>(TrampolineId)));
+    if (EC) {
+      Err = errorCodeToError(EC);
+      return;
+    }
   }
 
-  Error grow() override {
+  Error grow() {
     assert(this->AvailableTrampolines.empty() && "Growing prematurely?");
 
     std::error_code EC;
@@ -144,17 +153,17 @@ class LocalJITCompileCallbackManager : public JITCompileCallbackManager {
       return errorCodeToError(EC);
 
     unsigned NumTrampolines =
-        (sys::Process::getPageSize() - TargetT::PointerSize) /
-        TargetT::TrampolineSize;
+        (sys::Process::getPageSize() - ORCABI::PointerSize) /
+        ORCABI::TrampolineSize;
 
     uint8_t *TrampolineMem = static_cast<uint8_t *>(TrampolineBlock.base());
-    TargetT::writeTrampolines(TrampolineMem, ResolverBlock.base(),
-                              NumTrampolines);
+    ORCABI::writeTrampolines(TrampolineMem, ResolverBlock.base(),
+                             NumTrampolines);
 
     for (unsigned I = 0; I < NumTrampolines; ++I)
       this->AvailableTrampolines.push_back(
           static_cast<JITTargetAddress>(reinterpret_cast<uintptr_t>(
-              TrampolineMem + (I * TargetT::TrampolineSize))));
+              TrampolineMem + (I * ORCABI::TrampolineSize))));
 
     if (auto EC = sys::Memory::protectMappedMemory(
                     TrampolineBlock.getMemoryBlock(),
@@ -165,8 +174,87 @@ class LocalJITCompileCallbackManager : public JITCompileCallbackManager {
     return Error::success();
   }
 
+  GetTrampolineLandingFunction GetTrampolineLanding;
+
+  std::mutex LTPMutex;
   sys::OwningMemoryBlock ResolverBlock;
   std::vector<sys::OwningMemoryBlock> TrampolineBlocks;
+  std::vector<JITTargetAddress> AvailableTrampolines;
+};
+
+/// Target-independent base class for compile callback management.
+class JITCompileCallbackManager {
+public:
+  using CompileFunction = std::function<JITTargetAddress()>;
+
+  virtual ~JITCompileCallbackManager() = default;
+
+  /// Reserve a compile callback.
+  Expected<JITTargetAddress> getCompileCallback(CompileFunction Compile);
+
+  /// Execute the callback for the given trampoline id. Called by the JIT
+  ///        to compile functions on demand.
+  JITTargetAddress executeCompileCallback(JITTargetAddress TrampolineAddr);
+
+protected:
+  /// Construct a JITCompileCallbackManager.
+  JITCompileCallbackManager(std::unique_ptr<TrampolinePool> TP,
+                            ExecutionSession &ES,
+                            JITTargetAddress ErrorHandlerAddress)
+      : TP(std::move(TP)), ES(ES),
+        CallbacksJD(ES.createJITDylib("<Callbacks>")),
+        ErrorHandlerAddress(ErrorHandlerAddress) {}
+
+  void setTrampolinePool(std::unique_ptr<TrampolinePool> TP) {
+    this->TP = std::move(TP);
+  }
+
+private:
+  std::mutex CCMgrMutex;
+  std::unique_ptr<TrampolinePool> TP;
+  ExecutionSession &ES;
+  JITDylib &CallbacksJD;
+  JITTargetAddress ErrorHandlerAddress;
+  std::map<JITTargetAddress, SymbolStringPtr> AddrToSymbol;
+  size_t NextCallbackId = 0;
+};
+
+/// Manage compile callbacks for in-process JITs.
+template <typename ORCABI>
+class LocalJITCompileCallbackManager : public JITCompileCallbackManager {
+public:
+  /// Create a new LocalJITCompileCallbackManager.
+  static Expected<std::unique_ptr<LocalJITCompileCallbackManager>>
+  Create(ExecutionSession &ES, JITTargetAddress ErrorHandlerAddress) {
+    Error Err = Error::success();
+    auto CCMgr = std::unique_ptr<LocalJITCompileCallbackManager>(
+        new LocalJITCompileCallbackManager(ES, ErrorHandlerAddress, Err));
+    if (Err)
+      return std::move(Err);
+    return std::move(CCMgr);
+  }
+
+private:
+  /// Construct a InProcessJITCompileCallbackManager.
+  /// @param ErrorHandlerAddress The address of an error handler in the target
+  ///                            process to be used if a compile callback fails.
+  LocalJITCompileCallbackManager(ExecutionSession &ES,
+                                 JITTargetAddress ErrorHandlerAddress,
+                                 Error &Err)
+      : JITCompileCallbackManager(nullptr, ES, ErrorHandlerAddress) {
+    ErrorAsOutParameter _(&Err);
+    auto TP = LocalTrampolinePool<ORCABI>::Create(
+        [this](JITTargetAddress TrampolineAddr) {
+          return executeCompileCallback(TrampolineAddr);
+        });
+
+    if (!TP) {
+      Err = TP.takeError();
+      return;
+    }
+
+    setTrampolinePool(std::move(*TP));
+  }
 };
 
 /// Base class for managing collections of named indirect stubs.
@@ -299,7 +387,7 @@ class LocalIndirectStubsManager : public IndirectStubsManager {
 /// The given target triple will determine the ABI, and the given
 /// ErrorHandlerAddress will be used by the resulting compile callback
 /// manager if a compile callback fails.
-std::unique_ptr<JITCompileCallbackManager>
+Expected<std::unique_ptr<JITCompileCallbackManager>>
 createLocalCompileCallbackManager(const Triple &T, ExecutionSession &ES,
                                   JITTargetAddress ErrorHandlerAddress);
 

llvm/include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

@@ -447,16 +447,24 @@ class OrcRemoteTargetClient
     StringMap<std::pair<StubKey, JITSymbolFlags>> StubIndexes;
   };
 
-  /// Remote compile callback manager.
-  class RemoteCompileCallbackManager : public JITCompileCallbackManager {
+  class RemoteTrampolinePool : public TrampolinePool {
   public:
-    RemoteCompileCallbackManager(OrcRemoteTargetClient &Client,
-                                 ExecutionSession &ES,
-                                 JITTargetAddress ErrorHandlerAddress)
-        : JITCompileCallbackManager(ES, ErrorHandlerAddress), Client(Client) {}
+    RemoteTrampolinePool(OrcRemoteTargetClient &Client) : Client(Client) {}
+
+    Expected<JITTargetAddress> getTrampoline() override {
+      std::lock_guard<std::mutex> Lock(RTPMutex);
+      if (AvailableTrampolines.empty()) {
+        if (auto Err = grow())
+          return std::move(Err);
+      }
+      assert(!AvailableTrampolines.empty() && "Failed to grow trampoline pool");
+      auto TrampolineAddr = AvailableTrampolines.back();
+      AvailableTrampolines.pop_back();
+      return TrampolineAddr;
+    }
 
   private:
-    Error grow() override {
+    Error grow() {
       JITTargetAddress BlockAddr = 0;
       uint32_t NumTrampolines = 0;
       if (auto TrampolineInfoOrErr = Client.emitTrampolineBlock())
@@ -471,7 +479,20 @@ class OrcRemoteTargetClient
       return Error::success();
     }
 
+    std::mutex RTPMutex;
     OrcRemoteTargetClient &Client;
+    std::vector<JITTargetAddress> AvailableTrampolines;
+  };
+
+  /// Remote compile callback manager.
+  class RemoteCompileCallbackManager : public JITCompileCallbackManager {
+  public:
+    RemoteCompileCallbackManager(OrcRemoteTargetClient &Client,
+                                 ExecutionSession &ES,
+                                 JITTargetAddress ErrorHandlerAddress)
+        : JITCompileCallbackManager(
+              llvm::make_unique<RemoteTrampolinePool>(Client), ES,
+              ErrorHandlerAddress) {}
   };
 
   /// Create an OrcRemoteTargetClient.