WIP: Add a compile-on-demand layer

Makefile

@@ -190,7 +190,7 @@ else
 JL_PRIVATE_LIBS-$(USE_SYSTEM_ZLIB) += libz
 endif
 ifeq ($(USE_LLVM_SHLIB),1)
-JL_PRIVATE_LIBS-$(USE_SYSTEM_LLVM) += libLLVM libLLVM-13jl
+JL_PRIVATE_LIBS-$(USE_SYSTEM_LLVM) += libLLVM libLLVM-14jl
 endif
 JL_PRIVATE_LIBS-$(USE_SYSTEM_LIBUNWIND) += libunwind
 

src/Makefile

@@ -62,7 +62,7 @@ RUNTIME_CODEGEN_SRCS := jitlayers aotcompile debuginfo disasm llvm-simdloop llvm
 	llvm-final-gc-lowering llvm-pass-helpers llvm-late-gc-lowering \
 	llvm-lower-handlers llvm-gc-invariant-verifier llvm-propagate-addrspaces \
 	llvm-multiversioning llvm-alloc-opt llvm-alloc-helpers cgmemmgr llvm-remove-addrspaces \
-	llvm-remove-ni llvm-julia-licm llvm-demote-float16 llvm-cpufeatures
+	llvm-remove-ni llvm-julia-licm llvm-demote-float16 llvm-cpufeatures jitlinkmemmgr
 FLAGS += -I$(shell $(LLVM_CONFIG_HOST) --includedir)
 CG_LLVM_LIBS := all
 ifeq ($(USE_POLLY),1)
@@ -127,7 +127,7 @@ else
 ifeq ($(OS), Darwin)
 CG_LLVMLINK += $(LLVM_LDFLAGS) -lLLVM
 else
-CG_LLVMLINK += $(LLVM_LDFLAGS) -lLLVM-13jl
+CG_LLVMLINK += $(LLVM_LDFLAGS) -lLLVM-14jl
 endif
 endif
 endif

src/codegen.cpp

@@ -8215,7 +8215,7 @@ extern "C" void jl_init_llvm(void)
     defined(JL_USE_OPROFILE_JITEVENTS) || \
     defined(JL_USE_PERF_JITEVENTS)
 #ifdef JL_USE_JITLINK
-#error "JIT profiling support (JL_USE_*_JITEVENTS) not yet available on platforms that use JITLink"
+#warning "JIT profiling support (JL_USE_*_JITEVENTS) not yet available on platforms that use JITLink"
 #else
     const char *jit_profiling = getenv("ENABLE_JITPROFILING");
 

src/jitlayers.cpp

@@ -666,7 +666,7 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
                     continue;
                 }
                 auto SecName = Sec.getName().substr(SepPos + 1);
-                Info.SectionLoadAddresses[SecName] = jitlink::SectionRange(Sec).getStart();
+                Info.SectionLoadAddresses[SecName] = jitlink::SectionRange(Sec).getStart().getValue();
             }
             return Error::success();
         });
@@ -677,19 +677,17 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
 # ifdef LLVM_SHLIB
 class JLEHFrameRegistrar final : public jitlink::EHFrameRegistrar {
 public:
-    Error registerEHFrames(JITTargetAddress EHFrameSectionAddr,
-                         size_t EHFrameSectionSize) override {
+    Error registerEHFrames(ExecutorAddrRange EHFrameSectionAddr) override {
         register_eh_frames(
-            jitTargetAddressToPointer<uint8_t *>(EHFrameSectionAddr),
-            EHFrameSectionSize);
+            EHFrameSectionAddr.Start.toPtr<uint8_t*>(),
+            EHFrameSectionAddr.size());
         return Error::success();
     }
 
-    Error deregisterEHFrames(JITTargetAddress EHFrameSectionAddr,
-                           size_t EHFrameSectionSize) override {
+    Error deregisterEHFrames(ExecutorAddrRange EHFrameSectionAddr) override {
         deregister_eh_frames(
-            jitTargetAddressToPointer<uint8_t *>(EHFrameSectionAddr),
-            EHFrameSectionSize);
+            EHFrameSectionAddr.Start.toPtr<uint8_t*>(),
+            EHFrameSectionAddr.size());
         return Error::success();
     }
 };
@@ -884,7 +882,7 @@ namespace {
         .setCPU(TM.getTargetCPU().str())
         .setFeatures(TM.getTargetFeatureString())
         .setOptions(TM.Options)
-        .setRelocationModel(Reloc::Static)
+        .setRelocationModel(Reloc::PIC_)
         .setCodeModel(TM.getCodeModel())
         .setCodeGenOptLevel(CodeGenOptLevelFor(optlevel));
     }
@@ -897,6 +895,8 @@ llvm::DataLayout create_jl_data_layout(TargetMachine &TM) {
     return jl_data_layout;
 }
 
+Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>> CreateJuliaJITLinkMemMgr();
+
 JuliaOJIT::JuliaOJIT(LLVMContext *LLVMCtx)
   : TM(createTargetMachine()),
     DL(create_jl_data_layout(*TM)),
