Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 12, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 21 additions & 3 deletions orc-rt/include/orc-rt/SimpleNativeMemoryMap.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,13 @@ class SimpleNativeMemoryMap : public ResourceManager {

/// 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 {
Expand All @@ -74,6 +80,12 @@ class SimpleNativeMemoryMap : public ResourceManager {
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;

Expand All @@ -84,6 +96,10 @@ class SimpleNativeMemoryMap : public ResourceManager {
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);
Expand All @@ -100,15 +116,17 @@ 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);

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);

Expand Down
92 changes: 82 additions & 10 deletions orc-rt/lib/executor/SimpleNativeMemoryMap.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {

Expand Down Expand Up @@ -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.
Expand All @@ -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())
Expand Down Expand Up @@ -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(
Expand All @@ -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
36 changes: 20 additions & 16 deletions orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Expand All @@ -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()));
}

Expand Down Expand Up @@ -239,15 +242,16 @@ 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);
EXPECT_EQ(SentinelValue2, 42U);
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()));
}

Expand Down
Loading