Implement reserveAllocationSpace for SectionMemoryManager

Implements `reserveAllocationSpace` and provides an option to enable
`needsToReserveAllocationSpace` for large-memory environments with
AArch64.

The [AArch64 ABI](https://github.com/ARM-software/abi-aa/blob/main/sysvabi64/sysvabi64.rst)
has limits on the distance between sections as the instructions to
reference them are limited to 2 or 4GB. Allocating sections in multiple
blocks can result in distances greater than that on systems with lots of
memory. In those environments several projects using
SectionMemoryManager with MCJIT have run across assertion failures for
the R_AARCH64_ADR_PREL_PG_HI21 instruction as it attempts to address
across distances greater than 2GB (an int32).

Fixes #71963 by allocating all sections in a single contiguous memory
allocation, limiting the distance required for instruction offsets
similar to how pre-compiled binaries would be loaded into memory. Does
not change the default behavior of SectionMemoryManager.

Author: MikaelSmith

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();