[ORC] Add support for multithreaded compiles to LLJIT and LLLazyJIT.

LLJIT and LLLazyJIT can now be constructed with an optional NumCompileThreads
arguments. If this is non-zero then a thread-pool will be created with the
given number of threads, and compile tasks will be dispatched to the thread
pool.

To enable testing of this feature, two new flags are added to lli:

(1) -compile-threads=N (N = 0 by default) controls the number of compile threads
to use.

(2) -thread-entry can be used to execute code on additional threads. For each
-thread-entry argument supplied (multiple are allowed) a new thread will be
created and the given symbol called. These additional thread entry points are
called after static constructors are run, but before main.

llvm-svn: 343058

1 files changed, 82 insertions, 8 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index c79c47a0e33..ecdfd862c3b 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -12,13 +12,45 @@
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/Mangler.h"
 
+namespace {
+
+  // A SimpleCompiler that owns its TargetMachine.
+  class TMOwningSimpleCompiler : public llvm::orc::SimpleCompiler {
+  public:
+    TMOwningSimpleCompiler(std::unique_ptr<llvm::TargetMachine> TM)
+      : llvm::orc::SimpleCompiler(*TM), TM(std::move(TM)) {}
+  private:
+    // FIXME: shared because std::functions (and thus
+    // IRCompileLayer2::CompileFunction) are not moveable.
+    std::shared_ptr<llvm::TargetMachine> TM;
+  };
+
+} // end anonymous namespace
+
 namespace llvm {
 namespace orc {
 
+LLJIT::~LLJIT() {
+  if (CompileThreads)
+    CompileThreads->wait();
+}
+
 Expected<std::unique_ptr<LLJIT>>
-LLJIT::Create(std::unique_ptr<TargetMachine> TM, DataLayout DL) {
+LLJIT::Create(JITTargetMachineBuilder JTMB, DataLayout DL,
+              unsigned NumCompileThreads) {
+
+  if (NumCompileThreads == 0) {
+    // If NumCompileThreads == 0 then create a single-threaded LLJIT instance.
+    auto TM = JTMB.createTargetMachine();
+    if (!TM)
+      return TM.takeError();
+    return std::unique_ptr<LLJIT>(new LLJIT(llvm::make_unique<ExecutionSession>(),
+                                            std::move(*TM), std::move(DL)));
+  }
+
   return std::unique_ptr<LLJIT>(new LLJIT(llvm::make_unique<ExecutionSession>(),
-                                          std::move(TM), std::move(DL)));
+                                          std::move(JTMB), std::move(DL),
+                                          NumCompileThreads));
 }
 
 Error LLJIT::defineAbsolute(StringRef Name, JITEvaluatedSymbol Sym) {
@@ -52,12 +84,35 @@ Expected<JITEvaluatedSymbol> LLJIT::lookupLinkerMangled(JITDylib &JD,
 LLJIT::LLJIT(std::unique_ptr<ExecutionSession> ES,
              std::unique_ptr<TargetMachine> TM, DataLayout DL)
     : ES(std::move(ES)), Main(this->ES->createJITDylib("main")),
-      TM(std::move(TM)), DL(std::move(DL)),
+      DL(std::move(DL)),
       ObjLinkingLayer(*this->ES,
                       [this](VModuleKey K) { return getMemoryManager(K); }),
-      CompileLayer(*this->ES, ObjLinkingLayer, SimpleCompiler(*this->TM)),
+      CompileLayer(*this->ES, ObjLinkingLayer, TMOwningSimpleCompiler(std::move(TM))),
       CtorRunner(Main), DtorRunner(Main) {}
 
+LLJIT::LLJIT(std::unique_ptr<ExecutionSession> ES,
+             JITTargetMachineBuilder JTMB, DataLayout DL,
+             unsigned NumCompileThreads)
+    : ES(std::move(ES)), Main(this->ES->createJITDylib("main")),
+      DL(std::move(DL)),
+      ObjLinkingLayer(*this->ES,
+                      [this](VModuleKey K) { return getMemoryManager(K); }),
+      CompileLayer(*this->ES, ObjLinkingLayer, MultiThreadedSimpleCompiler(std::move(JTMB))),
+      CtorRunner(Main), DtorRunner(Main) {
+  assert(NumCompileThreads != 0 &&
+         "Multithreaded LLJIT instance can not be created with 0 threads");
+
+  CompileThreads = llvm::make_unique<ThreadPool>(NumCompileThreads);
+  this->ES->setDispatchMaterialization([this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
+      // FIXME: Switch to move capture once we have c++14.
+      auto SharedMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
+      auto Work = [SharedMU, &JD]() {
+        SharedMU->doMaterialize(JD);
+      };
+      CompileThreads->async(std::move(Work));
+    });
+}
+
 std::unique_ptr<RuntimeDyld::MemoryManager>
 LLJIT::getMemoryManager(VModuleKey K) {
   return llvm::make_unique<SectionMemoryManager>();
@@ -90,10 +145,11 @@ void LLJIT::recordCtorDtors(Module &M) {
 }
 
 Expected<std::unique_ptr<LLLazyJIT>>
-LLLazyJIT::Create(std::unique_ptr<TargetMachine> TM, DataLayout DL) {
+  LLLazyJIT::Create(JITTargetMachineBuilder JTMB, DataLayout DL,
+                    unsigned NumCompileThreads) {
   auto ES = llvm::make_unique<ExecutionSession>();
 
-  const Triple &TT = TM->getTargetTriple();
+  const Triple &TT = JTMB.getTargetTriple();
 
   auto CCMgr = createLocalCompileCallbackManager(TT, *ES, 0);
   if (!CCMgr)
@@ -107,9 +163,18 @@ LLLazyJIT::Create(std::unique_ptr<TargetMachine> TM, DataLayout DL) {
         std::string("No indirect stubs manager builder for ") + TT.str(),
         inconvertibleErrorCode());
 
-  return std::unique_ptr<LLLazyJIT>(
-      new LLLazyJIT(std::move(ES), std::move(TM), std::move(DL),
+  if (NumCompileThreads == 0) {
+    auto TM = JTMB.createTargetMachine();
+    if (!TM)
+      return TM.takeError();
+    return std::unique_ptr<LLLazyJIT>(
+      new LLLazyJIT(std::move(ES), std::move(*TM), std::move(DL),
                     std::move(CCMgr), std::move(ISMBuilder)));
+  }
+
+  return std::unique_ptr<LLLazyJIT>(
+      new LLLazyJIT(std::move(ES), std::move(JTMB), std::move(DL),
+                    NumCompileThreads, std::move(CCMgr), std::move(ISMBuilder)));
 }
 
 Error LLLazyJIT::addLazyIRModule(JITDylib &JD, ThreadSafeModule TSM) {
@@ -135,5 +200,14 @@ LLLazyJIT::LLLazyJIT(
       CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder)) {
 }
 
+LLLazyJIT::LLLazyJIT(
+    std::unique_ptr<ExecutionSession> ES, JITTargetMachineBuilder JTMB,
+    DataLayout DL, unsigned NumCompileThreads, std::unique_ptr<JITCompileCallbackManager> CCMgr,
+    std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder)
+    : LLJIT(std::move(ES), std::move(JTMB), std::move(DL), NumCompileThreads),
+      CCMgr(std::move(CCMgr)), TransformLayer(*this->ES, CompileLayer),
+      CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder)) {
+}
+
 } // End namespace orc.
 } // End namespace llvm.