(impala) 02/03: IMPALA-11542: Implement pre-allocation in LLVM memory manager

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

michaelsmith pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/impala.git

commit 67e4ff67cff6d6a77bfca43ef0c06b91f7a6661e
Author: Michael Smith <mi...@cloudera.com>
AuthorDate: Thu Nov 9 10:09:18 2023 -0800

    IMPALA-11542: Implement pre-allocation in LLVM memory manager
    
    Implements up-front allocation for the LLVM memory manager to avoid
    disparate sections on ARM which can exceed the 4GB limit for ADRP
    instructions and crash (or hit an assertion in debug mode).
    
    This is an LLVM issue that we're fixing by providing a custom memory
    manager. See the JIRA and
    https://discourse.llvm.org/t/llvm-rtdyld-aarch64-abi-relocation-restrictions/74616
    for further discussion.
    
    Testing:
    - passes debug test run on ARM
    - pre-commit tests
    
    Change-Id: I9f224edcdbdcb05fce663c18b4a5f03c8e985675
    Reviewed-on: http://gerrit.cloudera.org:8080/20692
    Tested-by: Michael Smith <mi...@cloudera.com>
    Reviewed-by: Joe McDonnell <jo...@cloudera.com>
---
 be/src/thirdparty/llvm/SectionMemoryManager.cpp | 105 ++++++++++++++++++++++++
 be/src/thirdparty/llvm/SectionMemoryManager.h   |  17 ++++
 2 files changed, 122 insertions(+)

diff --git a/be/src/thirdparty/llvm/SectionMemoryManager.cpp b/be/src/thirdparty/llvm/SectionMemoryManager.cpp
index a391a83b3..5964275ad 100644
--- a/be/src/thirdparty/llvm/SectionMemoryManager.cpp
+++ b/be/src/thirdparty/llvm/SectionMemoryManager.cpp
@@ -11,13 +11,118 @@
 // execution engine and RuntimeDyld
 //
 //===----------------------------------------------------------------------===//
+// Impala: Copied from the LLVM project to customize private portions of the
+// implementation.
 
 #include "SectionMemoryManager.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/Process.h"
 
