1 files changed, 40 insertions, 10 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
index ecc6793bbce..9e0d207e8bd 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
@@ -38,9 +38,21 @@ InProcessMemoryManager::allocate(const SegmentsRequestMap &Request) {
       OnFinalize(applyProtections());
     }
     Error deallocate() override {
-      for (auto &KV : SegBlocks)
-        if (auto EC = sys::Memory::releaseMappedMemory(KV.second))
-          return errorCodeToError(EC);
+      if (SegBlocks.empty())
+        return Error::success();
+      void *SlabStart = SegBlocks.begin()->second.base();
+      char *SlabEnd = (char *)SlabStart;
+      for (auto &KV : SegBlocks) {
+        SlabStart = std::min(SlabStart, KV.second.base());
+        SlabEnd = std::max(SlabEnd, (char *)(KV.second.base()) +
+                                        KV.second.allocatedSize());
+      }
+      size_t SlabSize = SlabEnd - (char *)SlabStart;
+      assert((SlabSize % sys::Process::getPageSizeEstimate()) == 0 &&
+             "Slab size is not a multiple of page size");
+      sys::MemoryBlock Slab(SlabStart, SlabSize);
+      if (auto EC = sys::Memory::releaseMappedMemory(Slab))
+        return errorCodeToError(EC);
       return Error::success();
     }
 
@@ -70,22 +82,40 @@ InProcessMemoryManager::allocate(const SegmentsRequestMap &Request) {
       static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
                                                 sys::Memory::MF_WRITE);
 
+  // Compute the total number of pages to allocate.
+  size_t TotalSize = 0;
   for (auto &KV : Request) {
-    auto &Seg = KV.second;
+    const auto &Seg = KV.second;
 
     if (Seg.getAlignment() > sys::Process::getPageSizeEstimate())
       return make_error<StringError>("Cannot request higher than page "
                                      "alignment",
                                      inconvertibleErrorCode());
 
-    uint64_t SegmentSize = Seg.getContentSize() + Seg.getZeroFillSize();
+    TotalSize = alignTo(TotalSize, sys::Process::getPageSizeEstimate());
+    TotalSize += Seg.getContentSize();
+    TotalSize += Seg.getZeroFillSize();
+  }
+
+  // Allocate one slab to cover all the segments.
+  std::error_code EC;
+  auto SlabRemaining =
+      sys::Memory::allocateMappedMemory(TotalSize, nullptr, ReadWrite, EC);
+
+  if (EC)
+    return errorCodeToError(EC);
+
+  // Allocate segment memory from the slab.
+  for (auto &KV : Request) {
+
+    const auto &Seg = KV.second;
 
-    std::error_code EC;
-    auto SegMem =
-        sys::Memory::allocateMappedMemory(SegmentSize, nullptr, ReadWrite, EC);
+    uint64_t SegmentSize = alignTo(Seg.getContentSize() + Seg.getZeroFillSize(),
+                                   sys::Process::getPageSizeEstimate());
 
-    if (EC)
-      return errorCodeToError(EC);
+    sys::MemoryBlock SegMem(SlabRemaining.base(), SegmentSize);
+    SlabRemaining = sys::MemoryBlock((char *)SlabRemaining.base() + SegmentSize,
+                                     SegmentSize);
 
     // Zero out the zero-fill memory.
     memset(static_cast<char *>(SegMem.base()) + Seg.getContentSize(), 0,