[ORC] LLJIT::Create() proposal: add flag to enable multithreaded codegen

[sgraenitz]

Summary

Pass a CompileFtorFactory to the LLJIT ctor. Depending on the Multithreaded parameter, this either produces a SimpleCompiler or a MultithreadedCompiler functor. Adjusted the implementation of LLLazyJIT::Create(), but did not yet add the flag there as it's used in lli.

[lhames]

Hi Stefan,

We actually want the compiler call to be able to proceed on multiple threads concurrently. I think we need a documentation fix here: It just isn't safe to use SimpleCompiler when you're running on multiple threads. You need to use MultiThreadedSimpleCompiler (also in include/llvm/ExecutionEngine/Orc/CompileUtils.h) instead.

Maybe it would be safer to name the multi-threaded variant "SimpleCompiler", and rename the current SimpleCompiler to SingleThreadedSimpleCompiler.

Cheers,
Lang.

[sgraenitz]

Ok, now I see the note on the SimpleCompiler class. I didn't consider that, because I didn't touch the Compile functor in IRCompileLayer2. It is the default in the LLJIT ctor.
Maybe LLJIT::Create() could have an option for it or an alternative like CreateMultiThreaded()?

[sgraenitz]

Add LLJIT::CreateMultiThreaded() instead.

[sgraenitz]

[ORC] LLJIT::Create() proposal: add flag to enable multithreaded codegen

[lhames]

I think the changes in r343122 provide a superset of this functionality, so this can be closed.

[lhames]

I suggest we close this: The code is stale now, though the idea could be applied (and would be a neater fit) in the new LLJITBuilder utility.

include/llvm/ExecutionEngine/Orc/LLJIT.h

 #include "llvm/Target/TargetMachine.h"
 
 namespace llvm {
 namespace orc {
 
 /// A pre-fabricated ORC JIT stack that can serve as an alternative to MCJIT.
 class LLJIT {
 public:
+  using CompileFtor = IRCompileLayer2::CompileFunction;
+  using CompileFtorFactory = std::function<CompileFtor(TargetMachine &)>;
+
   /// Create an LLJIT instance.
   static Expected<std::unique_ptr<LLJIT>>
   Create(std::unique_ptr<ExecutionSession> ES,
-         std::unique_ptr<TargetMachine> TM, DataLayout DL);
+         JITTargetMachineBuilder JTMB, DataLayout DL, bool Multithreaded);
 
   /// Returns a reference to the ExecutionSession for this JIT instance.
   ExecutionSession &getExecutionSession() { return *ES; }
 
   /// Returns a reference to the JITDylib representing the JIT'd main program.
   JITDylib &getMainJITDylib() { return Main; }
 
   /// Convenience method for defining an absolute symbol.
public:
   /// Runs all not-yet-run static constructors.
   Error runConstructors() { return CtorRunner.run(); }
 
   /// Runs all not-yet-run static destructors.
   Error runDestructors() { return DtorRunner.run(); }
 
 protected:
   LLJIT(std::unique_ptr<ExecutionSession> ES, std::unique_ptr<TargetMachine> TM,
-        DataLayout DL);
+        DataLayout DL, CompileFtorFactory MakeCompileFtor);
 
   std::shared_ptr<RuntimeDyld::MemoryManager> getMemoryManager(VModuleKey K);
 
   std::string mangle(StringRef UnmangledName);
 
   Error applyDataLayout(Module &M);
 
   void recordCtorDtors(Module &M);
public:
 
   /// Add a module to be lazily compiled to the main JITDylib.
   Error addLazyIRModule(std::unique_ptr<Module> M) {
     return addLazyIRModule(Main, std::move(M));
   }
 
 private:
   LLLazyJIT(std::unique_ptr<ExecutionSession> ES,
-            std::unique_ptr<TargetMachine> TM, DataLayout DL, LLVMContext &Ctx,
+            std::unique_ptr<TargetMachine> TM, DataLayout DL,
+            CompileFtorFactory MakeCompileFtor, LLVMContext &Ctx,
             std::unique_ptr<JITCompileCallbackManager> CCMgr,
             std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder);
 
   std::unique_ptr<JITCompileCallbackManager> CCMgr;
   std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder;
 
   IRTransformLayer2 TransformLayer;
   CompileOnDemandLayer2 CODLayer;
 };
 
 } // End namespace orc
 } // End namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_LLJIT_H

