diff --git a/include/llvm/ExecutionEngine/JITLink/JITLink.h b/include/llvm/ExecutionEngine/JITLink/JITLink.h index f0b4d9bcd49..48e64335613 100644 --- a/include/llvm/ExecutionEngine/JITLink/JITLink.h +++ b/include/llvm/ExecutionEngine/JITLink/JITLink.h @@ -1270,9 +1270,15 @@ class JITLinkContext { public: using LookupMap = DenseMap; + /// Create a JITLinkContext. + JITLinkContext(const JITLinkDylib *JD) : JD(JD) {} + /// Destroy a JITLinkContext. virtual ~JITLinkContext(); + /// Return the JITLinkDylib that this link is targeting, if any. + const JITLinkDylib *getJITLinkDylib() const { return JD; } + /// Return the MemoryManager to be used for this link. virtual JITLinkMemoryManager &getMemoryManager() = 0; @@ -1324,6 +1330,9 @@ public: /// Called by JITLink to modify the pass pipeline prior to linking. /// The default version performs no modification. virtual Error modifyPassConfig(const Triple &TT, PassConfiguration &Config); + +private: + const JITLinkDylib *JD = nullptr; }; /// Marks all symbols in a graph live. This can be used as a default, diff --git a/include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h b/include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h index a3dc6c1a700..cee7d6b09c4 100644 --- a/include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h +++ b/include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h @@ -14,10 +14,11 @@ #define LLVM_EXECUTIONENGINE_JITLINK_JITLINKMEMORYMANAGER_H #include "llvm/ADT/DenseMap.h" +#include "llvm/ExecutionEngine/JITLink/JITLinkDylib.h" #include "llvm/ExecutionEngine/JITSymbol.h" #include "llvm/Support/Error.h" -#include "llvm/Support/Memory.h" #include "llvm/Support/MSVCErrorWorkarounds.h" +#include "llvm/Support/Memory.h" #include #include @@ -93,15 +94,25 @@ public: virtual ~JITLinkMemoryManager(); /// Create an Allocation object. + /// + /// The JD argument represents the target JITLinkDylib, and can be used by + /// JITLinkMemoryManager implementers to manage per-dylib allocation pools + /// (e.g. one pre-reserved address space slab per dylib to ensure that all + /// allocations for the dylib are within a certain range). The JD argument + /// may be null (representing an allocation not associated with any + /// JITDylib. + /// + /// The request argument describes the segment sizes and permisssions being + /// requested. virtual Expected> - allocate(const SegmentsRequestMap &Request) = 0; + allocate(const JITLinkDylib *JD, const SegmentsRequestMap &Request) = 0; }; /// A JITLinkMemoryManager that allocates in-process memory. class InProcessMemoryManager : public JITLinkMemoryManager { public: Expected> - allocate(const SegmentsRequestMap &Request) override; + allocate(const JITLinkDylib *JD, const SegmentsRequestMap &Request) override; }; } // end namespace jitlink diff --git a/include/llvm/ExecutionEngine/Orc/Core.h b/include/llvm/ExecutionEngine/Orc/Core.h index 6256872b109..3020694ee73 100644 --- a/include/llvm/ExecutionEngine/Orc/Core.h +++ b/include/llvm/ExecutionEngine/Orc/Core.h @@ -17,6 +17,7 @@ #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/FunctionExtras.h" #include "llvm/ADT/IntrusiveRefCntPtr.h" +#include "llvm/ExecutionEngine/JITLink/JITLinkDylib.h" #include "llvm/ExecutionEngine/JITSymbol.h" #include "llvm/ExecutionEngine/Orc/SymbolStringPool.h" #include "llvm/ExecutionEngine/OrcV1Deprecation.h" @@ -887,7 +888,8 @@ public: /// their addresses may be used as keys for resource management. /// JITDylib state changes must be made via an ExecutionSession to guarantee /// that they are synchronized with respect to other JITDylib operations. -class JITDylib : public ThreadSafeRefCountedBase { +class JITDylib : public ThreadSafeRefCountedBase, + public jitlink::JITLinkDylib { friend class AsynchronousSymbolQuery; friend class ExecutionSession; friend class Platform; diff --git a/include/llvm/ExecutionEngine/Orc/OrcRPCTargetProcessControl.h b/include/llvm/ExecutionEngine/Orc/OrcRPCTargetProcessControl.h index 0bc209120ce..1856e0d7820 100644 --- a/include/llvm/ExecutionEngine/Orc/OrcRPCTargetProcessControl.h +++ b/include/llvm/ExecutionEngine/Orc/OrcRPCTargetProcessControl.h @@ -148,7 +148,8 @@ public: OrcRPCTPCJITLinkMemoryManager(OrcRPCTPCImplT &Parent) : Parent(Parent) {} Expected> - allocate(const SegmentsRequestMap &Request) override { + allocate(const jitlink::JITLinkDylib *JD, + const SegmentsRequestMap &Request) override { orcrpctpc::ReserveMemRequest RMR; HostAllocMap HostAllocs; diff --git a/include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h b/include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h index b255ddd6416..dfed5e09e4d 100644 --- a/include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h +++ b/include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h @@ -535,7 +535,8 @@ public: } Expected> - allocate(const SegmentsRequestMap &Request) override { + allocate(const jitlink::JITLinkDylib *JD, + const SegmentsRequestMap &Request) override { return RPCMMAlloc::Create(Client, Id, Request); } diff --git a/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp index 7c60d516771..393ed0ea709 100644 --- a/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp +++ b/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp @@ -266,7 +266,8 @@ Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) { } LLVM_DEBUG(dbgs() << " }\n"); - if (auto AllocOrErr = Ctx->getMemoryManager().allocate(Segments)) + if (auto AllocOrErr = + Ctx->getMemoryManager().allocate(Ctx->getJITLinkDylib(), Segments)) Alloc = std::move(*AllocOrErr); else return AllocOrErr.takeError(); diff --git a/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp index 68ec9d79af9..fbbb29e9164 100644 --- a/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp +++ b/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp @@ -17,7 +17,8 @@ JITLinkMemoryManager::~JITLinkMemoryManager() = default; JITLinkMemoryManager::Allocation::~Allocation() = default; Expected> -InProcessMemoryManager::allocate(const SegmentsRequestMap &Request) { +InProcessMemoryManager::allocate(const JITLinkDylib *JD, + const SegmentsRequestMap &Request) { using AllocationMap = DenseMap; diff --git a/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp index 5dd32fbc224..f8eb1dcfce4 100644 --- a/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp +++ b/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp @@ -28,7 +28,8 @@ public: ObjectLinkingLayer &Layer, std::unique_ptr MR, std::unique_ptr ObjBuffer) - : Layer(Layer), MR(std::move(MR)), ObjBuffer(std::move(ObjBuffer)) {} + : JITLinkContext(&MR->getTargetJITDylib()), Layer(Layer), + MR(std::move(MR)), ObjBuffer(std::move(ObjBuffer)) {} ~ObjectLinkingLayerJITLinkContext() { // If there is an object buffer return function then use it to diff --git a/lib/ExecutionEngine/Orc/TPCIndirectionUtils.cpp b/lib/ExecutionEngine/Orc/TPCIndirectionUtils.cpp index d827cf7b273..7989ec41952 100644 --- a/lib/ExecutionEngine/Orc/TPCIndirectionUtils.cpp +++ b/lib/ExecutionEngine/Orc/TPCIndirectionUtils.cpp @@ -109,7 +109,7 @@ Error TPCTrampolinePool::grow() { jitlink::JITLinkMemoryManager::SegmentsRequestMap Request; Request[TrampolinePagePermissions] = {PageSize, static_cast(PageSize), 0}; - auto Alloc = TPC.getMemMgr().allocate(Request); + auto Alloc = TPC.getMemMgr().allocate(nullptr, Request); if (!Alloc) return Alloc.takeError(); @@ -294,7 +294,7 @@ TPCIndirectionUtils::writeResolverBlock(JITTargetAddress ReentryFnAddr, jitlink::JITLinkMemoryManager::SegmentsRequestMap Request; Request[ResolverBlockPermissions] = {TPC.getPageSize(), static_cast(ResolverSize), 0}; - auto Alloc = TPC.getMemMgr().allocate(Request); + auto Alloc = TPC.getMemMgr().allocate(nullptr, Request); if (!Alloc) return Alloc.takeError(); @@ -364,7 +364,7 @@ TPCIndirectionUtils::getIndirectStubs(unsigned NumStubs) { Request[StubPagePermissions] = {PageSize, static_cast(StubBytes), 0}; Request[PointerPagePermissions] = {PageSize, 0, PointerBytes}; - auto Alloc = TPC.getMemMgr().allocate(Request); + auto Alloc = TPC.getMemMgr().allocate(nullptr, Request); if (!Alloc) return Alloc.takeError(); diff --git a/tools/llvm-jitlink/llvm-jitlink.cpp b/tools/llvm-jitlink/llvm-jitlink.cpp index 1ba5259988f..808a7db9c10 100644 --- a/tools/llvm-jitlink/llvm-jitlink.cpp +++ b/tools/llvm-jitlink/llvm-jitlink.cpp @@ -330,7 +330,7 @@ public: } Expected> - allocate(const SegmentsRequestMap &Request) override { + allocate(const JITLinkDylib *JD, const SegmentsRequestMap &Request) override { using AllocationMap = DenseMap;