[llvm] a31af32 - Reapply [Orc] Properly deallocate mapped memory in MapperJITLinkMemoryManager

[Anubhab Ghosh via llvm-commits]

Author: Anubhab Ghosh
Date: 2022-08-13T13:07:50+05:30
New Revision: a31af32183856ad6c3f24f5b0e89e1366180a780

URL: https://github.com/llvm/llvm-project/commit/a31af32183856ad6c3f24f5b0e89e1366180a780
DIFF: https://github.com/llvm/llvm-project/commit/a31af32183856ad6c3f24f5b0e89e1366180a780.diff

LOG: Reapply [Orc] Properly deallocate mapped memory in MapperJITLinkMemoryManager

When memory is deallocated from MapperJITLinkMemoryManager deinitialize
actions are run through mapper and in case of InProcessMapper, memory
protections of the region are reset to read/write as they were previously
changed and can be reused in future.

Differential Revision: https://reviews.llvm.org/D131768

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h
    llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
    llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
    llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h
index 0b4cda119cadd..80dac5aee80bc 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h
@@ -102,6 +102,7 @@ class InProcessMemoryMapper final : public MemoryMapper {
 
 private:
   struct Allocation {
+    size_t Size;
     std::vector<shared::WrapperFunctionCall> DeinitializationActions;
   };
   using AllocationMap = DenseMap<ExecutorAddr, Allocation>;

diff  --git a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
index b9328f0a16edf..bd773a10d9174 100644
--- a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
@@ -147,17 +147,35 @@ void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
 
 void MapperJITLinkMemoryManager::deallocate(
     std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
-  std::lock_guard<std::mutex> Lock(Mutex);
-
+  std::vector<ExecutorAddr> Bases;
+  Bases.reserve(Allocs.size());
   for (auto &FA : Allocs) {
     ExecutorAddr Addr = FA.getAddress();
-    ExecutorAddrDiff Size = UsedMemory[Addr];
-    UsedMemory.erase(Addr);
-
-    AvailableMemory.push_back({Addr, Addr + Size});
-    FA.release();
+    Bases.push_back(Addr);
   }
-  OnDeallocated(Error::success());
+
+  Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs),
+                               OnDeallocated = std::move(OnDeallocated)](
+                                  llvm::Error Err) mutable {
+    if (Err)
+      OnDeallocated(std::move(Err));
+
+    {
+      std::lock_guard<std::mutex> Lock(Mutex);
+
+      for (auto &FA : Allocs) {
+        ExecutorAddr Addr = FA.getAddress();
+        ExecutorAddrDiff Size = UsedMemory[Addr];
+
+        UsedMemory.erase(Addr);
+        AvailableMemory.push_back({Addr, Addr + Size});
+
+        FA.release();
+      }
+    }
+
+    OnDeallocated(Error::success());
+  });
 }
 
 } // end namespace orc

diff  --git a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
index ee92e5191b507..1e8ad27342d72 100644
--- a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
@@ -60,6 +60,7 @@ char *InProcessMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
 void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
                                        OnInitializedFunction OnInitialized) {
   ExecutorAddr MinAddr(~0ULL);
+  ExecutorAddr MaxAddr(0);
 
   for (auto &Segment : AI.Segments) {
     auto Base = AI.MappingBase + Segment.Offset;
@@ -68,6 +69,9 @@ void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
     if (Base < MinAddr)
       MinAddr = Base;
 
+    if (Base + Size > MaxAddr)
+      MaxAddr = Base + Size;
+
     std::memset((Base + Segment.ContentSize).toPtr<void *>(), 0,
                 Segment.ZeroFillSize);
 
@@ -85,6 +89,9 @@ void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
 
   {
     std::lock_guard<std::mutex> Lock(Mutex);
+
+    // This is the maximum range whose permission have been possibly modified
+    Allocations[MinAddr].Size = MaxAddr - MinAddr;
     Allocations[MinAddr].DeinitializationActions =
         std::move(*DeinitializeActions);
     Reservations[AI.MappingBase.toPtr<void *>()].Allocations.push_back(MinAddr);
@@ -108,6 +115,14 @@ void InProcessMemoryMapper::deinitialize(
         AllErr = joinErrors(std::move(AllErr), std::move(Err));
       }
 
+      // Reset protections to read/write so the area can be reused
+      if (auto EC = sys::Memory::protectMappedMemory(
+              {Base.toPtr<void *>(), Allocations[Base].Size},
+              sys::Memory::ProtectionFlags::MF_READ |
+                  sys::Memory::ProtectionFlags::MF_WRITE)) {
+        AllErr = joinErrors(std::move(AllErr), errorCodeToError(EC));
+      }
+
       Allocations.erase(Base);
     }
   }

diff  --git a/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp
index 74d8387701359..ca1b13c759f63 100644
--- a/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp
@@ -118,11 +118,17 @@ TEST(MapperJITLinkMemoryManagerTest, InProcess) {
   StringRef TargetHello2(TargetMem2, Hello.size());
   EXPECT_EQ(Hello, TargetHello2);
 
+  EXPECT_EQ(Counter->DeinitCount, 0);
+
   auto Err2 = MemMgr->deallocate(std::move(*FA1));
   EXPECT_THAT_ERROR(std::move(Err2), Succeeded());
 
+  EXPECT_EQ(Counter->DeinitCount, 1);
+
   auto Err3 = MemMgr->deallocate(std::move(*FA2));
   EXPECT_THAT_ERROR(std::move(Err3), Succeeded());
+
+  EXPECT_EQ(Counter->DeinitCount, 2);
 }
 
 } // namespace