[JITLink] Switch to slab allocation for InProcessMemoryManager, re-enable test.

InProcessMemoryManager used to make separate memory allocation calls for each
permission level (RW, RX, RO), which could lead to target-out-of-range errors
if data and code were placed too far apart (this was the source of failures in
the JITLink/AArch64 testcase when it was first landed).

This patch updates InProcessMemoryManager to allocate a single slab which is
subdivided between text and data. This should guarantee that accesses remain
in-range provided that individual object files do not exceed 1Mb in size.
This patch also re-enables the JITLink/AArch64 testcase.

llvm-svn: 374948

2 files changed, 42 insertions, 12 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,
diff --git a/llvm/test/ExecutionEngine/JITLink/AArch64/lit.local.cfg b/llvm/test/ExecutionEngine/JITLink/AArch64/lit.local.cfg
index 7cfadc4db86..304f2436ece 100644
--- a/llvm/test/ExecutionEngine/JITLink/AArch64/lit.local.cfg
+++ b/llvm/test/ExecutionEngine/JITLink/AArch64/lit.local.cfg
@@ -1,2 +1,2 @@
-# if not 'AArch64' in config.root.targets:
-config.unsupported = True
+if not 'AArch64' in config.root.targets:
+  config.unsupported = True