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

[Min-Yih Hsu via llvm-commits]

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

>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 1/2] [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))

>From 26bd4910f22bf8c586c6353406346566a264783f Mon Sep 17 00:00:00 2001
From: Min Hsu <min.hsu at sifive.com>
Date: Fri, 16 Feb 2024 09:05:49 -0800
Subject: [PATCH 2/2] Use empty list instead of std::optional in absent of
 DeallocAction

---
 .../JITLink/JITLinkMemoryManager.cpp          | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
index 2edeb6e2195dc1..dacf0e6c8aa4f2 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
@@ -442,18 +442,14 @@ 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::optional<std::vector<orc::shared::WrapperFunctionCall>>>
-      DeallocActionsList;
+  std::vector<std::vector<orc::shared::WrapperFunctionCall>> DeallocActionsList;
 
   {
     std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
     for (auto &Alloc : Allocs) {
       auto *FA = Alloc.release().toPtr<FinalizedAllocInfo *>();
       StandardSegmentsList.push_back(std::move(FA->StandardSegments));
-      if (!FA->DeallocActions.empty())
-        DeallocActionsList.push_back(std::move(FA->DeallocActions));
-      else
-        DeallocActionsList.push_back(std::nullopt);
+      DeallocActionsList.push_back(std::move(FA->DeallocActions));
       FA->~FinalizedAllocInfo();
       FinalizedAllocInfos.Deallocate(FA);
     }
@@ -466,12 +462,11 @@ void InProcessMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
     auto &StandardSegments = StandardSegmentsList.back();
 
     /// Run any deallocate calls.
-    if (DeallocActions.has_value())
-      while (!DeallocActions->empty()) {
-        if (auto Err = DeallocActions->back().runWithSPSRetErrorMerged())
-          DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
-        DeallocActions->pop_back();
-      }
+    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))