[orc-rt] Add multi-addr dealloc/release to SimpleNativeMemoryMap. (#163025)

In an ORC JIT it's common for multiple memory regions to be deallocated
at once, e.g. when a ResourceTracker covering multiple object files is
removed. This commit adds SimpleNativeMemoryMap::deallocateMultiple and
SimpleNativeMemoryMap::releaseMultiple APIs that can be used to reduce
the number of calls (and consequently IPC messages in cross-process
setups) in these cases.

Adding these operations will make it easier to write an
llvm::orc::MemoryMapper class that can use SimpleNativeMemoryMap as a
backend.
This commit is contained in:
Lang Hames
2025-10-12 12:03:11 +11:00
committed by GitHub
parent 765060be88
commit d4a4137976
3 changed files with 123 additions and 29 deletions

View File

@@ -50,7 +50,13 @@ public:
/// Release a slab of contiguous address space back to the system.
using OnReleaseCompleteFn = move_only_function<void(Error)>;
void release(OnReleaseCompleteFn &&OnComplete, void *Addr);
void release(OnReleaseCompleteFn &&OnComplete, void *Addrs);
/// Convenience method to release multiple slabs with one call. This can be
/// used to save on interprocess communication at the cost of less expressive
/// errors.
void releaseMultiple(OnReleaseCompleteFn &&OnComplete,
std::vector<void *> Addrs);
struct FinalizeRequest {
struct Segment {
@@ -74,6 +80,12 @@ public:
using OnDeallocateCompleteFn = move_only_function<void(Error)>;
void deallocate(OnDeallocateCompleteFn &&OnComplete, void *Base);
/// Convenience method to deallocate multiple regions with one call. This can
/// be used to save on interprocess communication at the cost of less
/// expressive errors.
void deallocateMultiple(OnDeallocateCompleteFn &&OnComplete,
std::vector<void *> Bases);
void detach(ResourceManager::OnCompleteFn OnComplete) override;
void shutdown(ResourceManager::OnCompleteFn OnComplete) override;
@@ -84,6 +96,10 @@ private:
std::unordered_map<void *, std::vector<AllocAction>> DeallocActions;
};
void releaseNext(OnReleaseCompleteFn &&OnComplete, std::vector<void *> Addrs,
bool AnyError, Error LastErr);
void deallocateNext(OnDeallocateCompleteFn &&OnComplete,
std::vector<void *> Bases, bool AnyError, Error LastErr);
void shutdownNext(OnCompleteFn OnComplete, std::vector<void *> Bases);
Error makeBadSlabError(void *Base, const char *Op);
SlabInfo *findSlabInfoFor(void *Base);
@@ -100,7 +116,8 @@ ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper(
orc_rt_SessionRef Session, void *CallCtx,
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_release_sps_wrapper(
ORC_RT_SPS_INTERFACE void
orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper(
orc_rt_SessionRef Session, void *CallCtx,
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);
@@ -108,7 +125,8 @@ ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper(
orc_rt_SessionRef Session, void *CallCtx,
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper(
ORC_RT_SPS_INTERFACE void
orc_rt_SimpleNativeMemoryMap_deallocateMultiple_sps_wrapper(
orc_rt_SessionRef Session, void *CallCtx,
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);

View File

@@ -116,6 +116,11 @@ void SimpleNativeMemoryMap::release(OnReleaseCompleteFn &&OnComplete,
OnComplete(hostOSMemoryRelease(Addr, SI->Size));
}
void SimpleNativeMemoryMap::releaseMultiple(OnReleaseCompleteFn &&OnComplete,
std::vector<void *> Addrs) {
releaseNext(std::move(OnComplete), std::move(Addrs), false, Error::success());
}
void SimpleNativeMemoryMap::finalize(OnFinalizeCompleteFn &&OnComplete,
FinalizeRequest FR) {
@@ -207,6 +212,12 @@ void SimpleNativeMemoryMap::deallocate(OnDeallocateCompleteFn &&OnComplete,
OnComplete(Error::success());
}
void SimpleNativeMemoryMap::deallocateMultiple(
OnDeallocateCompleteFn &&OnComplete, std::vector<void *> Bases) {
deallocateNext(std::move(OnComplete), std::move(Bases), false,
Error::success());
}
void SimpleNativeMemoryMap::detach(ResourceManager::OnCompleteFn OnComplete) {
// Detach is a noop for now: we just retain all actions to run at shutdown
// time.
@@ -228,6 +239,64 @@ void SimpleNativeMemoryMap::shutdown(ResourceManager::OnCompleteFn OnComplete) {
shutdownNext(std::move(OnComplete), std::move(Bases));
}
void SimpleNativeMemoryMap::releaseNext(OnReleaseCompleteFn &&OnComplete,
std::vector<void *> Addrs,
bool AnyError, Error LastErr) {
// TODO: Log error?
if (LastErr) {
consumeError(std::move(LastErr));
AnyError |= true;
}
if (Addrs.empty()) {
if (!AnyError)
return OnComplete(Error::success());
return OnComplete(
make_error<StringError>("Failed to release some addresses"));
}
void *NextAddr = Addrs.back();
Addrs.pop_back();
release(
[this, OnComplete = std::move(OnComplete), AnyError = AnyError,
Addrs = std::move(Addrs)](Error Err) mutable {
releaseNext(std::move(OnComplete), std::move(Addrs), AnyError,
std::move(Err));
},
NextAddr);
}
void SimpleNativeMemoryMap::deallocateNext(OnDeallocateCompleteFn &&OnComplete,
std::vector<void *> Addrs,
bool AnyError, Error LastErr) {
// TODO: Log error?
if (LastErr) {
consumeError(std::move(LastErr));
AnyError |= true;
}
if (Addrs.empty()) {
if (!AnyError)
return OnComplete(Error::success());
return OnComplete(
make_error<StringError>("Failed to deallocate some addresses"));
}
void *NextAddr = Addrs.back();
Addrs.pop_back();
deallocate(
[this, OnComplete = std::move(OnComplete), AnyError = AnyError,
Addrs = std::move(Addrs)](Error Err) mutable {
deallocateNext(std::move(OnComplete), std::move(Addrs), AnyError,
std::move(Err));
},
NextAddr);
}
void SimpleNativeMemoryMap::shutdownNext(
ResourceManager::OnCompleteFn OnComplete, std::vector<void *> Bases) {
if (Bases.empty())
@@ -303,14 +372,15 @@ ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper(
WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::reserve));
}
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_release_sps_wrapper(
ORC_RT_SPS_INTERFACE void
orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper(
orc_rt_SessionRef Session, void *CallCtx,
orc_rt_WrapperFunctionReturn Return,
orc_rt_WrapperFunctionBuffer ArgBytes) {
using Sig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
SPSWrapperFunction<Sig>::handle(
Session, CallCtx, Return, ArgBytes,
WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::release));
using Sig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
SPSWrapperFunction<Sig>::handle(Session, CallCtx, Return, ArgBytes,
WrapperFunction::handleWithAsyncMethod(
&SimpleNativeMemoryMap::releaseMultiple));
}
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper(
@@ -324,14 +394,16 @@ ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper(
WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::finalize));
}
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper(
ORC_RT_SPS_INTERFACE void
orc_rt_SimpleNativeMemoryMap_deallocateMultiple_sps_wrapper(
orc_rt_SessionRef Session, void *CallCtx,
orc_rt_WrapperFunctionReturn Return,
orc_rt_WrapperFunctionBuffer ArgBytes) {
using Sig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
SPSWrapperFunction<Sig>::handle(Session, CallCtx, Return, ArgBytes,
WrapperFunction::handleWithAsyncMethod(
&SimpleNativeMemoryMap::deallocate));
using Sig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
SPSWrapperFunction<Sig>::handle(
Session, CallCtx, Return, ArgBytes,
WrapperFunction::handleWithAsyncMethod(
&SimpleNativeMemoryMap::deallocateMultiple));
}
} // namespace orc_rt

View File

@@ -106,6 +106,17 @@ static void snmm_reserve(OnCompleteFn &&OnComplete,
std::forward<OnCompleteFn>(OnComplete), Instance, Size);
}
template <typename OnCompleteFn>
static void snmm_releaseMultiple(OnCompleteFn &&OnComplete,
SimpleNativeMemoryMap *Instance,
span<void *> Addr) {
using SPSSig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
SPSWrapperFunction<SPSSig>::call(
DirectCaller(nullptr,
orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper),
std::forward<OnCompleteFn>(OnComplete), Instance, Addr);
}
template <typename OnCompleteFn>
static void snmm_finalize(OnCompleteFn &&OnComplete,
SimpleNativeMemoryMap *Instance,
@@ -118,24 +129,16 @@ static void snmm_finalize(OnCompleteFn &&OnComplete,
}
template <typename OnCompleteFn>
static void snmm_deallocate(OnCompleteFn &&OnComplete,
SimpleNativeMemoryMap *Instance, void *Base) {
using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
static void snmm_deallocateMultiple(OnCompleteFn &&OnComplete,
SimpleNativeMemoryMap *Instance,
span<void *> Base) {
using SPSSig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
SPSWrapperFunction<SPSSig>::call(
DirectCaller(nullptr,
orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper),
orc_rt_SimpleNativeMemoryMap_deallocateMultiple_sps_wrapper),
std::forward<OnCompleteFn>(OnComplete), Instance, Base);
}
template <typename OnCompleteFn>
static void snmm_release(OnCompleteFn &&OnComplete,
SimpleNativeMemoryMap *Instance, void *Addr) {
using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
SPSWrapperFunction<SPSSig>::call(
DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_release_sps_wrapper),
std::forward<OnCompleteFn>(OnComplete), Instance, Addr);
}
TEST(SimpleNativeMemoryMapTest, ReserveAndRelease) {
// Test that we can reserve and release a slab of address space as expected,
// without finalizing any memory within it.
@@ -145,7 +148,7 @@ TEST(SimpleNativeMemoryMapTest, ReserveAndRelease) {
auto Addr = cantFail(cantFail(ReserveAddr.get()));
std::future<Expected<Error>> ReleaseResult;
snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
snmm_releaseMultiple(waitFor(ReleaseResult), SNMM.get(), {&Addr, 1});
cantFail(cantFail(ReleaseResult.get()));
}
@@ -239,7 +242,8 @@ TEST(SimpleNativeMemoryMap, FullPipelineForOneRWSegment) {
EXPECT_EQ(SentinelValue3, 0U);
std::future<Expected<Error>> DeallocResult;
snmm_deallocate(waitFor(DeallocResult), SNMM.get(), FinalizeKeyAddr);
snmm_deallocateMultiple(waitFor(DeallocResult), SNMM.get(),
{&FinalizeKeyAddr, 1});
cantFail(cantFail(DeallocResult.get()));
EXPECT_EQ(SentinelValue1, 42U);
@@ -247,7 +251,7 @@ TEST(SimpleNativeMemoryMap, FullPipelineForOneRWSegment) {
EXPECT_EQ(SentinelValue3, 0U);
std::future<Expected<Error>> ReleaseResult;
snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
snmm_releaseMultiple(waitFor(ReleaseResult), SNMM.get(), {&Addr, 1});
cantFail(cantFail(ReleaseResult.get()));
}