@@ -914,13 +914,17 @@ JuliaOJIT::JuliaOJIT(LLVMContext *LLVMCtx)
 #endif
     GlobalJD(ES.createBareJITDylib("JuliaGlobals")),
     JD(ES.createBareJITDylib("JuliaOJIT")),
+#if defined(JL_USE_JITLINK) && JL_LLVM_VERSION >= 140000
+    LCTM(cantFail(orc::createLocalLazyCallThroughManager(TM->getTargetTriple(), ES, 0))),
+#endif
 #ifdef JL_USE_JITLINK
     // TODO: Port our memory management optimisations to JITLink instead of using the
     // default InProcessMemoryManager.
 # if JL_LLVM_VERSION < 140000
     ObjectLayer(ES, std::make_unique<jitlink::InProcessMemoryManager>()),
 # else
-    ObjectLayer(ES, cantFail(jitlink::InProcessMemoryManager::Create())),
+    MemMgr(cantFail(CreateJuliaJITLinkMemMgr())),
+    ObjectLayer(ES, *MemMgr),
 # endif
 #else
     MemMgr(createRTDyldMemoryManager()),
@@ -943,6 +947,9 @@ JuliaOJIT::JuliaOJIT(LLVMContext *LLVMCtx)
         {ES, CompileLayer3, OptimizerT(PM3, 3)},
     },
     OptSelLayer(OptimizeLayers)
+#if defined(JL_USE_JITLINK) && JL_LLVM_VERSION >= 140000
+    ,JITLayer(ES, OptSelLayer, *LCTM, orc::createLocalIndirectStubsManagerBuilder(TM->getTargetTriple()))
+#endif
 {
 #ifdef JL_USE_JITLINK
 # if defined(_OS_DARWIN_) && defined(LLVM_SHLIB)
@@ -1043,7 +1050,13 @@ void JuliaOJIT::addModule(std::unique_ptr<Module> M)
     }
 #endif
     // TODO: what is the performance characteristics of this?
-    cantFail(OptSelLayer.add(JD, orc::ThreadSafeModule(std::move(M), TSCtx)));
+    cantFail(
+#if defined(JL_USE_JITLINK) && JL_LLVM_VERSION >= 140000
+        JITLayer
+#else
+        OptSelLayer
+#endif
+        .add(JD, orc::ThreadSafeModule(std::move(M), TSCtx)));
     // force eager compilation (for now), due to memory management specifics
     // (can't handle compilation recursion)
     for (auto Name : NewExports)

src/jitlayers.h

@@ -9,11 +9,14 @@
 
 #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>
 #include "julia_assert.h"
 
+#define FORCE_JITLINK
+
 // As of LLVM 13, there are two runtime JIT linker implementations, the older
 // RuntimeDyld (used via orc::RTDyldObjectLinkingLayer) and the newer JITLink
 // (used via orc::ObjectLinkingLayer).
@@ -30,7 +33,7 @@
 // and feature support (e.g. Windows, JITEventListeners for various profilers,
 // etc.). Thus, we currently only use JITLink where absolutely required, that is,
 // for Mac/aarch64.
-#if defined(_OS_DARWIN_) && defined(_CPU_AARCH64_)
+#if defined(_OS_DARWIN_) && defined(_CPU_AARCH64_) || defined(FORCE_JITLINK)
 # if JL_LLVM_VERSION < 130000
 #  warning "On aarch64-darwin, LLVM version >= 13 is required for JITLink; fallback suffers from occasional segfaults"
 # endif
@@ -188,6 +191,7 @@ class JuliaOJIT {
 #endif
     typedef orc::IRCompileLayer CompileLayerT;
     typedef orc::IRTransformLayer OptimizeLayerT;
+    typedef orc::CompileOnDemandLayer JITLayerT;
     typedef object::OwningBinary<object::ObjectFile> OwningObj;
 private:
     struct OptimizerT {
@@ -258,16 +262,24 @@ class JuliaOJIT {
     orc::JITDylib &GlobalJD;
     orc::JITDylib &JD;
 
-#ifndef JL_USE_JITLINK
+#ifdef JL_USE_JITLINK
+#if JL_LLVM_VERSION >= 140000
+    std::unique_ptr<orc::LazyCallThroughManager> LCTM;
+    std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr;
+#endif // JL_LLVM_VERSION >= 140000
+#else
     std::shared_ptr<RTDyldMemoryManager> MemMgr;
-#endif
+#endif // JL_USE_JITLINK
     ObjLayerT ObjectLayer;
     CompileLayerT CompileLayer0;
     CompileLayerT CompileLayer1;
     CompileLayerT CompileLayer2;
     CompileLayerT CompileLayer3;
     OptimizeLayerT OptimizeLayers[4];
     OptSelLayerT OptSelLayer;
+#if defined(JL_USE_JITLINK) && JL_LLVM_VERSION >= 140000
+    JITLayerT JITLayer;
+#endif
 
     DenseMap<void*, std::string> ReverseLocalSymbolTable;
 };