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

lhames · web-flow · commit d4a413797681 · 2025-10-12T12:03:11.000+11:00
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
@@ -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 {
@@ -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;
 
@@ -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);
@@ -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);
 
diff --git 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<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
diff --git a/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
@@ -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,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()));
 }
 