+#include "common/logging.h"
+
 namespace impala {
 
+// ---- Impala: llvm/llvm-project#71968 ----
+bool SectionMemoryManager::hasSpace(const MemoryGroup &MemGroup,
+                                    uintptr_t Size) const {
+  for (const FreeMemBlock &FreeMB : MemGroup.FreeMem) {
+    if (FreeMB.Free.size() >= Size)
+      return true;
+  }
+  return false;
+}
+
+static uintptr_t alignTo(uintptr_t Size, uint32_t Alignment) {
+  return (Size + Alignment - 1) & ~(uintptr_t)(Alignment - 1);
+}
+
+static uint32_t checkAlignment(uint32_t Alignment, unsigned PageSize) {
+  DCHECK_GT(Alignment, 0);
+  DCHECK(!(Alignment & (Alignment - 1))) << "Alignment must be a power of two.";
+  DCHECK_LT(Alignment, PageSize);
+  // 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 uint32_t StubAlign = 8;
+  return std::max(Alignment, StubAlign);
+}
+
+void SectionMemoryManager::reserveAllocationSpace(
+    uintptr_t CodeSize, uint32_t CodeAlign, uintptr_t RODataSize, uint32_t RODataAlign,
+    uintptr_t RWDataSize, uint32_t RWDataAlign) {
+  if (CodeSize == 0 && RODataSize == 0 && RWDataSize == 0) return;
+
+  static const unsigned PageSize = sys::Process::getPageSize();
+
+  CodeAlign = checkAlignment(CodeAlign, PageSize);
+  RODataAlign = checkAlignment(RODataAlign, PageSize);
+  RWDataAlign = checkAlignment(RWDataAlign, PageSize);
+
+  // Get space required for each section. Use the same calculation as
+  // allocateSection because we need to be able to satisfy it.
+  uintptr_t RequiredCodeSize = alignTo(CodeSize, CodeAlign) + CodeAlign;
+  uintptr_t RequiredRODataSize = alignTo(RODataSize, RODataAlign) + RODataAlign;
+  uintptr_t RequiredRWDataSize = alignTo(RWDataSize, RWDataAlign) + RWDataAlign;
+
+  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);
+  uintptr_t RequiredSize = RequiredCodeSize + RequiredRODataSize + RequiredRWDataSize;
+
+  std::error_code ec;
+  sys::MemoryBlock MB = sys::Memory::allocateMappedMemory(RequiredSize, nullptr,
+      sys::Memory::MF_READ | sys::Memory::MF_WRITE, ec);
+  if (ec) {
+    return;
+  }
+  // Request is page-aligned, so we should always get back exactly the request.
+  DCHECK_EQ(MB.size(), RequiredSize);
+  // 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) {
+    DCHECK_EQ(Addr, alignTo(Addr, CodeAlign));
+    FreeMB.Free = sys::MemoryBlock((void*)Addr, RequiredCodeSize);
+    CodeMem.FreeMem.push_back(FreeMB);
+    Addr += RequiredCodeSize;
+  }
+
+  if (RODataSize > 0) {
+    DCHECK_EQ(Addr, alignTo(Addr, RODataAlign));
+    FreeMB.Free = sys::MemoryBlock((void*)Addr, RequiredRODataSize);
+    RODataMem.FreeMem.push_back(FreeMB);
+    Addr += RequiredRODataSize;
+  }
+
+  if (RWDataSize > 0) {
+    DCHECK_EQ(Addr, alignTo(Addr, RWDataAlign));
+    FreeMB.Free = sys::MemoryBlock((void*)Addr, RequiredRWDataSize);
+    RWDataMem.FreeMem.push_back(FreeMB);
+  }
+}
+// ---- End Impala changes ----
+
 uint8_t *SectionMemoryManager::allocateDataSection(uintptr_t Size,
                                                    unsigned Alignment,
                                                    unsigned SectionID,
diff --git a/be/src/thirdparty/llvm/SectionMemoryManager.h b/be/src/thirdparty/llvm/SectionMemoryManager.h
index d8ce1fc66..23ba5de28 100644
--- a/be/src/thirdparty/llvm/SectionMemoryManager.h
+++ b/be/src/thirdparty/llvm/SectionMemoryManager.h
@@ -51,6 +51,20 @@ public:
   void operator=(const SectionMemoryManager&) = delete;
   ~SectionMemoryManager() override;
 
+  /// Impala: enable reserveAllocationSpace callback.
+  bool needsToReserveAllocationSpace() override { return true; }
+
+  /// Impala: Provides an option to reserveAllocationSpace and pre-allocate all
+  /// memory in a single block. This is required for ARM where ADRP instructions
+  /// have a limit of 4GB offsets. Large memory systems may allocate sections
+  /// further apart than this unless we pre-allocate.
+  ///
+  /// Should only be called once. Later calls might re-use free blocks rather
+  /// than allocating a new contiguous block.
+  void reserveAllocationSpace(uintptr_t CodeSize, uint32_t CodeAlign,
+      uintptr_t RODataSize, uint32_t RODataAlign,
+      uintptr_t RWDataSize, uint32_t RWDataAlign) override;
+
   /// \brief Allocates a memory block of (at least) the given size suitable for
   /// executable code.
   ///
@@ -122,6 +136,9 @@ private:
   std::error_code applyMemoryGroupPermissions(MemoryGroup &MemGroup,
                                               unsigned Permissions);
 
+  // Impala: added to identify possible MemoryBlock re-use
+  bool hasSpace(const MemoryGroup &MemGroup, uintptr_t Size) const;
+
   MemoryGroup CodeMem;
   MemoryGroup RWDataMem;
   MemoryGroup RODataMem;