[llvm] [JITLink] Always unmap standard segments in InProcessMemoryManager::deallocate (PR #81943)

[Min-Yih Hsu via llvm-commits]

https://github.com/mshockwave created https://github.com/llvm/llvm-project/pull/81943

Right now InProcessMemoryManager only releases a standard segment (via sys::Memory::releaseMappedMemory) in `deallocate` when there is a DeallocAction associated, leaving residual memory pages in the process until termination.
Despite being a de facto memory leak, it won't cause a major issue if users only create a single LLJIT instance per process, which is the most common use cases. It will, however, drain virtual memory pages if we create thousands of ephemeral LLJIT instances in the same process.

This patch fixes this issue by releasing every standard segments regardless of the attached DeallocAction.

>From 77ce3105fbee09ed1d082be872cb4b8c31700bdd Mon Sep 17 00:00:00 2001
From: Min Hsu <min.hsu at sifive.com>
Date: Thu, 15 Feb 2024 15:40:08 -0800
Subject: [PATCH] [JITLink] Always unmap standard segments in
 InProcessMemoryManager::deallocate

Right now InProcessMemoryManager only releases a standard segment (via
sys::Memory::releaseMappedMemory) in `deallocate` when there is a
DeallocAction associated, leaving residual memory pages in the process
until termination.
Despite being a de facto memory leak, it won't cause a major issue if
users only create a single LLJIT instance per process, which is the most
common use cases. It will, however, drain virtual memory pages if we create
thousands of ephemeral LLJIT instances in the same process.

This patch fixes this issue by releasing every standard segments
regardless of the attached DeallocAction.
---
 .../JITLink/JITLinkMemoryManager.cpp             | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
index 474a0b5160bcb0..2edeb6e2195dc1 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
@@ -442,7 +442,8 @@ void InProcessMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
 void InProcessMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
                                         OnDeallocatedFunction OnDeallocated) {
   std::vector<sys::MemoryBlock> StandardSegmentsList;
-  std::vector<std::vector<orc::shared::WrapperFunctionCall>> DeallocActionsList;
+  std::vector<std::optional<std::vector<orc::shared::WrapperFunctionCall>>>
+      DeallocActionsList;
 
   {
     std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
@@ -451,6 +452,8 @@ void InProcessMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
       StandardSegmentsList.push_back(std::move(FA->StandardSegments));
       if (!FA->DeallocActions.empty())
         DeallocActionsList.push_back(std::move(FA->DeallocActions));
+      else
+        DeallocActionsList.push_back(std::nullopt);
       FA->~FinalizedAllocInfo();
       FinalizedAllocInfos.Deallocate(FA);
     }
@@ -463,11 +466,12 @@ void InProcessMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
     auto &StandardSegments = StandardSegmentsList.back();
 
     /// Run any deallocate calls.
-    while (!DeallocActions.empty()) {
-      if (auto Err = DeallocActions.back().runWithSPSRetErrorMerged())
-        DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
-      DeallocActions.pop_back();
-    }
+    if (DeallocActions.has_value())
+      while (!DeallocActions->empty()) {
+        if (auto Err = DeallocActions->back().runWithSPSRetErrorMerged())
+          DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
+        DeallocActions->pop_back();
+      }
 
     /// Release the standard segments slab.
     if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))