Implement reserveAllocationSpace for SectionMemoryManager

llvm/include/llvm/ExecutionEngine/SectionMemoryManager.h

@@ -104,11 +104,24 @@ class SectionMemoryManager : public RTDyldMemoryManager {
   /// Creates a SectionMemoryManager instance with \p MM as the associated
   /// memory mapper.  If \p MM is nullptr then a default memory mapper is used
   /// that directly calls into the operating system.
-  SectionMemoryManager(MemoryMapper *MM = nullptr);
+  ///
+  /// If \p ReserveAlloc is true all memory will be pre-allocated, and any
+  /// attempts to allocate beyond pre-allocated memory will fail.
+  SectionMemoryManager(MemoryMapper *MM = nullptr, bool ReserveAlloc = false);
   SectionMemoryManager(const SectionMemoryManager &) = delete;
   void operator=(const SectionMemoryManager &) = delete;
   ~SectionMemoryManager() override;
 
+  /// Enable reserveAllocationSpace when requested.
+  bool needsToReserveAllocationSpace() override { return ReserveAllocation; }
+
+  /// Implements allocating all memory in a single block. This is required to
+  /// limit memory offsets to fit the ARM ABI; large memory systems may
+  /// otherwise allocate separate sections too far apart.
+  void reserveAllocationSpace(uintptr_t CodeSize, Align CodeAlign,
+                              uintptr_t RODataSize, Align RODataAlign,
+                              uintptr_t RWDataSize, Align RWDataAlign) override;
+
   /// Allocates a memory block of (at least) the given size suitable for
   /// executable code.
   ///
@@ -180,13 +193,16 @@ class SectionMemoryManager : public RTDyldMemoryManager {
   std::error_code applyMemoryGroupPermissions(MemoryGroup &MemGroup,
                                               unsigned Permissions);
 
+  bool hasSpace(const MemoryGroup &MemGroup, uintptr_t Size) const;
+
   void anchor() override;
 
   MemoryGroup CodeMem;
   MemoryGroup RWDataMem;
   MemoryGroup RODataMem;
   MemoryMapper *MMapper;
   std::unique_ptr<MemoryMapper> OwnedMMapper;
+  bool ReserveAllocation;
 };
 
 } // end namespace llvm

llvm/lib/ExecutionEngine/SectionMemoryManager.cpp

