From 29af589c7c59cf9fe07010a904be34aa3a69d659 Mon Sep 17 00:00:00 2001 From: Prem Chintalapudi Date: Fri, 11 Mar 2022 21:06:22 -0500 Subject: [PATCH 1/4] Use LLVM14 --- Makefile | 2 +- src/Makefile | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index df6e998d47c35..ff8d7039e47d2 100644 --- a/Makefile +++ b/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 diff --git a/src/Makefile b/src/Makefile index a4fcf67612e0e..7fe057e404435 100644 --- a/src/Makefile +++ b/src/Makefile @@ -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 From 801266752990cc4f4faffe218c97dc1a09b5b626 Mon Sep 17 00:00:00 2001 From: Prem Chintalapudi Date: Sun, 27 Feb 2022 15:12:59 -0500 Subject: [PATCH 2/4] Fix build on LLVM14/JITLink --- src/codegen.cpp | 2 +- src/jitlayers.cpp | 18 ++++++++---------- src/jitlayers.h | 4 +++- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/codegen.cpp b/src/codegen.cpp index d35a34500b41b..7bf095281e37d 100644 --- a/src/codegen.cpp +++ b/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"); diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp index 231b37778eae5..cf31ab09f3ae6 100644 --- a/src/jitlayers.cpp +++ b/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(EHFrameSectionAddr), - EHFrameSectionSize); + EHFrameSectionAddr.Start.toPtr(), + EHFrameSectionAddr.size()); return Error::success(); } - Error deregisterEHFrames(JITTargetAddress EHFrameSectionAddr, - size_t EHFrameSectionSize) override { + Error deregisterEHFrames(ExecutorAddrRange EHFrameSectionAddr) override { deregister_eh_frames( - jitTargetAddressToPointer(EHFrameSectionAddr), - EHFrameSectionSize); + EHFrameSectionAddr.Start.toPtr(), + 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)); } diff --git a/src/jitlayers.h b/src/jitlayers.h index 035e81bead88d..9a72dbe23d107 100644 --- a/src/jitlayers.h +++ b/src/jitlayers.h @@ -14,6 +14,8 @@ #include #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 +32,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 From 0583752ac3b5c9b1269536982fea789d56bb55d2 Mon Sep 17 00:00:00 2001 From: Prem Chintalapudi Date: Fri, 11 Mar 2022 21:48:47 -0500 Subject: [PATCH 3/4] Custom memory manager + CODLayer initialization --- src/Makefile | 2 +- src/jitlayers.cpp | 11 +- src/jitlayers.h | 14 +- src/jitlinkmemmgr.cpp | 526 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 549 insertions(+), 4 deletions(-) create mode 100644 src/jitlinkmemmgr.cpp diff --git a/src/Makefile b/src/Makefile index 7fe057e404435..2bc62f408daed 100644 --- a/src/Makefile +++ b/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) diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp index cf31ab09f3ae6..29874c4dfca0d 100644 --- a/src/jitlayers.cpp +++ b/src/jitlayers.cpp @@ -895,6 +895,8 @@ llvm::DataLayout create_jl_data_layout(TargetMachine &TM) { return jl_data_layout; } +Expected> CreateJuliaJITLinkMemMgr(); + JuliaOJIT::JuliaOJIT(LLVMContext *LLVMCtx) : TM(createTargetMachine()), DL(create_jl_data_layout(*TM)), @@ -912,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()), # else - ObjectLayer(ES, cantFail(jitlink::InProcessMemoryManager::Create())), + MemMgr(cantFail(CreateJuliaJITLinkMemMgr())), + ObjectLayer(ES, *MemMgr), # endif #else MemMgr(createRTDyldMemoryManager()), @@ -941,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) diff --git a/src/jitlayers.h b/src/jitlayers.h index 9a72dbe23d107..80fd79c29e859 100644 --- a/src/jitlayers.h +++ b/src/jitlayers.h @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -190,6 +191,7 @@ class JuliaOJIT { #endif typedef orc::IRCompileLayer CompileLayerT; typedef orc::IRTransformLayer OptimizeLayerT; + typedef orc::CompileOnDemandLayer JITLayerT; typedef object::OwningBinary OwningObj; private: struct OptimizerT { @@ -260,9 +262,14 @@ class JuliaOJIT { orc::JITDylib &GlobalJD; orc::JITDylib &JD; -#ifndef JL_USE_JITLINK +#ifdef JL_USE_JITLINK +#if JL_LLVM_VERSION >= 140000 + std::unique_ptr LCTM; + std::unique_ptr MemMgr; +#endif // JL_LLVM_VERSION >= 140000 +#else std::shared_ptr MemMgr; -#endif +#endif // JL_USE_JITLINK ObjLayerT ObjectLayer; CompileLayerT CompileLayer0; CompileLayerT CompileLayer1; @@ -270,6 +277,9 @@ class JuliaOJIT { CompileLayerT CompileLayer3; OptimizeLayerT OptimizeLayers[4]; OptSelLayerT OptSelLayer; +#if defined(JL_USE_JITLINK) && JL_LLVM_VERSION >= 140000 + JITLayerT JITLayer; +#endif DenseMap ReverseLocalSymbolTable; }; diff --git a/src/jitlinkmemmgr.cpp b/src/jitlinkmemmgr.cpp new file mode 100644 index 0000000000000..a4c67448b73fb --- /dev/null +++ b/src/jitlinkmemmgr.cpp @@ -0,0 +1,526 @@ +#include "llvm-version.h" +#include +#include +#include +#include +#ifdef _OS_LINUX_ +# include +# include +# include +#endif +#ifndef _OS_WINDOWS_ +# include +# include +# include +# include +# if defined(_OS_DARWIN_) && !defined(MAP_ANONYMOUS) +# define MAP_ANONYMOUS MAP_ANON +# endif +#endif +#ifdef _OS_FREEBSD_ +# include +# include +#endif + +using namespace llvm; + +namespace { + + static int getPosixProtectionFlags(unsigned Flags) { + switch (Flags & llvm::sys::Memory::MF_RWE_MASK) { + case llvm::sys::Memory::MF_READ: + return PROT_READ; + case llvm::sys::Memory::MF_WRITE: + return PROT_WRITE; + case llvm::sys::Memory::MF_READ|llvm::sys::Memory::MF_WRITE: + return PROT_READ | PROT_WRITE; + case llvm::sys::Memory::MF_READ|llvm::sys::Memory::MF_EXEC: + return PROT_READ | PROT_EXEC; + case llvm::sys::Memory::MF_READ | llvm::sys::Memory::MF_WRITE | + llvm::sys::Memory::MF_EXEC: + return PROT_READ | PROT_WRITE | PROT_EXEC; + case llvm::sys::Memory::MF_EXEC: + #if (defined(__FreeBSD__) || defined(__POWERPC__) || defined (__ppc__) || \ + defined(_POWER) || defined(_ARCH_PPC)) + // On PowerPC, having an executable page that has no read permission + // can have unintended consequences. The function InvalidateInstruction- + // Cache uses instructions dcbf and icbi, both of which are treated by + // the processor as loads. If the page has no read permissions, + // executing these instructions will result in a segmentation fault. + return PROT_READ | PROT_EXEC; + #else + return PROT_EXEC; + #endif + default: + llvm_unreachable("Illegal memory protection flag specified!"); + } + // Provide a default return value as required by some compilers. + return PROT_NONE; + } + + static Expected reserve_address_space(uint64_t size) { + auto ptr = mmap64(nullptr, size, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED) { + perror("reserve_address_space"); + abort(); + return errorCodeToError(std::make_error_code(std::errc::not_enough_memory)); + } + return ptr; + } + + static Error unreserve_address_space(void *addr, uint64_t size) { + auto result = munmap(addr, size); + if (result == -1) { + perror("unreserve_address_space"); + abort(); + return errorCodeToError(std::make_error_code(std::errc::bad_address)); + } + return Error::success(); + } + + static Error commit(void *addr, uint64_t size, llvm::sys::Memory::ProtectionFlags flags) { + auto result = mprotect(addr, size, getPosixProtectionFlags(flags)); + if (result == -1) { + perror("commit"); + abort(); + return errorCodeToError(std::make_error_code(std::errc::invalid_argument)); + } + return Error::success(); + } + + static Error uncommit(void *addr, uint64_t size) { + auto result = mprotect(addr, size, PROT_NONE); + if (result == -1) { + perror("uncommit"); + abort(); + return errorCodeToError(std::make_error_code(std::errc::invalid_argument)); + } + return Error::success(); + } + + static constexpr uint64_t RESERVED_PAGE_SIZE = 1ull << 32; // 4GB address space + + llvm::sys::Memory::ProtectionFlags prot() { + return llvm::sys::Memory::ProtectionFlags(0); + } + + template + llvm::sys::Memory::ProtectionFlags prot(llvm::sys::Memory::ProtectionFlags flag, Flags... flags) { + return llvm::sys::Memory::ProtectionFlags(flag | prot(flags...)); + } + + class Allocator { + public: + struct Allocation { + sys::MemoryBlock StandardSegsMem; + sys::MemoryBlock FinalizeSegsMem; + }; + + virtual Expected allocate(const jitlink::JITLinkDylib *JD, uint64_t standard, uint64_t finalize, sys::Memory::ProtectionFlags flags) = 0; + virtual Error deallocate(sys::MemoryBlock &block) = 0; + }; + + class DefaultAllocator : public Allocator { // Stateless allocator that delegates directly to the OS + public: + Expected allocate(const jitlink::JITLinkDylib *JD, uint64_t standard, uint64_t finalize, sys::Memory::ProtectionFlags flags) override { + (void) JD; + std::error_code EC; + auto Slab = sys::Memory::allocateMappedMemory(standard + finalize, nullptr, + flags, EC); + + if (EC) { + return errorCodeToError(EC); + } + Allocation allocation; + + // Zero-fill the whole slab up-front. + memset(Slab.base(), 0, Slab.allocatedSize()); + + allocation.StandardSegsMem = {Slab.base(), + static_cast(standard)}; + allocation.FinalizeSegsMem = {(void *)((char *)Slab.base() + standard), + static_cast(finalize)}; + return allocation; + } + + Error deallocate(sys::MemoryBlock &block) override { + return errorCodeToError(sys::Memory::releaseMappedMemory(block)); + } + }; + + class SlabAllocator : public Allocator { + typedef IntervalMap Reserver; + class Slab { + private: + typedef IntervalMap AllocationTracker; + public: + Slab(uint64_t PageSize, sys::MemoryBlock reserved) : PageSize(PageSize), reserved_block(reserved), allocation_mutex(), allocator(std::make_unique()), allocated(*allocator) { + allocated.insert(reserved.base(), static_cast(reserved.base()) + reserved.allocatedSize() - 1, false); + } + + static Expected> Create(uint64_t PageSize, uint64_t SlabSize) { + auto addr = reserve_address_space(SlabSize); + if (!addr) { + return addr.takeError(); + } + return std::make_unique(PageSize, sys::MemoryBlock(*addr, SlabSize)); + } + sys::MemoryBlock reserved() { + return reserved_block; + } + + ~Slab() { + cantFail(unreserve_address_space(reserved_block.base(), reserved_block.allocatedSize())); + } + + Expected allocate(uint64_t size, sys::Memory::ProtectionFlags flags) { + std::lock_guard lock(allocation_mutex); + size = std::max(size, uint64_t(1)); + size = ((size + PageSize - 1) / PageSize) * PageSize; + for (auto it = allocated.begin(), end = allocated.end(); it != end; ++it) { + if (!*it) { + //We found an unallocated block + //Is it big enough? + auto start = static_cast(it.start()); + //IntervalMap is [start, end] + if (static_cast((static_cast(it.stop()) + 1) - start) >= size) { + auto end = it.stop(); + it.erase(); + allocated.insert(start, start + size - 1, true); + if (start + size <= end) { + allocated.insert(start + size, end, false); + } + if (auto err = commit(start, size, flags)) { + return std::move(err); + } + return sys::MemoryBlock(start, size); + } + } + } + //We didn't find a big enough unallocated block :( + return errorCodeToError(std::make_error_code(std::errc::not_enough_memory)); + } + + Error deallocate(sys::MemoryBlock block) { + std::lock_guard lock(allocation_mutex); + void *block_end = static_cast(block.base()) + block.allocatedSize() - 1; + auto it = allocated.find(block.base()); + auto start = it.start(), end = it.stop(); + assert(block.base() >= start && block_end <= end && "Deallocated block was not a contiguous block!"); + assert(*it && "Attempted to deallocate an already-freed block!"); + it.erase(); + if (start != block.base()) { + allocated.insert(start, static_cast(block.base()) - 1, true); + } + allocated.insert(block.base(), block_end, false); + if (block_end != end) { + allocated.insert(static_cast(block_end) + 1, end, true); + } + if (auto err = uncommit(block.base(), block.allocatedSize())) { + return std::move(err); + } + return Error::success(); + } + private: + uint64_t PageSize; + sys::MemoryBlock reserved_block; + std::mutex allocation_mutex; + std::unique_ptr allocator; + AllocationTracker allocated; + }; + public: + + SlabAllocator(uint64_t PageSize, uint64_t SlabSize) : PageSize(PageSize), SlabSize(SlabSize), allocations(), reserve_allocator(), reserved(reserve_allocator) {} + + Expected reserve(const jitlink::JITLinkDylib *JD) { + std::lock_guard lock(allocation_mutex); + auto it = allocations.find(JD); + if (it == allocations.end()) { + auto slab = Slab::Create(PageSize, SlabSize); + if (!slab) { + return slab.takeError(); + } + auto reserved = (**slab).reserved(); + allocations[JD] = std::move(*slab); + this->reserved.insert(reserved.base(), static_cast(reserved.base()) + reserved.allocatedSize() - 1, JD); + it = allocations.find(JD); + assert(it != allocations.end()); + } + return it->second.get(); + } + + Expected allocate(const jitlink::JITLinkDylib *JD, uint64_t standard, uint64_t finalize, sys::Memory::ProtectionFlags flags) override { + auto maybe_slab = reserve(JD); + if (!maybe_slab) { + return maybe_slab.takeError(); + } + Allocation allocation; + auto StandardSlab = (**maybe_slab).allocate(standard, flags); + if (!StandardSlab) { + return StandardSlab.takeError(); + } + allocation.StandardSegsMem = *StandardSlab; + memset(StandardSlab->base(), 0, StandardSlab->allocatedSize()); + auto FinalizeSlab = (**maybe_slab).allocate(finalize, flags); + if (!FinalizeSlab) { + return joinErrors(FinalizeSlab.takeError(), (**maybe_slab).deallocate(allocation.StandardSegsMem)); + } + allocation.FinalizeSegsMem = *FinalizeSlab; + memset(FinalizeSlab->base(), 0, FinalizeSlab->allocatedSize()); + return allocation; + } + + Error deallocate(sys::MemoryBlock &block) override { + std::lock_guard lock(allocation_mutex); + auto it = reserved.find(block.base()); + assert(it != reserved.end()); + return allocations[*it]->deallocate(std::move(block)); + } + + private: + uint64_t PageSize; + uint64_t SlabSize; + std::mutex allocation_mutex; + std::map> allocations; + Reserver::Allocator reserve_allocator; + Reserver reserved; + }; + + class JuliaJITLinkMemoryManager : public jitlink::JITLinkMemoryManager { + class JuliaInFlightAlloc : public jitlink::JITLinkMemoryManager::InFlightAlloc { + public: + JuliaInFlightAlloc(JuliaJITLinkMemoryManager &MemMgr, jitlink::LinkGraph &G, + jitlink::BasicLayout BL, sys::MemoryBlock StandardSegments, sys::MemoryBlock FinalizationSegments) + : MemMgr(MemMgr), G(G), BL(std::move(BL)), StandardSegments(StandardSegments), + FinalizationSegments(FinalizationSegments) {} + + /// Called prior to finalization if the allocation should be abandoned. + void abandon(OnAbandonedFunction OnAbandoned) override { + OnAbandoned(joinErrors(MemMgr.allocator->deallocate(FinalizationSegments), MemMgr.allocator->deallocate(StandardSegments))); + } + + /// Called to transfer working memory to the target and apply finalization. + void finalize(OnFinalizedFunction OnFinalized) override { + // Apply memory protections to all segments. + if (auto Err = applyProtections()) { + OnFinalized(std::move(Err)); + return; + } + + // Run finalization actions. + auto DeallocActions = orc::shared::runFinalizeActions(G.allocActions()); + if (!DeallocActions) { + OnFinalized(DeallocActions.takeError()); + return; + } + + // Release the finalize segments slab. + if (auto err = MemMgr.allocator->deallocate(FinalizationSegments)) { + OnFinalized(std::move(err)); + return; + } + + // Continue with finalized allocation. + OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments), + std::move(*DeallocActions))); + } + private: + Error applyProtections() { + for (auto &KV : BL.segments()) { + const auto &AG = KV.first; + auto &Seg = KV.second; + + auto Prot = toSysMemoryProtectionFlags(AG.getMemProt()); + + uint64_t SegSize = + alignTo(Seg.ContentSize + Seg.ZeroFillSize, MemMgr.PageSize); + sys::MemoryBlock MB(Seg.WorkingMem, SegSize); + if (auto EC = sys::Memory::protectMappedMemory(MB, Prot)) + return errorCodeToError(EC); + if (Prot & sys::Memory::MF_EXEC) + sys::Memory::InvalidateInstructionCache(MB.base(), MB.allocatedSize()); + } + return Error::success(); + } + JuliaJITLinkMemoryManager &MemMgr; + jitlink::LinkGraph &G; + jitlink::BasicLayout BL; + sys::MemoryBlock StandardSegments; + sys::MemoryBlock FinalizationSegments; + }; + public: + + /// Attempts to auto-detect the host page size. + /// Will preallocate and dispatch slabs of size MaxSize + static Expected> Create(uint64_t MaxSize); + + /// Create an instance using the given page size. + JuliaJITLinkMemoryManager(uint64_t PageSize, std::unique_ptr allocator) : PageSize(PageSize), allocator(std::move(allocator)) {} + + void allocate(const jitlink::JITLinkDylib *JD, jitlink::LinkGraph &G, + OnAllocatedFunction OnAllocated) override { + + // FIXME: Just check this once on startup. + if (!isPowerOf2_64((uint64_t)PageSize)) { + OnAllocated(make_error("Page size is not a power of 2", + inconvertibleErrorCode())); + return; + } + + jitlink::BasicLayout BL(G); + + /// Scan the request and calculate the group and total sizes. + /// Check that segment size is no larger than a page. + auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(PageSize); + if (!SegsSizes) { + OnAllocated(SegsSizes.takeError()); + return; + } + + /// Check that the total size requested (including zero fill) is not larger + /// than a size_t. + if (SegsSizes->total() > std::numeric_limits::max()) { + OnAllocated(make_error( + "Total requested size " + formatv("{0:x}", SegsSizes->total()) + + " for graph " + G.getName() + " exceeds address space")); + return; + } + + // Allocate one slab for the whole thing (to make sure everything is + // in-range), then partition into standard and finalization blocks. + // + // FIXME: Make two separate allocations in the future to reduce + // fragmentation: finalization segments will usually be a single page, and + // standard segments are likely to be more than one page. Where multiple + // allocations are in-flight at once (likely) the current approach will leave + // a lot of single-page holes. + auto allocation = allocator->allocate(JD, SegsSizes->StandardSegs, SegsSizes->FinalizeSegs, prot(sys::Memory::MF_READ, sys::Memory::MF_WRITE)); + if (!allocation) { + OnAllocated(allocation.takeError()); + } + + auto NextStandardSegAddr = orc::ExecutorAddr::fromPtr(allocation->StandardSegsMem.base()); + auto NextFinalizeSegAddr = orc::ExecutorAddr::fromPtr(allocation->FinalizeSegsMem.base()); + + // LLVM_DEBUG({ + // dbgs() << "InProcessMemoryManager allocated:\n"; + // if (SegsSizes->StandardSegs) + // dbgs() << formatv(" [ {0:x16} -- {1:x16} ]", NextStandardSegAddr, + // NextStandardSegAddr + StandardSegsMem.allocatedSize()) + // << " to stardard segs\n"; + // else + // dbgs() << " no standard segs\n"; + // if (SegsSizes->FinalizeSegs) + // dbgs() << formatv(" [ {0:x16} -- {1:x16} ]", NextFinalizeSegAddr, + // NextFinalizeSegAddr + FinalizeSegsMem.allocatedSize()) + // << " to finalize segs\n"; + // else + // dbgs() << " no finalize segs\n"; + // }); + + // Build ProtMap, assign addresses. + for (auto &KV : BL.segments()) { + auto &AG = KV.first; + auto &Seg = KV.second; + + auto &SegAddr = (AG.getMemDeallocPolicy() == jitlink::MemDeallocPolicy::Standard) + ? NextStandardSegAddr + : NextFinalizeSegAddr; + + Seg.WorkingMem = SegAddr.toPtr(); + Seg.Addr = SegAddr; + + SegAddr += alignTo(Seg.ContentSize + Seg.ZeroFillSize, PageSize); + } + + if (auto Err = BL.apply()) { + OnAllocated(std::move(Err)); + return; + } + + OnAllocated(std::make_unique(*this, G, std::move(BL), + std::move(allocation->StandardSegsMem), + std::move(allocation->FinalizeSegsMem))); + } + + // Use overloads from base class. + using JITLinkMemoryManager::allocate; + + void deallocate(std::vector Allocs, + OnDeallocatedFunction OnDeallocated) override { + std::vector StandardSegmentsList; + std::vector> DeallocActionsList; + + { + std::lock_guard Lock(FinalizedAllocsMutex); + for (auto &Alloc : Allocs) { + auto *FA = Alloc.release().toPtr(); + StandardSegmentsList.push_back(std::move(FA->StandardSegments)); + if (!FA->DeallocActions.empty()) + DeallocActionsList.push_back(std::move(FA->DeallocActions)); + FA->~FinalizedAllocInfo(); + FinalizedAllocInfos.Deallocate(FA); + } + } + + Error DeallocErr = Error::success(); + + while (!DeallocActionsList.empty()) { + auto &DeallocActions = DeallocActionsList.back(); + auto &StandardSegments = StandardSegmentsList.back(); + + /// Run any deallocate calls. + while (!DeallocActions.empty()) { + if (auto Err = DeallocActions.back().runWithSPSRetErrorMerged()) + DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err)); + DeallocActions.pop_back(); + } + + /// Release the standard segments slab. + if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments)) + DeallocErr = joinErrors(std::move(DeallocErr), errorCodeToError(EC)); + + DeallocActionsList.pop_back(); + StandardSegmentsList.pop_back(); + } + + OnDeallocated(std::move(DeallocErr)); + } + + // Use overloads from base class. + using JITLinkMemoryManager::deallocate; + + private: + struct FinalizedAllocInfo { + sys::MemoryBlock StandardSegments; + std::vector DeallocActions; + }; + + FinalizedAlloc createFinalizedAlloc( + sys::MemoryBlock StandardSegments, + std::vector DeallocActions) { + std::lock_guard Lock(FinalizedAllocsMutex); + auto *FA = FinalizedAllocInfos.Allocate(); + new (FA) FinalizedAllocInfo( + {std::move(StandardSegments), std::move(DeallocActions)}); + return FinalizedAlloc(orc::ExecutorAddr::fromPtr(FA)); + } + + uint64_t PageSize; + std::unique_ptr allocator; + std::mutex FinalizedAllocsMutex; + RecyclingAllocator FinalizedAllocInfos; + }; + + Expected> JuliaJITLinkMemoryManager::Create(uint64_t MaxSize) { + if (auto PageSize = sys::Process::getPageSize()) { + return std::make_unique(*PageSize, std::make_unique(*PageSize, MaxSize)); + } else { + return PageSize.takeError(); + } + } +} // namespace + +Expected> CreateJuliaJITLinkMemMgr() { + return JuliaJITLinkMemoryManager::Create(RESERVED_PAGE_SIZE); +} From a6822134415e69f91866933ac5330ae2905f3698 Mon Sep 17 00:00:00 2001 From: Prem Chintalapudi Date: Fri, 11 Mar 2022 22:01:15 -0500 Subject: [PATCH 4/4] CODLayer (broken yet again) --- src/jitlayers.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp index 29874c4dfca0d..eb17ab6d44793 100644 --- a/src/jitlayers.cpp +++ b/src/jitlayers.cpp @@ -1050,7 +1050,13 @@ void JuliaOJIT::addModule(std::unique_ptr 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)