Implement LLVM compile-on-demand

JuliaLang · Apr 12, 2022 · 8aa06ba · 8aa06ba
1 parent 7f658ee
commit 8aa06ba
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp
@@ -986,6 +986,7 @@ JuliaOJIT::JuliaOJIT()
 #endif
     GlobalJD(ES.createBareJITDylib("JuliaGlobals")),
     JD(ES.createBareJITDylib("JuliaOJIT")),
+    LCTM(cantFail(orc::createLocalLazyCallThroughManager(TM->getTargetTriple(), ES, 0))),
     ContextPool([](){ return orc::ThreadSafeContext(std::make_unique<LLVMContext>()); }),
 #ifdef JL_USE_JITLINK
     // TODO: Port our memory management optimisations to JITLink instead of using the
@@ -1011,7 +1012,8 @@ JuliaOJIT::JuliaOJIT()
         std::make_unique<PipelineT>(ObjectLayer, *TM, 2),
         std::make_unique<PipelineT>(ObjectLayer, *TM, 3),
     },
-    OptSelLayer(Pipelines)
+    OptSelLayer(Pipelines),
+    CODLayer(ES, OptSelLayer, *LCTM, orc::createLocalIndirectStubsManagerBuilder(TM->getTargetTriple()))
 {
 #ifdef JL_USE_JITLINK
 # if defined(_OS_DARWIN_) && defined(LLVM_SHLIB)
@@ -1034,6 +1036,8 @@ JuliaOJIT::JuliaOJIT()
         });
 #endif
 
+    CODLayer.setPartitionFunction(CODLayerT::compileWholeModule);
+
     // Make sure SectionMemoryManager::getSymbolAddressInProcess can resolve
     // symbols in the program as well. The nullptr argument to the function
     // tells DynamicLibrary to load the program, not a library.
@@ -1112,11 +1116,11 @@ void JuliaOJIT::addModule(orc::ThreadSafeModule TSM)
 #endif
     });
     // TODO: what is the performance characteristics of this?
-    cantFail(OptSelLayer.add(JD, std::move(TSM)));
-    // force eager compilation (for now), due to memory management specifics
-    // (can't handle compilation recursion)
-    for (auto Name : NewExports)
-        cantFail(ES.lookup({&JD}, Name));
+    cantFail(CODLayer.add(JD, std::move(TSM)));
+    // // force eager compilation (for now), due to memory management specifics
+    // // (can't handle compilation recursion)
+    // for (auto Name : NewExports)
+    //     cantFail(ES.lookup({&JD}, Name));
 
 }
 

diff --git a/src/jitlayers.h b/src/jitlayers.h
@@ -9,6 +9,7 @@
 
 #include <llvm/ExecutionEngine/Orc/IRCompileLayer.h>
 #include <llvm/ExecutionEngine/Orc/IRTransformLayer.h>
+#include <llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h>
 #include <llvm/ExecutionEngine/JITEventListener.h>
 
 #include <llvm/Target/TargetMachine.h>
@@ -223,6 +224,7 @@ class JuliaOJIT {
 #endif
     typedef orc::IRCompileLayer CompileLayerT;
     typedef orc::IRTransformLayer OptimizeLayerT;
+    typedef orc::CompileOnDemandLayer CODLayerT;
     typedef object::OwningBinary<object::ObjectFile> OwningObj;
     template
     <typename ResourceT, size_t max = 0,
@@ -409,6 +411,8 @@ class JuliaOJIT {
     orc::JITDylib &GlobalJD;
     orc::JITDylib &JD;
 
+    std::unique_ptr<orc::LazyCallThroughManager> LCTM;
+
     ResourcePool<orc::ThreadSafeContext, 0, std::queue<orc::ThreadSafeContext>> ContextPool;
 
 #ifndef JL_USE_JITLINK
@@ -417,6 +421,7 @@ class JuliaOJIT {
     ObjLayerT ObjectLayer;
     const std::array<std::unique_ptr<PipelineT>, 4> Pipelines;
     OptSelLayerT OptSelLayer;
+    CODLayerT CODLayer;
 
     //Map and inc are guarded by RLST_mutex
     DenseMap<void*, std::string> ReverseLocalSymbolTable;