diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h new file mode 100644 index 00000000000000..d023bfbdb5b6f4 --- /dev/null +++ b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h @@ -0,0 +1,115 @@ +//===- MemoryMapper.h - Cross-process memory mapper -------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Cross-process (and in-process) memory mapping and transfer +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H +#define LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H + +#include "llvm/ExecutionEngine/Orc/Core.h" + +#include + +namespace llvm { +namespace orc { + +/// Manages mapping, content transfer and protections for JIT memory +class MemoryMapper { +public: + /// Represents a single allocation containing multiple segments and + /// initialization and deinitialization actions + struct AllocInfo { + struct SegInfo { + ExecutorAddrDiff Offset; + const char *WorkingMem; + size_t ContentSize; + size_t ZeroFillSize; + unsigned Prot; + }; + + ExecutorAddr MappingBase; + std::vector Segments; + shared::AllocActions Actions; + }; + + using OnReservedFunction = unique_function)>; + + /// Reserves address space in executor process + virtual void reserve(size_t NumBytes, OnReservedFunction OnReserved) = 0; + + /// Provides working memory + virtual char *prepare(ExecutorAddr Addr, size_t ContentSize) = 0; + + using OnInitializedFunction = unique_function)>; + + /// Ensures executor memory is synchronized with working copy memory, sends + /// functions to be called after initilization and before deinitialization and + /// applies memory protections + /// Returns a unique address identifying the allocation. This address should + /// be passed to deinitialize to run deallocation actions (and reset + /// permissions where possible). + virtual void initialize(AllocInfo &AI, + OnInitializedFunction OnInitialized) = 0; + + using OnDeinitializedFunction = unique_function; + + /// Runs previously specified deinitialization actions + /// Executor addresses returned by initialize should be passed + virtual void deinitialize(ArrayRef Allocations, + OnDeinitializedFunction OnDeInitialized) = 0; + + using OnReleasedFunction = unique_function; + + /// Release address space acquired through reserve() + virtual void release(ArrayRef Reservations, + OnReleasedFunction OnRelease) = 0; + + virtual ~MemoryMapper(); +}; + +class InProcessMemoryMapper final : public MemoryMapper { +public: + InProcessMemoryMapper() {} + + void reserve(size_t NumBytes, OnReservedFunction OnReserved) override; + + void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override; + + char *prepare(ExecutorAddr Addr, size_t ContentSize) override; + + void deinitialize(ArrayRef Allocations, + OnDeinitializedFunction OnDeInitialized) override; + + void release(ArrayRef Reservations, + OnReleasedFunction OnRelease) override; + + ~InProcessMemoryMapper() override; + +private: + struct Allocation { + std::vector DeinitializationActions; + }; + using AllocationMap = DenseMap; + + struct Reservation { + size_t Size; + std::vector Allocations; + }; + using ReservationMap = DenseMap; + + std::mutex Mutex; + ReservationMap Reservations; + AllocationMap Allocations; +}; + +} // namespace orc +} // end namespace llvm + +#endif // LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt index cf8e0c85840c95..1c9b9e8a1b31ec 100644 --- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt +++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt @@ -23,6 +23,7 @@ add_llvm_component_library(LLVMOrcJIT LookupAndRecordAddrs.cpp LLJIT.cpp MachOPlatform.cpp + MemoryMapper.cpp ELFNixPlatform.cpp Mangling.cpp ObjectLinkingLayer.cpp diff --git a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp new file mode 100644 index 00000000000000..f00cac3e2c3e98 --- /dev/null +++ b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp @@ -0,0 +1,152 @@ +//===- MemoryMapper.cpp - Cross-process memory mapper ------------*- C++ -*-==// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "llvm/ExecutionEngine/Orc/MemoryMapper.h" + +namespace llvm { +namespace orc { + +MemoryMapper::~MemoryMapper() {} + +void InProcessMemoryMapper::reserve(size_t NumBytes, + OnReservedFunction OnReserved) { + std::error_code EC; + auto MB = sys::Memory::allocateMappedMemory( + NumBytes, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC); + + if (EC) + return OnReserved(errorCodeToError(EC)); + + { + std::lock_guard Lock(Mutex); + Reservations[MB.base()].Size = MB.allocatedSize(); + } + + OnReserved( + ExecutorAddrRange(ExecutorAddr::fromPtr(MB.base()), MB.allocatedSize())); +} + +char *InProcessMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) { + return Addr.toPtr(); +} + +void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI, + OnInitializedFunction OnInitialized) { + ExecutorAddr MinAddr(~0ULL); + + for (auto &Segment : AI.Segments) { + auto Base = AI.MappingBase + Segment.Offset; + auto Size = Segment.ContentSize + Segment.ZeroFillSize; + + if (Base < MinAddr) + MinAddr = Base; + + std::memset((Base + Segment.ContentSize).toPtr(), 0, + Segment.ZeroFillSize); + + if (auto EC = sys::Memory::protectMappedMemory({Base.toPtr(), Size}, + Segment.Prot)) { + return OnInitialized(errorCodeToError(EC)); + } + if (Segment.Prot & sys::Memory::MF_EXEC) + sys::Memory::InvalidateInstructionCache(Base.toPtr(), Size); + } + + auto DeinitializeActions = shared::runFinalizeActions(AI.Actions); + if (!DeinitializeActions) + return OnInitialized(DeinitializeActions.takeError()); + + { + std::lock_guard Lock(Mutex); + Allocations[MinAddr].DeinitializationActions = + std::move(*DeinitializeActions); + Reservations[AI.MappingBase.toPtr()].Allocations.push_back(MinAddr); + } + + OnInitialized(MinAddr); +} + +void InProcessMemoryMapper::deinitialize( + ArrayRef Bases, + MemoryMapper::OnDeinitializedFunction OnDeinitialized) { + Error AllErr = Error::success(); + + { + std::lock_guard Lock(Mutex); + + for (auto Base : Bases) { + + if (Error Err = shared::runDeallocActions( + Allocations[Base].DeinitializationActions)) { + AllErr = joinErrors(std::move(AllErr), std::move(Err)); + } + + Allocations.erase(Base); + } + } + + OnDeinitialized(std::move(AllErr)); +} + +void InProcessMemoryMapper::release(ArrayRef Bases, + OnReleasedFunction OnReleased) { + Error Err = Error::success(); + + for (auto Base : Bases) { + std::vector AllocAddrs; + size_t Size; + { + std::lock_guard Lock(Mutex); + auto &R = Reservations[Base.toPtr()]; + Size = R.Size; + AllocAddrs.swap(R.Allocations); + } + + // deinitialize sub allocations + std::promise P; + auto F = P.get_future(); + deinitialize(AllocAddrs, [&](Error Err) { P.set_value(std::move(Err)); }); + if (Error E = F.get()) { + Err = joinErrors(std::move(Err), std::move(E)); + } + + // free the memory + auto MB = sys::MemoryBlock(Base.toPtr(), Size); + + auto EC = sys::Memory::releaseMappedMemory(MB); + if (EC) { + Err = joinErrors(std::move(Err), errorCodeToError(EC)); + } + + std::lock_guard Lock(Mutex); + Reservations.erase(Base.toPtr()); + } + + OnReleased(std::move(Err)); +} + +InProcessMemoryMapper::~InProcessMemoryMapper() { + std::vector ReservationAddrs; + { + std::lock_guard Lock(Mutex); + + ReservationAddrs.reserve(Reservations.size()); + for (const auto &R : Reservations) { + ReservationAddrs.push_back(ExecutorAddr::fromPtr(R.getFirst())); + } + } + + std::promise P; + auto F = P.get_future(); + release(ReservationAddrs, [&](Error Err) { P.set_value(std::move(Err)); }); + cantFail(F.get()); +} + +} // namespace orc + +} // namespace llvm diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt index 9d45957266b8a8..dd4028e5302ca9 100644 --- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt +++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt @@ -24,6 +24,7 @@ add_llvm_unittest(OrcJITTests JITTargetMachineBuilderTest.cpp LazyCallThroughAndReexportsTest.cpp LookupAndRecordAddrsTest.cpp + MemoryMapperTest.cpp ObjectLinkingLayerTest.cpp OrcCAPITest.cpp OrcTestCommon.cpp diff --git a/llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp b/llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp new file mode 100644 index 00000000000000..f352715429c022 --- /dev/null +++ b/llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp @@ -0,0 +1,203 @@ +//===------------------------ MemoryMapperTest.cpp ------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "llvm/ExecutionEngine/Orc/MemoryMapper.h" +#include "llvm/Support/Process.h" +#include "llvm/Testing/Support/Error.h" +#include "gtest/gtest.h" + +using namespace llvm; +using namespace llvm::orc; +using namespace llvm::orc::shared; + +namespace { + +Expected reserve(MemoryMapper &M, size_t NumBytes) { + std::promise> P; + auto F = P.get_future(); + M.reserve(NumBytes, [&](auto R) { P.set_value(std::move(R)); }); + return F.get(); +} + +Expected initialize(MemoryMapper &M, + MemoryMapper::AllocInfo &AI) { + std::promise> P; + auto F = P.get_future(); + M.initialize(AI, [&](auto R) { P.set_value(std::move(R)); }); + return F.get(); +} + +Error deinitialize(MemoryMapper &M, + const std::vector &Allocations) { + std::promise P; + auto F = P.get_future(); + M.deinitialize(Allocations, [&](auto R) { P.set_value(std::move(R)); }); + return F.get(); +} + +Error release(MemoryMapper &M, const std::vector &Reservations) { + std::promise P; + auto F = P.get_future(); + M.release(Reservations, [&](auto R) { P.set_value(std::move(R)); }); + return F.get(); +} + +// A basic function to be used as both initializer/deinitializer +orc::shared::CWrapperFunctionResult incrementWrapper(const char *ArgData, + size_t ArgSize) { + return WrapperFunction::handle( + ArgData, ArgSize, + [](ExecutorAddr A) -> Error { + *A.toPtr() += 1; + return Error::success(); + }) + .release(); +} + +TEST(MemoryMapperTest, InitializeDeinitialize) { + // These counters are used to track how many times the initializer and + // deinitializer functions are called + int InitializeCounter = 0; + int DeinitializeCounter = 0; + { + std::unique_ptr Mapper = + std::make_unique(); + + // We will do two separate allocations + auto PageSize = cantFail(sys::Process::getPageSize()); + auto TotalSize = PageSize * 2; + + // Reserve address space + auto Mem1 = reserve(*Mapper, TotalSize); + EXPECT_THAT_ERROR(Mem1.takeError(), Succeeded()); + + // Test string for memory transfer + std::string HW = "Hello, world!"; + + { + // Provide working memory + char *WA1 = Mapper->prepare(Mem1->Start, HW.size() + 1); + std::strcpy(static_cast(WA1), HW.c_str()); + } + + // A structure to be passed to initialize + MemoryMapper::AllocInfo Alloc1; + { + MemoryMapper::AllocInfo::SegInfo Seg1; + Seg1.Offset = 0; + Seg1.ContentSize = HW.size(); + Seg1.ZeroFillSize = PageSize - Seg1.ContentSize; + Seg1.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE; + + Alloc1.MappingBase = Mem1->Start; + Alloc1.Segments.push_back(Seg1); + Alloc1.Actions.push_back( + {cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&InitializeCounter))), + cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&DeinitializeCounter)))}); + } + + { + char *WA2 = Mapper->prepare(Mem1->Start + PageSize, HW.size() + 1); + std::strcpy(static_cast(WA2), HW.c_str()); + } + + MemoryMapper::AllocInfo Alloc2; + { + MemoryMapper::AllocInfo::SegInfo Seg2; + Seg2.Offset = PageSize; + Seg2.ContentSize = HW.size(); + Seg2.ZeroFillSize = PageSize - Seg2.ContentSize; + Seg2.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE; + + Alloc2.MappingBase = Mem1->Start; + Alloc2.Segments.push_back(Seg2); + Alloc2.Actions.push_back( + {cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&InitializeCounter))), + cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&DeinitializeCounter)))}); + } + + EXPECT_EQ(InitializeCounter, 0); + EXPECT_EQ(DeinitializeCounter, 0); + + // Set memory protections and run initializers + auto Init1 = initialize(*Mapper, Alloc1); + EXPECT_THAT_ERROR(Init1.takeError(), Succeeded()); + EXPECT_EQ(HW, std::string(static_cast(Init1->toPtr()))); + + EXPECT_EQ(InitializeCounter, 1); + EXPECT_EQ(DeinitializeCounter, 0); + + auto Init2 = initialize(*Mapper, Alloc2); + EXPECT_THAT_ERROR(Init2.takeError(), Succeeded()); + EXPECT_EQ(HW, std::string(static_cast(Init2->toPtr()))); + + EXPECT_EQ(InitializeCounter, 2); + EXPECT_EQ(DeinitializeCounter, 0); + + // Explicit deinitialization of first allocation + std::vector DeinitAddr = {*Init1}; + EXPECT_THAT_ERROR(deinitialize(*Mapper, DeinitAddr), Succeeded()); + + EXPECT_EQ(InitializeCounter, 2); + EXPECT_EQ(DeinitializeCounter, 1); + + // Test explicit release + { + auto Mem2 = reserve(*Mapper, PageSize); + EXPECT_THAT_ERROR(Mem2.takeError(), Succeeded()); + + char *WA = Mapper->prepare(Mem2->Start, HW.size() + 1); + std::strcpy(static_cast(WA), HW.c_str()); + + MemoryMapper::AllocInfo Alloc3; + { + MemoryMapper::AllocInfo::SegInfo Seg3; + Seg3.Offset = 0; + Seg3.ContentSize = HW.size(); + Seg3.ZeroFillSize = PageSize - Seg3.ContentSize; + Seg3.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE; + + Alloc3.MappingBase = Mem2->Start; + Alloc3.Segments.push_back(Seg3); + Alloc3.Actions.push_back( + {cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&InitializeCounter))), + cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&DeinitializeCounter)))}); + } + auto Init3 = initialize(*Mapper, Alloc3); + EXPECT_THAT_ERROR(Init3.takeError(), Succeeded()); + EXPECT_EQ(HW, std::string(static_cast(Init3->toPtr()))); + + EXPECT_EQ(InitializeCounter, 3); + EXPECT_EQ(DeinitializeCounter, 1); + + std::vector ReleaseAddrs = {Mem2->Start}; + EXPECT_THAT_ERROR(release(*Mapper, ReleaseAddrs), Succeeded()); + + EXPECT_EQ(InitializeCounter, 3); + EXPECT_EQ(DeinitializeCounter, 2); + } + } + + // Implicit deinitialization by the destructor + EXPECT_EQ(InitializeCounter, 3); + EXPECT_EQ(DeinitializeCounter, 3); +} + +} // namespace