llvm-project/llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryMana...

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

185 lines
6.7 KiB
C++
Raw Normal View History

//===---- EPCGenericJITLinkMemoryManager.cpp -- Mem management via EPC ----===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "llvm/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.h"
#include "llvm/ExecutionEngine/JITLink/JITLink.h"
#include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
#include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
#include <limits>
using namespace llvm::jitlink;
namespace llvm {
namespace orc {
class EPCGenericJITLinkMemoryManager::InFlightAlloc
: public jitlink::JITLinkMemoryManager::InFlightAlloc {
public:
// FIXME: The C++98 initializer is an attempt to work around compile failures
// due to http://www.open-std.org/jtc1/sc22/wg21/docs/cwg_defects.html#1397.
// We should be able to switch this back to member initialization once that
// issue is fixed.
struct SegInfo {
SegInfo() : WorkingMem(nullptr), ContentSize(0), ZeroFillSize(0) {}
char *WorkingMem;
ExecutorAddr Addr;
uint64_t ContentSize;
uint64_t ZeroFillSize;
};
using SegInfoMap = AllocGroupSmallMap<SegInfo>;
InFlightAlloc(EPCGenericJITLinkMemoryManager &Parent, LinkGraph &G,
ExecutorAddr AllocAddr, SegInfoMap Segs)
: Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
void finalize(OnFinalizedFunction OnFinalize) override {
tpctypes::FinalizeRequest FR;
for (auto &KV : Segs) {
assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
FR.Segments.push_back(tpctypes::SegFinalizeRequest{
tpctypes::toWireProtectionFlags(
toSysMemoryProtectionFlags(KV.first.getMemProt())),
KV.second.Addr,
alignTo(KV.second.ContentSize + KV.second.ZeroFillSize,
Parent.EPC.getPageSize()),
{KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
}
// Transfer allocation actions.
// FIXME: Merge JITLink and ORC SupportFunctionCall and Action list types,
// turn this into a std::swap.
FR.Actions.reserve(G.allocActions().size());
for (auto &ActPair : G.allocActions())
FR.Actions.push_back({{ExecutorAddr(ActPair.Finalize.FnAddr),
{ExecutorAddr(ActPair.Finalize.CtxAddr),
ExecutorAddrDiff(ActPair.Finalize.CtxSize)}},
{ExecutorAddr(ActPair.Dealloc.FnAddr),
{ExecutorAddr(ActPair.Dealloc.CtxAddr),
ExecutorAddrDiff(ActPair.Dealloc.CtxSize)}}});
G.allocActions().clear();
Parent.EPC.callSPSWrapperAsync<
rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
Parent.SAs.Finalize,
[OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
Error SerializationErr, Error FinalizeErr) mutable {
// FIXME: Release abandoned alloc.
if (SerializationErr) {
cantFail(std::move(FinalizeErr));
OnFinalize(std::move(SerializationErr));
} else if (FinalizeErr)
OnFinalize(std::move(FinalizeErr));
else
OnFinalize(FinalizedAlloc(AllocAddr.getValue()));
},
Parent.SAs.Allocator, std::move(FR));
}
void abandon(OnAbandonedFunction OnAbandoned) override {
// FIXME: Return memory to pool instead.
Parent.EPC.callSPSWrapperAsync<
rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
Parent.SAs.Deallocate,
[OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
Error DeallocateErr) mutable {
if (SerializationErr) {
cantFail(std::move(DeallocateErr));
OnAbandoned(std::move(SerializationErr));
} else
OnAbandoned(std::move(DeallocateErr));
},
Parent.SAs.Allocator, ArrayRef<ExecutorAddr>(AllocAddr));
}
private:
EPCGenericJITLinkMemoryManager &Parent;
LinkGraph &G;
ExecutorAddr AllocAddr;
SegInfoMap Segs;
};
void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
LinkGraph &G,
OnAllocatedFunction OnAllocated) {
BasicLayout BL(G);
auto Pages = BL.getContiguousPageBasedLayoutSizes(EPC.getPageSize());
if (!Pages)
return OnAllocated(Pages.takeError());
EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
SAs.Reserve,
[this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
if (SerializationErr) {
cantFail(AllocAddr.takeError());
return OnAllocated(std::move(SerializationErr));
}
if (!AllocAddr)
return OnAllocated(AllocAddr.takeError());
completeAllocation(*AllocAddr, std::move(BL), std::move(OnAllocated));
},
SAs.Allocator, Pages->total());
}
void EPCGenericJITLinkMemoryManager::deallocate(
std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
EPC.callSPSWrapperAsync<
rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
SAs.Deallocate,
[OnDeallocated = std::move(OnDeallocated)](Error SerErr,
Error DeallocErr) mutable {
if (SerErr) {
cantFail(std::move(DeallocErr));
OnDeallocated(std::move(SerErr));
} else
OnDeallocated(std::move(DeallocErr));
},
SAs.Allocator, Allocs);
for (auto &A : Allocs)
A.release();
}
[JITLink][ORC] Major JITLinkMemoryManager refactor. This commit substantially refactors the JITLinkMemoryManager API to: (1) add asynchronous versions of key operations, (2) give memory manager implementations full control over link graph address layout, (3) enable more efficient tracking of allocated memory, and (4) support "allocation actions" and finalize-lifetime memory. Together these changes provide a more usable API, and enable more powerful and efficient memory manager implementations. To support these changes the JITLinkMemoryManager::Allocation inner class has been split into two new classes: InFlightAllocation, and FinalizedAllocation. The allocate method returns an InFlightAllocation that tracks memory (both working and executor memory) prior to finalization. The finalize method returns a FinalizedAllocation object, and the InFlightAllocation is discarded. Breaking Allocation into InFlightAllocation and FinalizedAllocation allows InFlightAllocation subclassses to be written more naturally, and FinalizedAlloc to be implemented and used efficiently (see (3) below). In addition to the memory manager changes this commit also introduces a new MemProt type to represent memory protections (MemProt replaces use of sys::Memory::ProtectionFlags in JITLink), and a new MemDeallocPolicy type that can be used to indicate when a section should be deallocated (see (4) below). Plugin/pass writers who were using sys::Memory::ProtectionFlags will have to switch to MemProt -- this should be straightworward. Clients with out-of-tree memory managers will need to update their implementations. Clients using in-tree memory managers should mostly be able to ignore it. Major features: (1) More asynchrony: The allocate and deallocate methods are now asynchronous by default, with synchronous convenience wrappers supplied. The asynchronous versions allow clients (including JITLink) to request and deallocate memory without blocking. (2) Improved control over graph address layout: Instead of a SegmentRequestMap, JITLinkMemoryManager::allocate now takes a reference to the LinkGraph to be allocated. The memory manager is responsible for calculating the memory requirements for the graph, and laying out the graph (setting working and executor memory addresses) within the allocated memory. This gives memory managers full control over JIT'd memory layout. For clients that don't need or want this degree of control the new "BasicLayout" utility can be used to get a segment-based view of the graph, similar to the one provided by SegmentRequestMap. Once segment addresses are assigned the BasicLayout::apply method can be used to automatically lay out the graph. (3) Efficient tracking of allocated memory. The FinalizedAlloc type is a wrapper for an ExecutorAddr and requires only 64-bits to store in the controller. The meaning of the address held by the FinalizedAlloc is left up to the memory manager implementation, but the FinalizedAlloc type enforces a requirement that deallocate be called on any non-default values prior to destruction. The deallocate method takes a vector<FinalizedAlloc>, allowing for bulk deallocation of many allocations in a single call. Memory manager implementations will typically store the address of some allocation metadata in the executor in the FinalizedAlloc, as holding this metadata in the executor is often cheaper and may allow for clean deallocation even in failure cases where the connection with the controller is lost. (4) Support for "allocation actions" and finalize-lifetime memory. Allocation actions are pairs (finalize_act, deallocate_act) of JITTargetAddress triples (fn, arg_buffer_addr, arg_buffer_size), that can be attached to a finalize request. At finalization time, after memory protections have been applied, each of the "finalize_act" elements will be called in order (skipping any elements whose fn value is zero) as ((char*(*)(const char *, size_t))fn)((const char *)arg_buffer_addr, (size_t)arg_buffer_size); At deallocation time the deallocate elements will be run in reverse order (again skipping any elements where fn is zero). The returned char * should be null to indicate success, or a non-null heap-allocated string error message to indicate failure. These actions allow finalization and deallocation to be extended to include operations like registering and deregistering eh-frames, TLS sections, initializer and deinitializers, and language metadata sections. Previously these operations required separate callWrapper invocations. Compared to callWrapper invocations, actions require no extra IPC/RPC, reducing costs and eliminating a potential source of errors. Finalize lifetime memory can be used to support finalize actions: Sections with finalize lifetime should be destroyed by memory managers immediately after finalization actions have been run. Finalize memory can be used to support finalize actions (e.g. with extra-metadata, or synthesized finalize actions) without incurring permanent memory overhead.
2021-10-11 08:39:24 +08:00
void EPCGenericJITLinkMemoryManager::completeAllocation(
ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
InFlightAlloc::SegInfoMap SegInfos;
ExecutorAddr NextSegAddr = AllocAddr;
for (auto &KV : BL.segments()) {
const auto &AG = KV.first;
auto &Seg = KV.second;
Seg.Addr = NextSegAddr.getValue();
KV.second.WorkingMem = BL.getGraph().allocateBuffer(Seg.ContentSize).data();
NextSegAddr += ExecutorAddrDiff(
alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize()));
auto &SegInfo = SegInfos[AG];
SegInfo.ContentSize = Seg.ContentSize;
SegInfo.ZeroFillSize = Seg.ZeroFillSize;
SegInfo.Addr = ExecutorAddr(Seg.Addr);
SegInfo.WorkingMem = Seg.WorkingMem;
}
if (auto Err = BL.apply())
return OnAllocated(std::move(Err));
OnAllocated(std::make_unique<InFlightAlloc>(*this, BL.getGraph(), AllocAddr,
std::move(SegInfos)));
}
} // end namespace orc
} // end namespace llvm