[Orc][LLJIT] Add optimizer support to LLJIT

[hintonda]

Summary

Add ability to optimize LLJIT modules by moving the TransformLayer
from LLLazyJIT up to LLJIT so it can be used by both. The LLLazyJIT
interface and functionality remains unchanged.

This is a minimal patch that reuses TransformLayer for both LLJIT and
LLLazyJIT. This shouldn't be a problem, since LLLazyJIT functionality
remains unchanged. However, an alternative would be to add a module
specific TransformLayer.

[alexr]

Seems simple enough, but Lang should weigh in.

[hintonda]

@lhames took care of this in e9e26c01cd865da678b1af6ba5f417c713956a66 yesterday.

[lhames]

Actually I missed that clone-to-new-context update in e9e26c01cd865da678b1af6ba5f417c713956a66. I've fixed it in c75180258e6. Thanks for this Don!

llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h

 
 public:
   static Expected<std::unique_ptr<LLJIT>> Create(LLJITBuilderState &S);
 
   /// Destruct this instance. If a multi-threaded instance, waits for all
   /// compile threads to complete.
   ~LLJIT();
 
+  /// Set an IR transform (e.g. pass manager pipeline) to run when it is
+  /// compiled.
+  void setCompileTransform(IRTransformLayer::TransformFunction Transform) {
+    TransformLayer->setTransform(std::move(Transform));
+  }
+
   /// Returns the ExecutionSession for this instance.
   ExecutionSession &getExecutionSession() { return *ES; }
 
   /// Returns a reference to the DataLayout for this instance.
   const DataLayout &getDataLayout() const { return DL; }
 
   /// Returns a reference to the JITDylib representing the JIT'd main program.
   JITDylib &getMainJITDylib() { return Main; }
protected:
   JITDylib &Main;
 
   DataLayout DL;
   std::unique_ptr<ThreadPool> CompileThreads;
 
   std::unique_ptr<ObjectLayer> ObjLinkingLayer;
   ObjectTransformLayer ObjTransformLayer;
   std::unique_ptr<IRCompileLayer> CompileLayer;
+  std::unique_ptr<IRTransformLayer> TransformLayer;
 
   CtorDtorRunner CtorRunner, DtorRunner;
 };
 
 /// An extended version of LLJIT that supports lazy function-at-a-time
 /// compilation of LLVM IR.
 class LLLazyJIT : public LLJIT {
   template <typename, typename, typename> friend class LLJITBuilderSetters;
public:
   }
 
 private:
 
   // Create a single-threaded LLLazyJIT instance.
   LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
 
   std::unique_ptr<LazyCallThroughManager> LCTMgr;
-  std::unique_ptr<IRTransformLayer> TransformLayer;
   std::unique_ptr<CompileOnDemandLayer> CODLayer;
 };
 
 class LLJITBuilderState {
 public:
   using ObjectLinkingLayerCreator = std::function<std::unique_ptr<ObjectLayer>(
       ExecutionSession &, const Triple &TT)>;
 

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

 
 Error LLJIT::addIRModule(JITDylib &JD, ThreadSafeModule TSM) {
   assert(TSM && "Can not add null module");
 
   if (auto Err =
           TSM.withModuleDo([&](Module &M) { return applyDataLayout(M); }))
     return Err;
 
-  return CompileLayer->add(JD, std::move(TSM), ES->allocateVModule());
+  return TransformLayer->add(JD, std::move(TSM), ES->allocateVModule());
 }
 
 Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
   assert(Obj && "Can not add null object");
 
   return ObjTransformLayer.add(JD, std::move(Obj), ES->allocateVModule());
 }
 
LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
   {
     auto CompileFunction = createCompileFunction(S, std::move(*S.JTMB));
     if (!CompileFunction) {
       Err = CompileFunction.takeError();
       return;
     }
     CompileLayer = std::make_unique<IRCompileLayer>(
         *ES, ObjTransformLayer, std::move(*CompileFunction));
+
+    // Create the transform layer.
+    TransformLayer = std::make_unique<IRTransformLayer>(*ES, *CompileLayer);
   }
 
   if (S.NumCompileThreads > 0) {
-    CompileLayer->setCloneToNewContextOnEmit(true);
+    TransformLayer->setCloneToNewContextOnEmit(true);
     CompileThreads = std::make_unique<ThreadPool>(S.NumCompileThreads);
     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));
         });
LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
   if (!ISMBuilder) {
     Err = make_error<StringError>("Could not construct "
                                   "IndirectStubsManagerBuilder for target " +
                                       S.TT.str(),
                                   inconvertibleErrorCode());
     return;
   }
 
-  // Create the transform layer.
-  TransformLayer = std::make_unique<IRTransformLayer>(*ES, *CompileLayer);
-
   // Create the COD layer.
   CODLayer = std::make_unique<CompileOnDemandLayer>(
       *ES, *TransformLayer, *LCTMgr, std::move(ISMBuilder));
 
   if (S.NumCompileThreads > 0)
     CODLayer->setCloneToNewContextOnEmit(true);
 }
 
 } // End namespace orc.
 } // End namespace llvm.