[ORC] Update JITCompileCallbackManager to support multi-threaded code.

Previously JITCompileCallbackManager only supported single threaded code. This
patch embeds a VSO (see include/llvm/ExecutionEngine/Orc/Core.h) in the callback
manager. The VSO ensures that the compile callback is only executed once and that
the resulting address cached for use by subsequent re-entries.

llvm-svn: 333490

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h

@@ -76,8 +76,8 @@ public:
                       [this](std::unique_ptr<Module> M) {
                         return optimizeModule(std::move(M));
                       }),
-        CompileCallbackManager(
+        CompileCallbackManager(orc::createLocalCompileCallbackManager(
-            orc::createLocalCompileCallbackManager(TM->getTargetTriple(), 0)),
+            TM->getTargetTriple(), ES, 0)),
         CODLayer(ES, OptimizeLayer,
                  [&](orc::VModuleKey K) { return Resolvers[K]; },
                  [&](orc::VModuleKey K, std::shared_ptr<SymbolResolver> R) {

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter4/KaleidoscopeJIT.h

@@ -116,8 +116,8 @@ public:
                       [this](std::unique_ptr<Module> M) {
                         return optimizeModule(std::move(M));
                       }),
-        CompileCallbackMgr(
+        CompileCallbackMgr(orc::createLocalCompileCallbackManager(
-            orc::createLocalCompileCallbackManager(TM->getTargetTriple(), 0)) {
+            TM->getTargetTriple(), ES, 0)) {
     auto IndirectStubsMgrBuilder =
       orc::createLocalIndirectStubsManagerBuilder(TM->getTargetTriple());
     IndirectStubsMgr = IndirectStubsMgrBuilder();
@@ -134,22 +134,6 @@ public:
   }
 
   Error addFunctionAST(std::unique_ptr<FunctionAST> FnAST) {
-    // Create a CompileCallback - this is the re-entry point into the compiler
-    // for functions that haven't been compiled yet.
-    auto CCInfo = cantFail(CompileCallbackMgr->getCompileCallback());
-
-    // Create an indirect stub. This serves as the functions "canonical
-    // definition" - an unchanging (constant address) entry point to the
-    // function implementation.
-    // Initially we point the stub's function-pointer at the compile callback
-    // that we just created. In the compile action for the callback (see below)
-    // we will update the stub's function pointer to point at the function
-    // implementation that we just implemented.
-    if (auto Err = IndirectStubsMgr->createStub(mangle(FnAST->getName()),
-                                                CCInfo.getAddress(),
-                                                JITSymbolFlags::Exported))
-      return Err;
-
     // Move ownership of FnAST to a shared pointer - C++11 lambdas don't support
     // capture-by-move, which is be required for unique_ptr.
     auto SharedFnAST = std::shared_ptr<FunctionAST>(std::move(FnAST));
@@ -170,23 +154,37 @@ public:
     //     The JIT runtime (the resolver block) will use the return address of
     //     this function as the address to continue at once it has reset the
     //     CPU state to what it was immediately before the call.
-    CCInfo.setCompileAction(
+    auto CompileAction = [this, SharedFnAST]() {
-      [this, SharedFnAST]() {
       auto M = irgenAndTakeOwnership(*SharedFnAST, "$impl");
       addModule(std::move(M));
       auto Sym = findSymbol(SharedFnAST->getName() + "$impl");
       assert(Sym && "Couldn't find compiled function?");
       JITTargetAddress SymAddr = cantFail(Sym.getAddress());
-        if (auto Err =
+      if (auto Err = IndirectStubsMgr->updatePointer(
-              IndirectStubsMgr->updatePointer(mangle(SharedFnAST->getName()),
+              mangle(SharedFnAST->getName()), SymAddr)) {
-                                              SymAddr)) {
         logAllUnhandledErrors(std::move(Err), errs(),
                               "Error updating function pointer: ");
         exit(1);
       }
 
       return SymAddr;
-      });
+    };
+
+    // Create a CompileCallback using the CompileAction - this is the re-entry
+    // point into the compiler for functions that haven't been compiled yet.
+    auto CCAddr = cantFail(
+        CompileCallbackMgr->getCompileCallback(std::move(CompileAction)));
+
+    // Create an indirect stub. This serves as the functions "canonical
+    // definition" - an unchanging (constant address) entry point to the
+    // function implementation.
+    // Initially we point the stub's function-pointer at the compile callback
+    // that we just created. When the compile action for the callback is run we
+    // will update the stub's function pointer to point at the function
+    // implementation that we just implemented.
+    if (auto Err = IndirectStubsMgr->createStub(
+            mangle(SharedFnAST->getName()), CCAddr, JITSymbolFlags::Exported))
+      return Err;
 
     return Error::success();
   }

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter5/KaleidoscopeJIT.h

@@ -78,7 +78,7 @@ using MyRemote = remote::OrcRemoteTargetClient;
 
 class KaleidoscopeJIT {
 private:
-  ExecutionSession ES;
+  ExecutionSession &ES;
   std::shared_ptr<SymbolResolver> Resolver;
   std::unique_ptr<TargetMachine> TM;
   const DataLayout DL;
@@ -95,8 +95,9 @@ private:
   MyRemote &Remote;
 
 public:
-  KaleidoscopeJIT(MyRemote &Remote)
+  KaleidoscopeJIT(ExecutionSession &ES, MyRemote &Remote)
-      : Resolver(createLegacyLookupResolver(
+      : ES(ES),
+        Resolver(createLegacyLookupResolver(
             ES,
             [this](const std::string &Name) -> JITSymbol {
               if (auto Sym = IndirectStubsMgr->findStub(Name, false))
@@ -146,22 +147,6 @@ public:
   }
 
   Error addFunctionAST(std::unique_ptr<FunctionAST> FnAST) {
-    // Create a CompileCallback - this is the re-entry point into the compiler
-    // for functions that haven't been compiled yet.
-    auto CCInfo = cantFail(CompileCallbackMgr->getCompileCallback());
-
-    // Create an indirect stub. This serves as the functions "canonical
-    // definition" - an unchanging (constant address) entry point to the
-    // function implementation.
-    // Initially we point the stub's function-pointer at the compile callback
-    // that we just created. In the compile action for the callback (see below)
-    // we will update the stub's function pointer to point at the function
-    // implementation that we just implemented.
-    if (auto Err = IndirectStubsMgr->createStub(mangle(FnAST->getName()),
-                                                CCInfo.getAddress(),
-                                                JITSymbolFlags::Exported))
-      return Err;
-
     // Move ownership of FnAST to a shared pointer - C++11 lambdas don't support
     // capture-by-move, which is be required for unique_ptr.
     auto SharedFnAST = std::shared_ptr<FunctionAST>(std::move(FnAST));
@@ -182,23 +167,37 @@ public:
     //     The JIT runtime (the resolver block) will use the return address of
     //     this function as the address to continue at once it has reset the
     //     CPU state to what it was immediately before the call.
-    CCInfo.setCompileAction(
+    auto CompileAction = [this, SharedFnAST]() {
-      [this, SharedFnAST]() {
       auto M = irgenAndTakeOwnership(*SharedFnAST, "$impl");
       addModule(std::move(M));
       auto Sym = findSymbol(SharedFnAST->getName() + "$impl");
       assert(Sym && "Couldn't find compiled function?");
       JITTargetAddress SymAddr = cantFail(Sym.getAddress());
-        if (auto Err =
+      if (auto Err = IndirectStubsMgr->updatePointer(
-              IndirectStubsMgr->updatePointer(mangle(SharedFnAST->getName()),
+              mangle(SharedFnAST->getName()), SymAddr)) {
-                                              SymAddr)) {
         logAllUnhandledErrors(std::move(Err), errs(),
                               "Error updating function pointer: ");
         exit(1);
       }
 
       return SymAddr;
-      });
+    };
+
+    // Create a CompileCallback suing the CompileAction - this is the re-entry
+    // point into the compiler for functions that haven't been compiled yet.
+    auto CCAddr = cantFail(
+        CompileCallbackMgr->getCompileCallback(std::move(CompileAction)));
+
+    // Create an indirect stub. This serves as the functions "canonical
+    // definition" - an unchanging (constant address) entry point to the
+    // function implementation.
+    // Initially we point the stub's function-pointer at the compile callback
+    // that we just created. In the compile action for the callback we will
+    // update the stub's function pointer to point at the function
+    // implementation that we just implemented.
+    if (auto Err = IndirectStubsMgr->createStub(
+            mangle(SharedFnAST->getName()), CCAddr, JITSymbolFlags::Exported))
+      return Err;
 
     return Error::success();
   }

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter5/toy.cpp

@@ -1243,7 +1243,9 @@ std::unique_ptr<FDRPCChannel> connect() {
   sockaddr_in servAddr;
   memset(&servAddr, 0, sizeof(servAddr));
   servAddr.sin_family = PF_INET;
-  memcpy(&servAddr.sin_addr.s_addr, server->h_addr, server->h_length);
+  char *src;
+  memcpy(&src, &server->h_addr, sizeof(char *));
+  memcpy(&servAddr.sin_addr.s_addr, src, server->h_length);
   servAddr.sin_port = htons(Port);
   if (connect(sockfd, reinterpret_cast<sockaddr*>(&servAddr),
               sizeof(servAddr)) < 0) {
@@ -1276,9 +1278,10 @@ int main(int argc, char *argv[]) {
   BinopPrecedence['-'] = 20;
   BinopPrecedence['*'] = 40; // highest.
 
+  ExecutionSession ES;
   auto TCPChannel = connect();
-  auto Remote = ExitOnErr(MyRemote::Create(*TCPChannel, ExitOnErr));
+  auto Remote = ExitOnErr(MyRemote::Create(*TCPChannel, ES));
-  TheJIT = llvm::make_unique<KaleidoscopeJIT>(*Remote);
+  TheJIT = llvm::make_unique<KaleidoscopeJIT>(ES, *Remote);
 
   // Automatically inject a definition for 'printExprResult'.
   FunctionProtos["printExprResult"] =

llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

@@ -349,12 +349,7 @@ private:
         // Create a callback, associate it with the stub for the function,
         // and set the compile action to compile the partition containing the
         // function.
-        if (auto CCInfoOrErr = CompileCallbackMgr.getCompileCallback()) {
+        auto CompileAction = [this, &LD, LMId, &F]() -> JITTargetAddress {
-          auto &CCInfo = *CCInfoOrErr;
-          StubInits[MangledName] =
-            std::make_pair(CCInfo.getAddress(),
-                           JITSymbolFlags::fromGlobalValue(F));
-          CCInfo.setCompileAction([this, &LD, LMId, &F]() -> JITTargetAddress {
           if (auto FnImplAddrOrErr = this->extractAndCompile(LD, LMId, F))
             return *FnImplAddrOrErr;
           else {
@@ -362,9 +357,13 @@ private:
             consumeError(FnImplAddrOrErr.takeError());
             return 0;
           }
-            });
+        };
-        } else
+        if (auto CCAddr =
-          return CCInfoOrErr.takeError();
+                CompileCallbackMgr.getCompileCallback(std::move(CompileAction)))
+          StubInits[MangledName] =
+              std::make_pair(*CCAddr, JITSymbolFlags::fromGlobalValue(F));
+        else
+          return CCAddr.takeError();
       }
 
       if (auto Err = LD.StubsMgr->createStubs(StubInits))

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

@@ -610,10 +610,10 @@ private:
 /// VSOs will be searched in order and no VSO pointer may be null.
 /// All symbols must be found within the given VSOs or an error
 /// will be returned.
-Expected<SymbolMap> lookup(const std::vector<VSO *> &VSOs, SymbolNameSet Names);
+Expected<SymbolMap> lookup(const VSO::VSOList &VSOs, SymbolNameSet Names);
 
 /// Look up a symbol by searching a list of VSOs.
-Expected<JITEvaluatedSymbol> lookup(const std::vector<VSO *> VSOs,
+Expected<JITEvaluatedSymbol> lookup(const VSO::VSOList &VSOs,
                                     SymbolStringPtr Name);
 
 } // End namespace orc

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

@@ -18,6 +18,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Memory.h"
 #include "llvm/Support/Process.h"
@@ -49,95 +50,26 @@ namespace orc {
 /// Target-independent base class for compile callback management.
 class JITCompileCallbackManager {
 public:
-  using CompileFtor = std::function<JITTargetAddress()>;
+  using CompileFunction = std::function<JITTargetAddress()>;
-
-  /// Handle to a newly created compile callback. Can be used to get an
-  ///        IR constant representing the address of the trampoline, and to set
-  ///        the compile action for the callback.
-  class CompileCallbackInfo {
-  public:
-    CompileCallbackInfo(JITTargetAddress Addr, CompileFtor &Compile)
-        : Addr(Addr), Compile(Compile) {}
-
-    JITTargetAddress getAddress() const { return Addr; }
-    void setCompileAction(CompileFtor Compile) {
-      this->Compile = std::move(Compile);
-    }
-
-  private:
-    JITTargetAddress Addr;
-    CompileFtor &Compile;
-  };
 
   /// Construct a JITCompileCallbackManager.
   /// @param ErrorHandlerAddress The address of an error handler in the target
   ///                            process to be used if a compile callback fails.
-  JITCompileCallbackManager(JITTargetAddress ErrorHandlerAddress)
+  JITCompileCallbackManager(ExecutionSession &ES,
-      : ErrorHandlerAddress(ErrorHandlerAddress) {}
+                            JITTargetAddress ErrorHandlerAddress)
+      : ES(ES), CallbacksVSO(ES.createVSO("<Callbacks>")),
+        ErrorHandlerAddress(ErrorHandlerAddress) {}
 
   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) {
+  JITTargetAddress executeCompileCallback(JITTargetAddress TrampolineAddr);
-    auto I = ActiveTrampolines.find(TrampolineAddr);
-    // FIXME: Also raise an error in the Orc error-handler when we finally have
-    //        one.
-    if (I == ActiveTrampolines.end())
-      return ErrorHandlerAddress;
-
-    // Found a callback handler. Yank this trampoline out of the active list and
-    // put it back in the available trampolines list, then try to run the
-    // handler's compile and update actions.
-    // Moving the trampoline ID back to the available list first means there's
-    // at
-    // least one available trampoline if the compile action triggers a request
-    // for
-    // a new one.
-    auto Compile = std::move(I->second);
-    ActiveTrampolines.erase(I);
-    AvailableTrampolines.push_back(TrampolineAddr);
-
-    if (auto Addr = Compile())
-      return Addr;
-
-    return ErrorHandlerAddress;
-  }
-
-  /// Reserve a compile callback.
-  Expected<CompileCallbackInfo> getCompileCallback() {
-    if (auto TrampolineAddrOrErr = getAvailableTrampolineAddr()) {
-      const auto &TrampolineAddr = *TrampolineAddrOrErr;
-      auto &Compile = this->ActiveTrampolines[TrampolineAddr];
-      return CompileCallbackInfo(TrampolineAddr, Compile);
-    } else
-      return TrampolineAddrOrErr.takeError();
-  }
-
-  /// Get a CompileCallbackInfo for an existing callback.
-  CompileCallbackInfo getCompileCallbackInfo(JITTargetAddress TrampolineAddr) {
-    auto I = ActiveTrampolines.find(TrampolineAddr);
-    assert(I != ActiveTrampolines.end() && "Not an active trampoline.");
-    return CompileCallbackInfo(I->first, I->second);
-  }
-
-  /// Release a compile callback.
-  ///
-  ///   Note: Callbacks are auto-released after they execute. This method should
-  /// only be called to manually release a callback that is not going to
-  /// execute.
-  void releaseCompileCallback(JITTargetAddress TrampolineAddr) {
-    auto I = ActiveTrampolines.find(TrampolineAddr);
-    assert(I != ActiveTrampolines.end() && "Not an active trampoline.");
-    ActiveTrampolines.erase(I);
-    AvailableTrampolines.push_back(TrampolineAddr);
-  }
 
 protected:
-  JITTargetAddress ErrorHandlerAddress;
-
-  using TrampolineMapT = std::map<JITTargetAddress, CompileFtor>;
-  TrampolineMapT ActiveTrampolines;
   std::vector<JITTargetAddress> AvailableTrampolines;
 
 private:
@@ -156,6 +88,13 @@ private:
   virtual Error grow() = 0;
 
   virtual void anchor();
+
+  std::mutex CCMgrMutex;
+  ExecutionSession &ES;
+  VSO &CallbacksVSO;
+  JITTargetAddress ErrorHandlerAddress;
+  std::map<JITTargetAddress, SymbolStringPtr> AddrToSymbol;
+  size_t NextCallbackId = 0;
 };
 
 /// Manage compile callbacks for in-process JITs.
@@ -165,8 +104,9 @@ 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(JITTargetAddress ErrorHandlerAddress)
+  LocalJITCompileCallbackManager(ExecutionSession &ES,
-      : JITCompileCallbackManager(ErrorHandlerAddress) {
+                                 JITTargetAddress ErrorHandlerAddress)
+      : JITCompileCallbackManager(ES, ErrorHandlerAddress) {
     /// Set up the resolver block.
     std::error_code EC;
     ResolverBlock = sys::OwningMemoryBlock(sys::Memory::allocateMappedMemory(
@@ -360,7 +300,7 @@ private:
 /// ErrorHandlerAddress will be used by the resulting compile callback
 /// manager if a compile callback fails.
 std::unique_ptr<JITCompileCallbackManager>
-createLocalCompileCallbackManager(const Triple &T,
+createLocalCompileCallbackManager(const Triple &T, ExecutionSession &ES,
                                   JITTargetAddress ErrorHandlerAddress);
 
 /// Create a local indriect stubs manager builder.

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

@@ -451,8 +451,9 @@ public:
   class RemoteCompileCallbackManager : public JITCompileCallbackManager {
   public:
     RemoteCompileCallbackManager(OrcRemoteTargetClient &Client,
+                                 ExecutionSession &ES,
                                  JITTargetAddress ErrorHandlerAddress)
-        : JITCompileCallbackManager(ErrorHandlerAddress), Client(Client) {}
+        : JITCompileCallbackManager(ES, ErrorHandlerAddress), Client(Client) {}
 
   private:
     Error grow() override {
@@ -477,10 +478,10 @@ public:
   /// Channel is the ChannelT instance to communicate on. It is assumed that
   /// the channel is ready to be read from and written to.
   static Expected<std::unique_ptr<OrcRemoteTargetClient>>
-  Create(rpc::RawByteChannel &Channel, std::function<void(Error)> ReportError) {
+  Create(rpc::RawByteChannel &Channel, ExecutionSession &ES) {
     Error Err = Error::success();
     auto Client = std::unique_ptr<OrcRemoteTargetClient>(
-        new OrcRemoteTargetClient(Channel, std::move(ReportError), Err));
+        new OrcRemoteTargetClient(Channel, ES, Err));
     if (Err)
       return std::move(Err);
     return std::move(Client);
@@ -534,12 +535,14 @@ public:
 
   Expected<RemoteCompileCallbackManager &>
   enableCompileCallbacks(JITTargetAddress ErrorHandlerAddress) {
+    assert(!CallbackManager && "CallbackManager already obtained");
+
     // Emit the resolver block on the JIT server.
     if (auto Err = callB<stubs::EmitResolverBlock>())
       return std::move(Err);
 
     // Create the callback manager.
-    CallbackManager.emplace(*this, ErrorHandlerAddress);
+    CallbackManager.emplace(*this, ES, ErrorHandlerAddress);
     RemoteCompileCallbackManager &Mgr = *CallbackManager;
     return Mgr;
   }
@@ -557,10 +560,10 @@ public:
   Error terminateSession() { return callB<utils::TerminateSession>(); }
 
 private:
-  OrcRemoteTargetClient(rpc::RawByteChannel &Channel,
+  OrcRemoteTargetClient(rpc::RawByteChannel &Channel, ExecutionSession &ES,
-                        std::function<void(Error)> ReportError, Error &Err)
+                        Error &Err)
       : rpc::SingleThreadedRPCEndpoint<rpc::RawByteChannel>(Channel, true),
-        ReportError(std::move(ReportError)) {
+        ES(ES) {
     ErrorAsOutParameter EAO(&Err);
 
     addHandler<utils::RequestCompile>(
@@ -580,7 +583,7 @@ private:
 
   void deregisterEHFrames(JITTargetAddress Addr, uint32_t Size) {
     if (auto Err = callB<eh::RegisterEHFrames>(Addr, Size))
-      ReportError(std::move(Err));
+      ES.reportError(std::move(Err));
   }
 
   void destroyRemoteAllocator(ResourceIdMgr::ResourceId Id) {
@@ -595,7 +598,7 @@ private:
   void destroyIndirectStubsManager(ResourceIdMgr::ResourceId Id) {
     IndirectStubOwnerIds.release(Id);
     if (auto Err = callB<stubs::DestroyIndirectStubsOwner>(Id))
-      ReportError(std::move(Err));
+      ES.reportError(std::move(Err));
   }
 
   Expected<std::tuple<JITTargetAddress, JITTargetAddress, uint32_t>>
@@ -628,7 +631,7 @@ private:
     if (auto AddrOrErr = callB<mem::ReserveMem>(Id, Size, Align))
       return *AddrOrErr;
     else {
-      ReportError(AddrOrErr.takeError());
+      ES.reportError(AddrOrErr.takeError());
       return 0;
     }
   }
@@ -636,7 +639,7 @@ private:
   bool setProtections(ResourceIdMgr::ResourceId Id,
                       JITTargetAddress RemoteSegAddr, unsigned ProtFlags) {
     if (auto Err = callB<mem::SetProtections>(Id, RemoteSegAddr, ProtFlags)) {
-      ReportError(std::move(Err));
+      ES.reportError(std::move(Err));
       return true;
     } else
       return false;
@@ -644,7 +647,7 @@ private:
 
   bool writeMem(JITTargetAddress Addr, const char *Src, uint64_t Size) {
     if (auto Err = callB<mem::WriteMem>(DirectBufferWriter(Src, Addr, Size))) {
-      ReportError(std::move(Err));
+      ES.reportError(std::move(Err));
       return true;
     } else
       return false;
@@ -656,6 +659,7 @@ private:
 
   static Error doNothing() { return Error::success(); }
 
+  ExecutionSession &ES;
   std::function<void(Error)> ReportError;
   std::string RemoteTargetTriple;
   uint32_t RemotePointerSize = 0;

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -869,7 +869,7 @@ VSO &ExecutionSession::createVSO(std::string Name) {
   });
 }
 
-Expected<SymbolMap> lookup(const std::vector<VSO *> &VSOs, SymbolNameSet Names) {
+Expected<SymbolMap> lookup(const VSO::VSOList &VSOs, SymbolNameSet Names) {
 #if LLVM_ENABLE_THREADS
   // In the threaded case we use promises to return the results.
   std::promise<SymbolMap> PromisedResult;
@@ -975,7 +975,7 @@ Expected<SymbolMap> lookup(const std::vector<VSO *> &VSOs, SymbolNameSet Names)
 }
 
 /// Look up a symbol by searching a list of VSOs.
-Expected<JITEvaluatedSymbol> lookup(const std::vector<VSO *> VSOs,
+Expected<JITEvaluatedSymbol> lookup(const VSO::VSOList &VSOs,
                                     SymbolStringPtr Name) {
   SymbolNameSet Names({Name});
   if (auto ResultMap = lookup(VSOs, std::move(Names))) {

llvm/lib/ExecutionEngine/Orc/IndirectionUtils.cpp

@@ -13,38 +13,123 @@
 #include "llvm/ExecutionEngine/Orc/OrcABISupport.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/Support/Format.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include <sstream>
 
+using namespace llvm;
+using namespace llvm::orc;
+
+namespace {
+
+class CompileCallbackMaterializationUnit : public orc::MaterializationUnit {
+public:
+  using CompileFunction = JITCompileCallbackManager::CompileFunction;
+
+  CompileCallbackMaterializationUnit(SymbolStringPtr Name,
+                                     CompileFunction Compile)
+      : MaterializationUnit(SymbolFlagsMap({{Name, JITSymbolFlags::Exported}})),
+        Name(std::move(Name)), Compile(std::move(Compile)) {}
+
+private:
+  void materialize(MaterializationResponsibility R) {
+    SymbolMap Result;
+    Result[Name] = JITEvaluatedSymbol(Compile(), JITSymbolFlags::Exported);
+    R.resolve(Result);
+    R.finalize();
+  }
+
+  void discard(const VSO &V, SymbolStringPtr Name) {
+    llvm_unreachable("Discard should never occur on a LMU?");
+  }
+
+  SymbolStringPtr Name;
+  CompileFunction Compile;
+};
+
+} // namespace
+
 namespace llvm {
 namespace orc {
 
 void JITCompileCallbackManager::anchor() {}
 void IndirectStubsManager::anchor() {}
 
+Expected<JITTargetAddress>
+JITCompileCallbackManager::getCompileCallback(CompileFunction Compile) {
+  if (auto TrampolineAddr = getAvailableTrampolineAddr()) {
+    auto CallbackName = ES.getSymbolStringPool().intern(
+        std::string("cc") + std::to_string(++NextCallbackId));
+
+    std::lock_guard<std::mutex> Lock(CCMgrMutex);
+    AddrToSymbol[*TrampolineAddr] = CallbackName;
+    cantFail(
+        CallbacksVSO.define(make_unique<CompileCallbackMaterializationUnit>(
+            std::move(CallbackName), std::move(Compile))));
+    return *TrampolineAddr;
+  } else
+    return TrampolineAddr.takeError();
+}
+
+JITTargetAddress JITCompileCallbackManager::executeCompileCallback(
+    JITTargetAddress TrampolineAddr) {
+  SymbolStringPtr Name;
+
+  {
+    std::unique_lock<std::mutex> Lock(CCMgrMutex);
+    auto I = AddrToSymbol.find(TrampolineAddr);
+
+    // If this address is not associated with a compile callback then report an
+    // error to the execution session and return ErrorHandlerAddress to the
+    // callee.
+    if (I == AddrToSymbol.end()) {
+      Lock.unlock();
+      std::string ErrMsg;
+      {
+        raw_string_ostream ErrMsgStream(ErrMsg);
+        ErrMsgStream << "No compile callback for trampoline at "
+                     << format("0x%016x", TrampolineAddr);
+      }
+      ES.reportError(
+          make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode()));
+      return ErrorHandlerAddress;
+    } else
+      Name = I->second;
+  }
+
+  if (auto Sym = lookup({&CallbacksVSO}, Name))
+    return Sym->getAddress();
+  else {
+    // If anything goes wrong materializing Sym then report it to the session
+    // and return the ErrorHandlerAddress;
+    ES.reportError(Sym.takeError());
+    return ErrorHandlerAddress;
+  }
+}
+
 std::unique_ptr<JITCompileCallbackManager>
-createLocalCompileCallbackManager(const Triple &T,
+createLocalCompileCallbackManager(const Triple &T, ExecutionSession &ES,
                                   JITTargetAddress ErrorHandlerAddress) {
   switch (T.getArch()) {
     default: return nullptr;
 
     case Triple::aarch64: {
       typedef orc::LocalJITCompileCallbackManager<orc::OrcAArch64> CCMgrT;
-      return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+      return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
     }
 
     case Triple::x86: {
       typedef orc::LocalJITCompileCallbackManager<orc::OrcI386> CCMgrT;
-      return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+      return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
     }
 
     case Triple::x86_64: {
       if ( T.getOS() == Triple::OSType::Win32 ) {
         typedef orc::LocalJITCompileCallbackManager<orc::OrcX86_64_Win32> CCMgrT;
-        return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+        return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
       } else {
         typedef orc::LocalJITCompileCallbackManager<orc::OrcX86_64_SysV> CCMgrT;
-        return llvm::make_unique<CCMgrT>(ErrorHandlerAddress);
+        return llvm::make_unique<CCMgrT>(ES, ErrorHandlerAddress);
       }
     }
 

llvm/lib/ExecutionEngine/Orc/OrcCBindings.cpp

@@ -18,12 +18,11 @@ LLVMOrcJITStackRef LLVMOrcCreateInstance(LLVMTargetMachineRef TM) {
 
   Triple T(TM2->getTargetTriple());
 
-  auto CompileCallbackMgr = orc::createLocalCompileCallbackManager(T, 0);
   auto IndirectStubsMgrBuilder =
       orc::createLocalIndirectStubsManagerBuilder(T);
 
-  OrcCBindingsStack *JITStack = new OrcCBindingsStack(
+  OrcCBindingsStack *JITStack =
-      *TM2, std::move(CompileCallbackMgr), IndirectStubsMgrBuilder);
+      new OrcCBindingsStack(*TM2, std::move(IndirectStubsMgrBuilder));
 
   return wrap(JITStack);
 }

llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h

@@ -200,12 +200,10 @@ private:
   };
 
 public:
-
   OrcCBindingsStack(TargetMachine &TM,
-                    std::unique_ptr<CompileCallbackMgr> CCMgr,
                     IndirectStubsManagerBuilder IndirectStubsMgrBuilder)
-      : DL(TM.createDataLayout()),
+      : CCMgr(createLocalCompileCallbackManager(TM.getTargetTriple(), ES, 0)),
-        IndirectStubsMgr(IndirectStubsMgrBuilder()), CCMgr(std::move(CCMgr)),
+        DL(TM.createDataLayout()), IndirectStubsMgr(IndirectStubsMgrBuilder()),
         ObjectLayer(ES,
                     [this](orc::VModuleKey K) {
                       auto ResolverI = Resolvers.find(K);
@@ -217,7 +215,8 @@ public:
                           std::make_shared<SectionMemoryManager>(), Resolver};
                     },
                     nullptr,
-		    [this](orc::VModuleKey K, const object::ObjectFile &Obj, const RuntimeDyld::LoadedObjectInfo &LoadedObjInfo) {
+                    [this](orc::VModuleKey K, const object::ObjectFile &Obj,
+                           const RuntimeDyld::LoadedObjectInfo &LoadedObjInfo) {
 		      this->notifyFinalized(K, Obj, LoadedObjInfo);
                     },
                     [this](orc::VModuleKey K, const object::ObjectFile &Obj) {
@@ -270,15 +269,15 @@ public:
   createLazyCompileCallback(JITTargetAddress &RetAddr,
                             LLVMOrcLazyCompileCallbackFn Callback,
                             void *CallbackCtx) {
-    if (auto CCInfoOrErr = CCMgr->getCompileCallback()) {
+    auto WrappedCallback = [=]() -> JITTargetAddress {
-      auto &CCInfo = *CCInfoOrErr;
-      CCInfo.setCompileAction([=]() -> JITTargetAddress {
       return Callback(wrap(this), CallbackCtx);
-        });
+    };
-      RetAddr = CCInfo.getAddress();
+
+    if (auto CCAddr = CCMgr->getCompileCallback(std::move(WrappedCallback))) {
+      RetAddr = *CCAddr;
       return LLVMOrcErrSuccess;
     } else
-      return mapError(CCInfoOrErr.takeError());
+      return mapError(CCAddr.takeError());
   }
 
   LLVMOrcErrorCode createIndirectStub(StringRef StubName,
@@ -484,6 +483,7 @@ private:
   }
 
   orc::ExecutionSession ES;
+  std::unique_ptr<CompileCallbackMgr> CCMgr;
 
   std::vector<JITEventListener *> EventListeners;
 
@@ -492,7 +492,6 @@ private:
 
   std::unique_ptr<orc::IndirectStubsManager> IndirectStubsMgr;
 
-  std::unique_ptr<CompileCallbackMgr> CCMgr;
   ObjLayerT ObjectLayer;
   CompileLayerT CompileLayer;
   CODLayerT CODLayer;

llvm/tools/lli/OrcLazyJIT.cpp

@@ -121,15 +121,6 @@ int llvm::runOrcLazyJIT(std::vector<std::unique_ptr<Module>> Ms,
   EB.setOptLevel(getOptLevel());
   auto TM = std::unique_ptr<TargetMachine>(EB.selectTarget());
   Triple T(TM->getTargetTriple());
-  auto CompileCallbackMgr = orc::createLocalCompileCallbackManager(T, 0);
-
-  // If we couldn't build the factory function then there must not be a callback
-  // manager for this target. Bail out.
-  if (!CompileCallbackMgr) {
-    errs() << "No callback manager available for target '"
-           << TM->getTargetTriple().str() << "'.\n";
-    return 1;
-  }
 
   auto IndirectStubsMgrBuilder = orc::createLocalIndirectStubsManagerBuilder(T);
 
@@ -141,8 +132,7 @@ int llvm::runOrcLazyJIT(std::vector<std::unique_ptr<Module>> Ms,
   }
 
   // Everything looks good. Build the JIT.
-  OrcLazyJIT J(std::move(TM), std::move(CompileCallbackMgr),
+  OrcLazyJIT J(std::move(TM), std::move(IndirectStubsMgrBuilder),
-               std::move(IndirectStubsMgrBuilder),
                OrcInlineStubs);
 
   // Add the module, look up main and run it.

llvm/tools/lli/OrcLazyJIT.h

@@ -57,12 +57,11 @@ public:
   using IndirectStubsManagerBuilder = CODLayerT::IndirectStubsManagerBuilderT;
 
   OrcLazyJIT(std::unique_ptr<TargetMachine> TM,
-             std::unique_ptr<CompileCallbackMgr> CCMgr,
              IndirectStubsManagerBuilder IndirectStubsMgrBuilder,
              bool InlineStubs)
-      : TM(std::move(TM)),
+      : TM(std::move(TM)), DL(this->TM->createDataLayout()),
-        DL(this->TM->createDataLayout()),
+        CCMgr(orc::createLocalCompileCallbackManager(
-        CCMgr(std::move(CCMgr)),
+            this->TM->getTargetTriple(), ES, 0)),
         ObjectLayer(ES,
                     [this](orc::VModuleKey K) {
                       auto ResolverI = Resolvers.find(K);

llvm/tools/lli/lli.cpp

@@ -612,8 +612,10 @@ int main(int argc, char **argv, char * const *envp) {
     }
 
     // Create a remote target client running over the channel.
+    llvm::orc::ExecutionSession ES;
+    ES.setErrorReporter([&](Error Err) { ExitOnErr(std::move(Err)); });
     typedef orc::remote::OrcRemoteTargetClient MyRemote;
-    auto R = ExitOnErr(MyRemote::Create(*C, ExitOnErr));
+    auto R = ExitOnErr(MyRemote::Create(*C, ES));
 
     // Create a remote memory manager.
     auto RemoteMM = ExitOnErr(R->createRemoteMemoryManager());

llvm/unittests/ExecutionEngine/Orc/CompileOnDemandLayerTest.cpp

@@ -18,7 +18,8 @@ namespace {
 
 class DummyCallbackManager : public orc::JITCompileCallbackManager {
 public:
-  DummyCallbackManager() : JITCompileCallbackManager(0) {}
+  DummyCallbackManager(ExecutionSession &ES)
+      : JITCompileCallbackManager(ES, 0) {}
 
 public:
   Error grow() override { llvm_unreachable("not implemented"); }
@@ -57,9 +58,9 @@ TEST(CompileOnDemandLayerTest, FindSymbol) {
       return JITSymbol(nullptr);
     };
 
-  DummyCallbackManager CallbackMgr;
 
   ExecutionSession ES(std::make_shared<SymbolStringPool>());
+  DummyCallbackManager CallbackMgr(ES);
 
   auto GetResolver =
       [](orc::VModuleKey) -> std::shared_ptr<llvm::orc::SymbolResolver> {

llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h

@@ -71,11 +71,12 @@ public:
       // Use ability to create callback manager to detect whether Orc
       // has indirection support on this platform. This way the test
       // and Orc code do not get out of sync.
-      SupportsIndirection = !!orc::createLocalCompileCallbackManager(TT, 0);
+      SupportsIndirection = !!orc::createLocalCompileCallbackManager(TT, ES, 0);
     }
   };
 
 protected:
+  orc::ExecutionSession ES;
   LLVMContext Context;
   std::unique_ptr<TargetMachine> TM;
   bool SupportsJIT = false;