[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.
diff --git a/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h b/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h
index 6dbc0c0..0e516ee 100644
--- a/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h
+++ b/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h
@@ -50,7 +50,13 @@
/// 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 @@
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 @@
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_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_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);
diff --git a/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp b/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp
index 10cdcf5..987bd85 100644
--- a/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp
+++ b/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp
@@ -116,6 +116,11 @@
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 @@
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 @@
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 @@
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 @@
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
diff --git a/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
index b7ef7f0..c54d791 100644
--- a/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
+++ b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
@@ -107,6 +107,17 @@
}
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,
TestSNMMFinalizeRequest FR) {
@@ -118,24 +129,16 @@
}
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 @@
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 @@
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 @@
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()));
}
