[llvm] Implement reserveAllocationSpace for SectionMemoryManager (PR #71968)
Michael Smith via llvm-commits
llvm-commits at lists.llvm.org
Tue Nov 14 08:42:23 PST 2023
https://github.com/MikaelSmith updated https://github.com/llvm/llvm-project/pull/71968
>From 237a5215a6710d7ed68050893fb81efa6bd20d6e Mon Sep 17 00:00:00 2001
From: Michael Smith <michael.smith at cloudera.com>
Date: Fri, 10 Nov 2023 09:48:28 -0800
Subject: [PATCH] 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.
---
.../ExecutionEngine/SectionMemoryManager.h | 18 +-
.../ExecutionEngine/SectionMemoryManager.cpp | 96 +++++++-
.../MCJIT/MCJITMemoryManagerTest.cpp | 225 +++++++++++++++++-
3 files changed, 324 insertions(+), 15 deletions(-)
diff --git a/llvm/include/llvm/ExecutionEngine/SectionMemoryManager.h b/llvm/include/llvm/ExecutionEngine/SectionMemoryManager.h
index fa1b2355528dd0b..e0c4cede204816f 100644
--- a/llvm/include/llvm/ExecutionEngine/SectionMemoryManager.h
+++ b/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,6 +193,8 @@ 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;
@@ -187,6 +202,7 @@ class SectionMemoryManager : public RTDyldMemoryManager {
MemoryGroup RODataMem;
MemoryMapper *MMapper;
std::unique_ptr<MemoryMapper> OwnedMMapper;
+ bool ReserveAllocation;
};
} // end namespace llvm
diff --git a/llvm/lib/ExecutionEngine/SectionMemoryManager.cpp b/llvm/lib/ExecutionEngine/SectionMemoryManager.cpp
index 436888730bfb246..9400a2b4a9746e4 100644
--- a/llvm/lib/ExecutionEngine/SectionMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/SectionMemoryManager.cpp
@@ -18,6 +18,91 @@
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();
+
+ // 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(RequiredCodeSize, PageSize);
+ uint64_t RequiredSize =
+ RequiredCodeSize + RequiredRODataSize + RequiredRWDataSize;
+ printf("Allocating %llu\n", RequiredSize);
+
+ 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.
+ assert(MB.allocatedSize() == 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) {
+ 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,
@@ -91,6 +176,11 @@ uint8_t *SectionMemoryManager::allocateSection(
}
}
+ // All memory should be pre-allocated if needsToReserveAllocationSpace().
+ if (needsToReserveAllocationSpace()) {
+ return nullptr;
+ }
+
// No pre-allocated free block was large enough. Allocate a new memory region.
// Note that all sections get allocated as read-write. The permissions will
// be updated later based on memory group.
@@ -265,8 +355,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();
diff --git a/llvm/unittests/ExecutionEngine/MCJIT/MCJITMemoryManagerTest.cpp b/llvm/unittests/ExecutionEngine/MCJIT/MCJITMemoryManagerTest.cpp
index 7a756a707160adb..ae20be50d3d230b 100644
--- a/llvm/unittests/ExecutionEngine/MCJIT/MCJITMemoryManagerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/MCJIT/MCJITMemoryManagerTest.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/Support/Process.h"
#include "gtest/gtest.h"
using namespace llvm;
@@ -16,15 +17,17 @@ namespace {
TEST(MCJITMemoryManagerTest, BasicAllocations) {
std::unique_ptr<SectionMemoryManager> MemMgr(new SectionMemoryManager());
+ EXPECT_FALSE(MemMgr->needsToReserveAllocationSpace());
+
uint8_t *code1 = MemMgr->allocateCodeSection(256, 0, 1, "");
uint8_t *data1 = MemMgr->allocateDataSection(256, 0, 2, "", true);
uint8_t *code2 = MemMgr->allocateCodeSection(256, 0, 3, "");
uint8_t *data2 = MemMgr->allocateDataSection(256, 0, 4, "", false);
- EXPECT_NE((uint8_t*)nullptr, code1);
- EXPECT_NE((uint8_t*)nullptr, code2);
- EXPECT_NE((uint8_t*)nullptr, data1);
- EXPECT_NE((uint8_t*)nullptr, data2);
+ EXPECT_NE((uint8_t *)nullptr, code1);
+ EXPECT_NE((uint8_t *)nullptr, code2);
+ EXPECT_NE((uint8_t *)nullptr, data1);
+ EXPECT_NE((uint8_t *)nullptr, data2);
// Initialize the data
for (unsigned i = 0; i < 256; ++i) {
@@ -54,10 +57,10 @@ TEST(MCJITMemoryManagerTest, LargeAllocations) {
uint8_t *code2 = MemMgr->allocateCodeSection(0x100000, 0, 3, "");
uint8_t *data2 = MemMgr->allocateDataSection(0x100000, 0, 4, "", false);
- EXPECT_NE((uint8_t*)nullptr, code1);
- EXPECT_NE((uint8_t*)nullptr, code2);
- EXPECT_NE((uint8_t*)nullptr, data1);
- EXPECT_NE((uint8_t*)nullptr, data2);
+ EXPECT_NE((uint8_t *)nullptr, code1);
+ EXPECT_NE((uint8_t *)nullptr, code2);
+ EXPECT_NE((uint8_t *)nullptr, data1);
+ EXPECT_NE((uint8_t *)nullptr, data2);
// Initialize the data
for (unsigned i = 0; i < 0x100000; ++i) {
@@ -82,8 +85,8 @@ TEST(MCJITMemoryManagerTest, LargeAllocations) {
TEST(MCJITMemoryManagerTest, ManyAllocations) {
std::unique_ptr<SectionMemoryManager> MemMgr(new SectionMemoryManager());
- uint8_t* code[10000];
- uint8_t* data[10000];
+ uint8_t *code[10000];
+ uint8_t *data[10000];
for (unsigned i = 0; i < 10000; ++i) {
const bool isReadOnly = i % 2 == 0;
@@ -117,8 +120,8 @@ TEST(MCJITMemoryManagerTest, ManyAllocations) {
TEST(MCJITMemoryManagerTest, ManyVariedAllocations) {
std::unique_ptr<SectionMemoryManager> MemMgr(new SectionMemoryManager());
- uint8_t* code[10000];
- uint8_t* data[10000];
+ uint8_t *code[10000];
+ uint8_t *data[10000];
for (unsigned i = 0; i < 10000; ++i) {
uintptr_t CodeSize = i % 16 + 1;
@@ -165,5 +168,203 @@ TEST(MCJITMemoryManagerTest, ManyVariedAllocations) {
}
}
+TEST(MCJITMemoryManagerTest, PreAllocation) {
+ std::unique_ptr<SectionMemoryManager> MemMgr(
+ new SectionMemoryManager(nullptr, true));
+
+ EXPECT_TRUE(MemMgr->needsToReserveAllocationSpace());
+
+ llvm::Align Align{16};
+ MemMgr->reserveAllocationSpace(512, Align, 256, Align, 256, Align);
+
+ uint8_t *code1 = MemMgr->allocateCodeSection(256, 0, 1, "");
+ uint8_t *data1 = MemMgr->allocateDataSection(256, 0, 2, "", true);
+ uint8_t *code2 = MemMgr->allocateCodeSection(256, 0, 3, "");
+ uint8_t *data2 = MemMgr->allocateDataSection(256, 0, 4, "", false);
+
+ EXPECT_NE((uint8_t *)nullptr, code1);
+ EXPECT_NE((uint8_t *)nullptr, code2);
+ EXPECT_NE((uint8_t *)nullptr, data1);
+ EXPECT_NE((uint8_t *)nullptr, data2);
+
+ // Initialize the data
+ for (unsigned i = 0; i < 256; ++i) {
+ code1[i] = 1;
+ code2[i] = 2;
+ data1[i] = 3;
+ data2[i] = 4;
+ }
+
+ // Verify the data (this is checking for overlaps in the addresses)
+ for (unsigned i = 0; i < 256; ++i) {
+ EXPECT_EQ(1, code1[i]);
+ EXPECT_EQ(2, code2[i]);
+ EXPECT_EQ(3, data1[i]);
+ EXPECT_EQ(4, data2[i]);
+ }
+
+ std::string Error;
+ EXPECT_FALSE(MemMgr->finalizeMemory(&Error));
+}
+
+TEST(MCJITMemoryManagerTest, PreAllocationReuse) {
+ std::unique_ptr<SectionMemoryManager> MemMgr(
+ new SectionMemoryManager(nullptr, true));
+
+ EXPECT_TRUE(MemMgr->needsToReserveAllocationSpace());
+
+ // Reserve PageSize, because finalizeMemory eliminates blocks that aren't a
+ // full page size. Alignment adjustment will ensure that 2 pages are
+ // allocated for each section.
+ const size_t PageSize = sys::Process::getPageSizeEstimate();
+ EXPECT_GE(PageSize, 512u);
+ llvm::Align Align{16};
+ MemMgr->reserveAllocationSpace(PageSize, Align, PageSize, Align, PageSize,
+ Align);
+
+ uint8_t *code1 = MemMgr->allocateCodeSection(256, 0, 1, "");
+ uint8_t *data1 = MemMgr->allocateDataSection(256, 0, 2, "", true);
+ uint8_t *code2 = MemMgr->allocateCodeSection(256, 0, 3, "");
+ uint8_t *data2 = MemMgr->allocateDataSection(256, 0, 4, "", false);
+
+ uint8_t *minAddr = std::min({code1, data1, code2, data2});
+ uint8_t *maxAddr = std::max({code1, data1, code2, data2});
+
+ EXPECT_NE((uint8_t *)nullptr, code1);
+ EXPECT_NE((uint8_t *)nullptr, code2);
+ EXPECT_NE((uint8_t *)nullptr, data1);
+ EXPECT_NE((uint8_t *)nullptr, data2);
+
+ // Initialize the data
+ for (unsigned i = 0; i < 256; ++i) {
+ code1[i] = 1;
+ code2[i] = 2;
+ data1[i] = 3;
+ data2[i] = 4;
+ }
+
+ // Verify the data (this is checking for overlaps in the addresses)
+ for (unsigned i = 0; i < 256; ++i) {
+ EXPECT_EQ(1, code1[i]);
+ EXPECT_EQ(2, code2[i]);
+ EXPECT_EQ(3, data1[i]);
+ EXPECT_EQ(4, data2[i]);
+ }
+
+ std::string Error;
+ EXPECT_FALSE(MemMgr->finalizeMemory(&Error));
+
+ // Each type of data is allocated on PageSize (usually 4KB). Allocate again
+ // and guarantee we get requests in the same block.
+ MemMgr->reserveAllocationSpace(512, Align, 256, Align, 256, Align);
+
+ code1 = MemMgr->allocateCodeSection(256, 0, 5, "");
+ data1 = MemMgr->allocateDataSection(256, 0, 6, "", true);
+ code2 = MemMgr->allocateCodeSection(256, 0, 7, "");
+ data2 = MemMgr->allocateDataSection(256, 0, 8, "", false);
+
+ // Validate difference is less than 6x PageSize
+ minAddr = std::min({minAddr, code1, data1, code2, data2});
+ maxAddr = std::max({maxAddr, code1, data1, code2, data2});
+ EXPECT_LT(maxAddr - minAddr, 6 * sys::Process::getPageSizeEstimate());
+
+ EXPECT_NE((uint8_t *)nullptr, code1);
+ EXPECT_NE((uint8_t *)nullptr, code2);
+ EXPECT_NE((uint8_t *)nullptr, data1);
+ EXPECT_NE((uint8_t *)nullptr, data2);
+
+ // Initialize the data
+ for (unsigned i = 0; i < 256; ++i) {
+ code1[i] = 1;
+ code2[i] = 2;
+ data1[i] = 3;
+ data2[i] = 4;
+ }
+
+ // Verify the data (this is checking for overlaps in the addresses)
+ for (unsigned i = 0; i < 256; ++i) {
+ EXPECT_EQ(1, code1[i]);
+ EXPECT_EQ(2, code2[i]);
+ EXPECT_EQ(3, data1[i]);
+ EXPECT_EQ(4, data2[i]);
+ }
+
+ EXPECT_FALSE(MemMgr->finalizeMemory(&Error));
+}
+
+TEST(MCJITMemoryManagerTest, ManyPreAllocation) {
+ std::unique_ptr<SectionMemoryManager> MemMgr(
+ new SectionMemoryManager(nullptr, true));
+
+ uint8_t *code[10000];
+ uint8_t *data[10000];
+
+ // Total size computation needs to take into account how much memory will be
+ // used including alignment.
+ uintptr_t CodeSize = 0, RODataSize = 0, RWDataSize = 0;
+ for (unsigned i = 0; i < 10000; ++i) {
+ unsigned Align = 8 << (i % 4);
+ CodeSize += alignTo(i % 16 + 1, Align);
+ if (i % 3 == 0) {
+ RODataSize += alignTo(i % 8 + 1, Align);
+ } else {
+ RWDataSize += alignTo(i % 8 + 1, Align);
+ }
+ }
+ llvm::Align Align = llvm::Align(8);
+ MemMgr->reserveAllocationSpace(CodeSize, Align, RODataSize, Align, RWDataSize,
+ Align);
+ uint8_t *LowAddr = (uint8_t *)std::numeric_limits<uintptr_t>::max();
+ uint8_t *HighAddr = (uint8_t *)std::numeric_limits<uintptr_t>::min();
+
+ for (unsigned i = 0; i < 10000; ++i) {
+ uintptr_t CodeSize = i % 16 + 1;
+ uintptr_t DataSize = i % 8 + 1;
+
+ bool isReadOnly = i % 3 == 0;
+ unsigned Align = 8 << (i % 4);
+
+ code[i] = MemMgr->allocateCodeSection(CodeSize, Align, i, "");
+ data[i] =
+ MemMgr->allocateDataSection(DataSize, Align, i + 10000, "", isReadOnly);
+ LowAddr = std::min({LowAddr, code[i], data[i]});
+ HighAddr = std::max({HighAddr, code[i], data[i]});
+
+ EXPECT_NE((uint8_t *)nullptr, code[i]);
+ EXPECT_NE((uint8_t *)nullptr, data[i]);
+
+ for (unsigned j = 0; j < CodeSize; j++) {
+ code[i][j] = 1 + (i % 254);
+ }
+
+ for (unsigned j = 0; j < DataSize; j++) {
+ data[i][j] = 2 + (i % 254);
+ }
+
+ uintptr_t CodeAlign = Align ? (uintptr_t)code[i] % Align : 0;
+ uintptr_t DataAlign = Align ? (uintptr_t)data[i] % Align : 0;
+
+ EXPECT_EQ((uintptr_t)0, CodeAlign);
+ EXPECT_EQ((uintptr_t)0, DataAlign);
+ }
+
+ EXPECT_LT(HighAddr - LowAddr, 1024 * 1024 * 1024);
+
+ for (unsigned i = 0; i < 10000; ++i) {
+ uintptr_t CodeSize = i % 16 + 1;
+ uintptr_t DataSize = i % 8 + 1;
+
+ for (unsigned j = 0; j < CodeSize; j++) {
+ uint8_t ExpectedCode = 1 + (i % 254);
+ EXPECT_EQ(ExpectedCode, code[i][j]);
+ }
+
+ for (unsigned j = 0; j < DataSize; j++) {
+ uint8_t ExpectedData = 2 + (i % 254);
+ EXPECT_EQ(ExpectedData, data[i][j]);
+ }
+ }
+}
+
} // Namespace
More information about the llvm-commits
mailing list