@@ -18,6 +18,97 @@
 
 namespace llvm {
 
+bool SectionMemoryManager::hasSpace(const MemoryGroup &MemGroup,
+                                    uintptr_t Size) const {
+  for (const FreeMemBlock &FreeMB : MemGroup.FreeMem) {
+    if (FreeMB.Free.allocatedSize() >= Size)
+      return true;
+  }
+  return false;
+}
+
+void SectionMemoryManager::reserveAllocationSpace(
+    uintptr_t CodeSize, Align CodeAlign, uintptr_t RODataSize,
+    Align RODataAlign, uintptr_t RWDataSize, Align RWDataAlign) {
+  if (CodeSize == 0 && RODataSize == 0 && RWDataSize == 0)
+    return;
+
+  static const size_t PageSize = sys::Process::getPageSizeEstimate();
+
+  // Code alignment needs to be at least the stub alignment - however, we
+  // don't have an easy way to get that here so as a workaround, we assume
+  // it's 8, which is the largest value I observed across all platforms.
+  constexpr uint64_t StubAlign = 8;
+  CodeAlign = Align(std::max(CodeAlign.value(), StubAlign));
+  RODataAlign = Align(std::max(RODataAlign.value(), StubAlign));
+  RWDataAlign = Align(std::max(RWDataAlign.value(), StubAlign));
+
+  // Get space required for each section. Use the same calculation as
+  // allocateSection because we need to be able to satisfy it.
+  uint64_t RequiredCodeSize = alignTo(CodeSize, CodeAlign) + CodeAlign.value();
+  uint64_t RequiredRODataSize =
+      alignTo(RODataSize, RODataAlign) + RODataAlign.value();
+  uint64_t RequiredRWDataSize =
+      alignTo(RWDataSize, RWDataAlign) + RWDataAlign.value();
+
+  if (hasSpace(CodeMem, RequiredCodeSize) &&
+      hasSpace(RODataMem, RequiredRODataSize) &&
+      hasSpace(RWDataMem, RequiredRWDataSize)) {
+    // Sufficient space in contiguous block already available.
+    return;
+  }
+
+  // MemoryManager does not have functions for releasing memory after it's
+  // allocated. Normally it tries to use any excess blocks that were allocated
+  // due to page alignment, but if we have insufficient free memory for the
+  // request this can lead to allocating disparate memory that can violate the
+  // ARM ABI. Clear free memory so only the new allocations are used, but do
+  // not release allocated memory as it may still be in-use.
+  CodeMem.FreeMem.clear();
+  RODataMem.FreeMem.clear();
+  RWDataMem.FreeMem.clear();
+
+  // Round up to the nearest page size. Blocks must be page-aligned.
+  RequiredCodeSize = alignTo(RequiredCodeSize, PageSize);
+  RequiredRODataSize = alignTo(RequiredRODataSize, PageSize);
+  RequiredRWDataSize = alignTo(RequiredRWDataSize, PageSize);
+  uint64_t RequiredSize =
+      RequiredCodeSize + RequiredRODataSize + RequiredRWDataSize;
+
+  std::error_code ec;
+  sys::MemoryBlock MB = MMapper->allocateMappedMemory(
+      AllocationPurpose::RWData, RequiredSize, nullptr,
+      sys::Memory::MF_READ | sys::Memory::MF_WRITE, ec);
+  if (ec) {
+    return;
+  }
+  // CodeMem will arbitrarily own this MemoryBlock to handle cleanup.
+  CodeMem.AllocatedMem.push_back(MB);
+  uintptr_t Addr = (uintptr_t)MB.base();
+  FreeMemBlock FreeMB;
+  FreeMB.PendingPrefixIndex = (unsigned)-1;
+
+  if (CodeSize > 0) {
+    assert(isAddrAligned(CodeAlign, (void *)Addr));
+    FreeMB.Free = sys::MemoryBlock((void *)Addr, RequiredCodeSize);
+    CodeMem.FreeMem.push_back(FreeMB);
+    Addr += RequiredCodeSize;
+  }
+
+  if (RODataSize > 0) {
+    assert(isAddrAligned(RODataAlign, (void *)Addr));
+    FreeMB.Free = sys::MemoryBlock((void *)Addr, RequiredRODataSize);
+    RODataMem.FreeMem.push_back(FreeMB);
+    Addr += RequiredRODataSize;
+  }
+
+  if (RWDataSize > 0) {
+    assert(isAddrAligned(RWDataAlign, (void *)Addr));
+    FreeMB.Free = sys::MemoryBlock((void *)Addr, RequiredRWDataSize);
+    RWDataMem.FreeMem.push_back(FreeMB);
+  }
+}
+
 uint8_t *SectionMemoryManager::allocateDataSection(uintptr_t Size,
                                                    unsigned Alignment,
                                                    unsigned SectionID,
@@ -265,8 +356,10 @@ class DefaultMMapper final : public SectionMemoryManager::MemoryMapper {
 };
 } // namespace
 
-SectionMemoryManager::SectionMemoryManager(MemoryMapper *UnownedMM)
-    : MMapper(UnownedMM), OwnedMMapper(nullptr) {
+SectionMemoryManager::SectionMemoryManager(MemoryMapper *UnownedMM,
+                                           bool ReserveAlloc)
+    : MMapper(UnownedMM), OwnedMMapper(nullptr),
+      ReserveAllocation(ReserveAlloc) {
   if (!MMapper) {
     OwnedMMapper = std::make_unique<DefaultMMapper>();
     MMapper = OwnedMMapper.get();