diff --git a/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h b/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h index 6dbc0c0b6b910..0e516eeaaea60 100644 --- a/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h +++ b/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h @@ -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 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 Addrs); struct FinalizeRequest { struct Segment { @@ -74,6 +80,12 @@ class SimpleNativeMemoryMap : public ResourceManager { using OnDeallocateCompleteFn = move_only_function; 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 Bases); + void detach(ResourceManager::OnCompleteFn OnComplete) override; void shutdown(ResourceManager::OnCompleteFn OnComplete) override; @@ -84,6 +96,10 @@ class SimpleNativeMemoryMap : public ResourceManager { std::unordered_map> DeallocActions; }; + void releaseNext(OnReleaseCompleteFn &&OnComplete, std::vector Addrs, + bool AnyError, Error LastErr); + void deallocateNext(OnDeallocateCompleteFn &&OnComplete, + std::vector Bases, bool AnyError, Error LastErr); void shutdownNext(OnCompleteFn OnComplete, std::vector 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); diff --git a/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp b/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp index 10cdcf5d8add8..987bd857defe6 100644 --- a/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp +++ b/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp @@ -116,6 +116,11 @@ void SimpleNativeMemoryMap::release(OnReleaseCompleteFn &&OnComplete, OnComplete(hostOSMemoryRelease(Addr, SI->Size)); } +void SimpleNativeMemoryMap::releaseMultiple(OnReleaseCompleteFn &&OnComplete, + std::vector 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 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 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("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 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("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 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::handle( - Session, CallCtx, Return, ArgBytes, - WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::release)); + using Sig = SPSError(SPSExecutorAddr, SPSSequence); + SPSWrapperFunction::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::handle(Session, CallCtx, Return, ArgBytes, - WrapperFunction::handleWithAsyncMethod( - &SimpleNativeMemoryMap::deallocate)); + using Sig = SPSError(SPSExecutorAddr, SPSSequence); + SPSWrapperFunction::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 b7ef7f08fbc01..c54d79154cc30 100644 --- a/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp +++ b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp @@ -106,6 +106,17 @@ static void snmm_reserve(OnCompleteFn &&OnComplete, std::forward(OnComplete), Instance, Size); } +template +static void snmm_releaseMultiple(OnCompleteFn &&OnComplete, + SimpleNativeMemoryMap *Instance, + span Addr) { + using SPSSig = SPSError(SPSExecutorAddr, SPSSequence); + SPSWrapperFunction::call( + DirectCaller(nullptr, + orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper), + std::forward(OnComplete), Instance, Addr); +} + template static void snmm_finalize(OnCompleteFn &&OnComplete, SimpleNativeMemoryMap *Instance, @@ -118,24 +129,16 @@ static void snmm_finalize(OnCompleteFn &&OnComplete, } template -static void snmm_deallocate(OnCompleteFn &&OnComplete, - SimpleNativeMemoryMap *Instance, void *Base) { - using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr); +static void snmm_deallocateMultiple(OnCompleteFn &&OnComplete, + SimpleNativeMemoryMap *Instance, + span Base) { + using SPSSig = SPSError(SPSExecutorAddr, SPSSequence); SPSWrapperFunction::call( DirectCaller(nullptr, - orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper), + orc_rt_SimpleNativeMemoryMap_deallocateMultiple_sps_wrapper), std::forward(OnComplete), Instance, Base); } -template -static void snmm_release(OnCompleteFn &&OnComplete, - SimpleNativeMemoryMap *Instance, void *Addr) { - using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr); - SPSWrapperFunction::call( - DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_release_sps_wrapper), - std::forward(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> 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> 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> ReleaseResult; - snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr); + snmm_releaseMultiple(waitFor(ReleaseResult), SNMM.get(), {&Addr, 1}); cantFail(cantFail(ReleaseResult.get())); }