Skip to content

Commit d4a4137

Browse files
lhameslhames
authored
[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.
1 parent 765060b commit d4a4137

File tree

3 files changed

+123
-29
lines changed

3 files changed

+123
-29
lines changed

‎orc-rt/include/orc-rt/SimpleNativeMemoryMap.h‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,13 @@ class SimpleNativeMemoryMap : public ResourceManager {
5050

5151
/// Release a slab of contiguous address space back to the system.
5252
using OnReleaseCompleteFn = move_only_function<void(Error)>;
53-
void release(OnReleaseCompleteFn &&OnComplete, void *Addr);
53+
void release(OnReleaseCompleteFn &&OnComplete, void *Addrs);
54+
55+
/// Convenience method to release multiple slabs with one call. This can be
56+
/// used to save on interprocess communication at the cost of less expressive
57+
/// errors.
58+
void releaseMultiple(OnReleaseCompleteFn &&OnComplete,
59+
std::vector<void *> Addrs);
5460

5561
struct FinalizeRequest {
5662
struct Segment {
@@ -74,6 +80,12 @@ class SimpleNativeMemoryMap : public ResourceManager {
7480
using OnDeallocateCompleteFn = move_only_function<void(Error)>;
7581
void deallocate(OnDeallocateCompleteFn &&OnComplete, void *Base);
7682

83+
/// Convenience method to deallocate multiple regions with one call. This can
84+
/// be used to save on interprocess communication at the cost of less
85+
/// expressive errors.
86+
void deallocateMultiple(OnDeallocateCompleteFn &&OnComplete,
87+
std::vector<void *> Bases);
88+
7789
void detach(ResourceManager::OnCompleteFn OnComplete) override;
7890
void shutdown(ResourceManager::OnCompleteFn OnComplete) override;
7991

@@ -84,6 +96,10 @@ class SimpleNativeMemoryMap : public ResourceManager {
8496
std::unordered_map<void *, std::vector<AllocAction>> DeallocActions;
8597
};
8698

99+
void releaseNext(OnReleaseCompleteFn &&OnComplete, std::vector<void *> Addrs,
100+
bool AnyError, Error LastErr);
101+
void deallocateNext(OnDeallocateCompleteFn &&OnComplete,
102+
std::vector<void *> Bases, bool AnyError, Error LastErr);
87103
void shutdownNext(OnCompleteFn OnComplete, std::vector<void *> Bases);
88104
Error makeBadSlabError(void *Base, const char *Op);
89105
SlabInfo *findSlabInfoFor(void *Base);
@@ -100,15 +116,17 @@ ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper(
100116
orc_rt_SessionRef Session, void *CallCtx,
101117
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);
102118

103-
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_release_sps_wrapper(
119+
ORC_RT_SPS_INTERFACE void
120+
orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper(
104121
orc_rt_SessionRef Session, void *CallCtx,
105122
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);
106123

107124
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper(
108125
orc_rt_SessionRef Session, void *CallCtx,
109126
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);
110127

111-
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper(
128+
ORC_RT_SPS_INTERFACE void
129+
orc_rt_SimpleNativeMemoryMap_deallocateMultiple_sps_wrapper(
112130
orc_rt_SessionRef Session, void *CallCtx,
113131
orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes);
114132

‎orc-rt/lib/executor/SimpleNativeMemoryMap.cpp‎

Lines changed: 82 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,11 @@ void SimpleNativeMemoryMap::release(OnReleaseCompleteFn &&OnComplete,
116116
OnComplete(hostOSMemoryRelease(Addr, SI->Size));
117117
}
118118

119+
void SimpleNativeMemoryMap::releaseMultiple(OnReleaseCompleteFn &&OnComplete,
120+
std::vector<void *> Addrs) {
121+
releaseNext(std::move(OnComplete), std::move(Addrs), false, Error::success());
122+
}
123+
119124
void SimpleNativeMemoryMap::finalize(OnFinalizeCompleteFn &&OnComplete,
120125
FinalizeRequest FR) {
121126

@@ -207,6 +212,12 @@ void SimpleNativeMemoryMap::deallocate(OnDeallocateCompleteFn &&OnComplete,
207212
OnComplete(Error::success());
208213
}
209214

215+
void SimpleNativeMemoryMap::deallocateMultiple(
216+
OnDeallocateCompleteFn &&OnComplete, std::vector<void *> Bases) {
217+
deallocateNext(std::move(OnComplete), std::move(Bases), false,
218+
Error::success());
219+
}
220+
210221
void SimpleNativeMemoryMap::detach(ResourceManager::OnCompleteFn OnComplete) {
211222
// Detach is a noop for now: we just retain all actions to run at shutdown
212223
// time.
@@ -228,6 +239,64 @@ void SimpleNativeMemoryMap::shutdown(ResourceManager::OnCompleteFn OnComplete) {
228239
shutdownNext(std::move(OnComplete), std::move(Bases));
229240
}
230241

242+
void SimpleNativeMemoryMap::releaseNext(OnReleaseCompleteFn &&OnComplete,
243+
std::vector<void *> Addrs,
244+
bool AnyError, Error LastErr) {
245+
// TODO: Log error?
246+
if (LastErr) {
247+
consumeError(std::move(LastErr));
248+
AnyError |= true;
249+
}
250+
251+
if (Addrs.empty()) {
252+
if (!AnyError)
253+
return OnComplete(Error::success());
254+
255+
return OnComplete(
256+
make_error<StringError>("Failed to release some addresses"));
257+
}
258+
259+
void *NextAddr = Addrs.back();
260+
Addrs.pop_back();
261+
262+
release(
263+
[this, OnComplete = std::move(OnComplete), AnyError = AnyError,
264+
Addrs = std::move(Addrs)](Error Err) mutable {
265+
releaseNext(std::move(OnComplete), std::move(Addrs), AnyError,
266+
std::move(Err));
267+
},
268+
NextAddr);
269+
}
270+
271+
void SimpleNativeMemoryMap::deallocateNext(OnDeallocateCompleteFn &&OnComplete,
272+
std::vector<void *> Addrs,
273+
bool AnyError, Error LastErr) {
274+
// TODO: Log error?
275+
if (LastErr) {
276+
consumeError(std::move(LastErr));
277+
AnyError |= true;
278+
}
279+
280+
if (Addrs.empty()) {
281+
if (!AnyError)
282+
return OnComplete(Error::success());
283+
284+
return OnComplete(
285+
make_error<StringError>("Failed to deallocate some addresses"));
286+
}
287+
288+
void *NextAddr = Addrs.back();
289+
Addrs.pop_back();
290+
291+
deallocate(
292+
[this, OnComplete = std::move(OnComplete), AnyError = AnyError,
293+
Addrs = std::move(Addrs)](Error Err) mutable {
294+
deallocateNext(std::move(OnComplete), std::move(Addrs), AnyError,
295+
std::move(Err));
296+
},
297+
NextAddr);
298+
}
299+
231300
void SimpleNativeMemoryMap::shutdownNext(
232301
ResourceManager::OnCompleteFn OnComplete, std::vector<void *> Bases) {
233302
if (Bases.empty())
@@ -303,14 +372,15 @@ ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper(
303372
WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::reserve));
304373
}
305374

306-
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_release_sps_wrapper(
375+
ORC_RT_SPS_INTERFACE void
376+
orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper(
307377
orc_rt_SessionRef Session, void *CallCtx,
308378
orc_rt_WrapperFunctionReturn Return,
309379
orc_rt_WrapperFunctionBuffer ArgBytes) {
310-
using Sig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
311-
SPSWrapperFunction<Sig>::handle(
312-
Session, CallCtx, Return, ArgBytes,
313-
WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::release));
380+
using Sig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
381+
SPSWrapperFunction<Sig>::handle(Session, CallCtx, Return, ArgBytes,
382+
WrapperFunction::handleWithAsyncMethod(
383+
&SimpleNativeMemoryMap::releaseMultiple));
314384
}
315385

316386
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(
324394
WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::finalize));
325395
}
326396

327-
ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper(
397+
ORC_RT_SPS_INTERFACE void
398+
orc_rt_SimpleNativeMemoryMap_deallocateMultiple_sps_wrapper(
328399
orc_rt_SessionRef Session, void *CallCtx,
329400
orc_rt_WrapperFunctionReturn Return,
330401
orc_rt_WrapperFunctionBuffer ArgBytes) {
331-
using Sig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
332-
SPSWrapperFunction<Sig>::handle(Session, CallCtx, Return, ArgBytes,
333-
WrapperFunction::handleWithAsyncMethod(
334-
&SimpleNativeMemoryMap::deallocate));
402+
using Sig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
403+
SPSWrapperFunction<Sig>::handle(
404+
Session, CallCtx, Return, ArgBytes,
405+
WrapperFunction::handleWithAsyncMethod(
406+
&SimpleNativeMemoryMap::deallocateMultiple));
335407
}
336408

337409
} // namespace orc_rt

‎orc-rt/unittests/SimpleNativeMemoryMapTest.cpp‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,17 @@ static void snmm_reserve(OnCompleteFn &&OnComplete,
106106
std::forward<OnCompleteFn>(OnComplete), Instance, Size);
107107
}
108108

109+
template <typename OnCompleteFn>
110+
static void snmm_releaseMultiple(OnCompleteFn &&OnComplete,
111+
SimpleNativeMemoryMap *Instance,
112+
span<void *> Addr) {
113+
using SPSSig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
114+
SPSWrapperFunction<SPSSig>::call(
115+
DirectCaller(nullptr,
116+
orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper),
117+
std::forward<OnCompleteFn>(OnComplete), Instance, Addr);
118+
}
119+
109120
template <typename OnCompleteFn>
110121
static void snmm_finalize(OnCompleteFn &&OnComplete,
111122
SimpleNativeMemoryMap *Instance,
@@ -118,24 +129,16 @@ static void snmm_finalize(OnCompleteFn &&OnComplete,
118129
}
119130

120131
template <typename OnCompleteFn>
121-
static void snmm_deallocate(OnCompleteFn &&OnComplete,
122-
SimpleNativeMemoryMap *Instance, void *Base) {
123-
using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
132+
static void snmm_deallocateMultiple(OnCompleteFn &&OnComplete,
133+
SimpleNativeMemoryMap *Instance,
134+
span<void *> Base) {
135+
using SPSSig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);
124136
SPSWrapperFunction<SPSSig>::call(
125137
DirectCaller(nullptr,
126-
orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper),
138+
orc_rt_SimpleNativeMemoryMap_deallocateMultiple_sps_wrapper),
127139
std::forward<OnCompleteFn>(OnComplete), Instance, Base);
128140
}
129141

130-
template <typename OnCompleteFn>
131-
static void snmm_release(OnCompleteFn &&OnComplete,
132-
SimpleNativeMemoryMap *Instance, void *Addr) {
133-
using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
134-
SPSWrapperFunction<SPSSig>::call(
135-
DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_release_sps_wrapper),
136-
std::forward<OnCompleteFn>(OnComplete), Instance, Addr);
137-
}
138-
139142
TEST(SimpleNativeMemoryMapTest, ReserveAndRelease) {
140143
// Test that we can reserve and release a slab of address space as expected,
141144
// without finalizing any memory within it.
@@ -145,7 +148,7 @@ TEST(SimpleNativeMemoryMapTest, ReserveAndRelease) {
145148
auto Addr = cantFail(cantFail(ReserveAddr.get()));
146149

147150
std::future<Expected<Error>> ReleaseResult;
148-
snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
151+
snmm_releaseMultiple(waitFor(ReleaseResult), SNMM.get(), {&Addr, 1});
149152
cantFail(cantFail(ReleaseResult.get()));
150153
}
151154

@@ -239,15 +242,16 @@ TEST(SimpleNativeMemoryMap, FullPipelineForOneRWSegment) {
239242
EXPECT_EQ(SentinelValue3, 0U);
240243

241244
std::future<Expected<Error>> DeallocResult;
242-
snmm_deallocate(waitFor(DeallocResult), SNMM.get(), FinalizeKeyAddr);
245+
snmm_deallocateMultiple(waitFor(DeallocResult), SNMM.get(),
246+
{&FinalizeKeyAddr, 1});
243247
cantFail(cantFail(DeallocResult.get()));
244248

245249
EXPECT_EQ(SentinelValue1, 42U);
246250
EXPECT_EQ(SentinelValue2, 42U);
247251
EXPECT_EQ(SentinelValue3, 0U);
248252

249253
std::future<Expected<Error>> ReleaseResult;
250-
snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
254+
snmm_releaseMultiple(waitFor(ReleaseResult), SNMM.get(), {&Addr, 1});
251255
cantFail(cantFail(ReleaseResult.get()));
252256
}
253257

0 commit comments

Comments
 (0)