diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h index 0b4cda119cadd..80dac5aee80bc 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h @@ -102,6 +102,7 @@ class InProcessMemoryMapper final : public MemoryMapper { private: struct Allocation { + size_t Size; std::vector DeinitializationActions; }; using AllocationMap = DenseMap; diff --git a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp index b9328f0a16edf..bd773a10d9174 100644 --- a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp +++ b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp @@ -147,17 +147,35 @@ void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G, void MapperJITLinkMemoryManager::deallocate( std::vector Allocs, OnDeallocatedFunction OnDeallocated) { - std::lock_guard Lock(Mutex); - + std::vector Bases; + Bases.reserve(Allocs.size()); for (auto &FA : Allocs) { ExecutorAddr Addr = FA.getAddress(); - ExecutorAddrDiff Size = UsedMemory[Addr]; - UsedMemory.erase(Addr); - - AvailableMemory.push_back({Addr, Addr + Size}); - FA.release(); + Bases.push_back(Addr); } - OnDeallocated(Error::success()); + + Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs), + OnDeallocated = std::move(OnDeallocated)]( + llvm::Error Err) mutable { + if (Err) + OnDeallocated(std::move(Err)); + + { + std::lock_guard Lock(Mutex); + + for (auto &FA : Allocs) { + ExecutorAddr Addr = FA.getAddress(); + ExecutorAddrDiff Size = UsedMemory[Addr]; + + UsedMemory.erase(Addr); + AvailableMemory.push_back({Addr, Addr + Size}); + + FA.release(); + } + } + + OnDeallocated(Error::success()); + }); } } // end namespace orc diff --git a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp index ee92e5191b507..1e8ad27342d72 100644 --- a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp +++ b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp @@ -60,6 +60,7 @@ char *InProcessMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) { void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI, OnInitializedFunction OnInitialized) { ExecutorAddr MinAddr(~0ULL); + ExecutorAddr MaxAddr(0); for (auto &Segment : AI.Segments) { auto Base = AI.MappingBase + Segment.Offset; @@ -68,6 +69,9 @@ void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI, if (Base < MinAddr) MinAddr = Base; + if (Base + Size > MaxAddr) + MaxAddr = Base + Size; + std::memset((Base + Segment.ContentSize).toPtr(), 0, Segment.ZeroFillSize); @@ -85,6 +89,9 @@ void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI, { std::lock_guard Lock(Mutex); + + // This is the maximum range whose permission have been possibly modified + Allocations[MinAddr].Size = MaxAddr - MinAddr; Allocations[MinAddr].DeinitializationActions = std::move(*DeinitializeActions); Reservations[AI.MappingBase.toPtr()].Allocations.push_back(MinAddr); @@ -108,6 +115,14 @@ void InProcessMemoryMapper::deinitialize( AllErr = joinErrors(std::move(AllErr), std::move(Err)); } + // Reset protections to read/write so the area can be reused + if (auto EC = sys::Memory::protectMappedMemory( + {Base.toPtr(), Allocations[Base].Size}, + sys::Memory::ProtectionFlags::MF_READ | + sys::Memory::ProtectionFlags::MF_WRITE)) { + AllErr = joinErrors(std::move(AllErr), errorCodeToError(EC)); + } + Allocations.erase(Base); } } diff --git a/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp index 74d8387701359..ca1b13c759f63 100644 --- a/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp +++ b/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp @@ -118,11 +118,17 @@ TEST(MapperJITLinkMemoryManagerTest, InProcess) { StringRef TargetHello2(TargetMem2, Hello.size()); EXPECT_EQ(Hello, TargetHello2); + EXPECT_EQ(Counter->DeinitCount, 0); + auto Err2 = MemMgr->deallocate(std::move(*FA1)); EXPECT_THAT_ERROR(std::move(Err2), Succeeded()); + EXPECT_EQ(Counter->DeinitCount, 1); + auto Err3 = MemMgr->deallocate(std::move(*FA2)); EXPECT_THAT_ERROR(std::move(Err3), Succeeded()); + + EXPECT_EQ(Counter->DeinitCount, 2); } } // namespace