lib/ExecutionEngine/Orc/LLJIT.cpp

 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/Mangler.h"
 
 namespace llvm {
 namespace orc {
 
 Expected<std::unique_ptr<LLJIT>>
 LLJIT::Create(std::unique_ptr<ExecutionSession> ES,
-              std::unique_ptr<TargetMachine> TM, DataLayout DL) {
+              JITTargetMachineBuilder JTMB, DataLayout DL, bool Multithreaded) {
+  auto TM = JTMB.createTargetMachine();
+  if (!TM)
+    return TM.takeError();
+  
+  CompileFtorFactory MakeCompileFtor;
+  if (Multithreaded) {
+    MakeCompileFtor = [JTMB](TargetMachine &) -> CompileFtor {
+      return MultiThreadedSimpleCompiler(JTMB);
+    };
+  } else {
+    MakeCompileFtor = [](TargetMachine &TM) -> CompileFtor {
+      return SimpleCompiler(TM);
+    };
+  }
+
   return std::unique_ptr<LLJIT>(
-      new LLJIT(std::move(ES), std::move(TM), std::move(DL)));
+      new LLJIT(std::move(ES), std::move(*TM), std::move(DL),
+                std::move(MakeCompileFtor)));
 }
 
 Error LLJIT::defineAbsolute(StringRef Name, JITEvaluatedSymbol Sym) {
   auto InternedName = ES->getSymbolStringPool().intern(Name);
   SymbolMap Symbols({{InternedName, Sym}});
   return Main.define(absoluteSymbols(std::move(Symbols)));
 }
 

 }
 
 Expected<JITEvaluatedSymbol> LLJIT::lookupLinkerMangled(JITDylib &JD,
                                                         StringRef Name) {
   return llvm::orc::lookup({&JD}, ES->getSymbolStringPool().intern(Name));
 }
 
 LLJIT::LLJIT(std::unique_ptr<ExecutionSession> ES,
-             std::unique_ptr<TargetMachine> TM, DataLayout DL)
+             std::unique_ptr<TargetMachine> TM, DataLayout DL,
+             CompileFtorFactory MakeCompileFtor)
     : ES(std::move(ES)), Main(this->ES->createJITDylib("main")),
       TM(std::move(TM)), DL(std::move(DL)),
       ObjLinkingLayer(*this->ES,
                       [this](VModuleKey K) { return getMemoryManager(K); }),
-      CompileLayer(*this->ES, ObjLinkingLayer, SimpleCompiler(*this->TM)),
+      CompileLayer(*this->ES, ObjLinkingLayer, MakeCompileFtor(*this->TM)),
       CtorRunner(Main), DtorRunner(Main) {}
 
 std::shared_ptr<RuntimeDyld::MemoryManager>
 LLJIT::getMemoryManager(VModuleKey K) {
   return llvm::make_unique<SectionMemoryManager>();
 }
 
 std::string LLJIT::mangle(StringRef UnmangledName) {
return make_error<StringError>(
         inconvertibleErrorCode());
 
   auto ISMBuilder = createLocalIndirectStubsManagerBuilder(TT);
   if (!ISMBuilder)
     return make_error<StringError>(
         std::string("No indirect stubs manager builder for ") + TT.str(),
         inconvertibleErrorCode());
 
+  auto MakeCompileFtor = [](TargetMachine &TM) -> CompileFtor {
+    return SimpleCompiler(TM);
+  };
+
   return std::unique_ptr<LLLazyJIT>(
-      new LLLazyJIT(std::move(ES), std::move(TM), std::move(DL), Ctx,
+      new LLLazyJIT(std::move(ES), std::move(TM), std::move(DL),
+                    std::move(MakeCompileFtor), Ctx,
                     std::move(CCMgr), std::move(ISMBuilder)));
 }
 
 Error LLLazyJIT::addLazyIRModule(JITDylib &JD, std::unique_ptr<Module> M) {
   assert(M && "Can not add null module");
 
   if (auto Err = applyDataLayout(*M))
     return Err;
 
   makeAllSymbolsExternallyAccessible(*M);
 
   recordCtorDtors(*M);
 
   auto K = ES->allocateVModule();
   return CODLayer.add(JD, K, std::move(M));
 }
 
 LLLazyJIT::LLLazyJIT(
     std::unique_ptr<ExecutionSession> ES, std::unique_ptr<TargetMachine> TM,
-    DataLayout DL, LLVMContext &Ctx,
+    DataLayout DL, CompileFtorFactory MakeCompileFtor, LLVMContext &Ctx,
     std::unique_ptr<JITCompileCallbackManager> CCMgr,
     std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder)
-    : LLJIT(std::move(ES), std::move(TM), std::move(DL)),
+    : LLJIT(std::move(ES), std::move(TM), std::move(DL),
+            std::move(MakeCompileFtor)),
       CCMgr(std::move(CCMgr)), TransformLayer(*this->ES, CompileLayer),
       CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder),
                [&]() -> LLVMContext & { return Ctx; }) {}
 
 } // End namespace orc.
 } // End namespace llvm.