diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h @@ -0,0 +1,144 @@ +//===- 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 chunk of memory reserved on executor and corresponding + /// working memory on controller process + struct Mapping { + ExecutorAddr RemoteAddr; + void *LocalAddr; + size_t Size; + }; + + /// Represents an unit of memory to finalize its contents and protections + /// The offest is realtive to a reservation represented by a Mapping + /// It ensures the offset is same on both executor and controller side for + /// in-process or shared memory use cases + struct SegInfo { + ExecutorAddrDiff Offset; + size_t ContentSize; + size_t ZeroFillSize; + unsigned Prot; + shared::AllocActions Actions; + }; + + using OnReservedFunction = unique_function)>; + + /// Reserves address space in executor process and provides working copy for + /// controller process + virtual void reserve(size_t NumBytes, OnReservedFunction OnReserved) = 0; + + using OnInitializedFunction = unique_function; + + /// Ensures working copy memory is synchronized with executor memory, sends + /// functions to be called after initilization and deinitialization and + /// applies memory protections + virtual void initialize(Mapping &M, std::vector &Segments, + OnInitializedFunction OnInitialized) = 0; + + using OnDeinitializedFunction = unique_function; + + /// Runs previously specified denitialization actions + virtual void deinitialize(Mapping &M, std::vector &Segments, + OnDeinitializedFunction OnDeInitialized) = 0; + + using OnReleasedFunction = unique_function; + + /// Unmaps address space from both process + virtual void release(std::vector &Allocs, + OnReleasedFunction OnRelease) = 0; + + Expected reserve(size_t NumBytes) { + std::promise> P; + auto F = P.get_future(); + reserve(NumBytes, [&](auto R) { P.set_value(std::move(R)); }); + return F.get(); + } + + Error initialize(Mapping &M, std::vector Segments) { + std::promise P; + auto F = P.get_future(); + initialize(M, Segments, [&](Error R) { P.set_value(std::move(R)); }); + return F.get(); + } + + Error deinitialize(Mapping &M, std::vector &Bases) { + std::promise P; + auto F = P.get_future(); + deinitialize(M, Bases, [&](Error R) { P.set_value(std::move(R)); }); + return F.get(); + } + + Error release(std::vector &Allocs) { + std::promise P; + auto F = P.get_future(); + release(Allocs, [&](Error R) { P.set_value(std::move(R)); }); + return F.get(); + } + + virtual ~MemoryMapper(); +}; + +class InProcessMemoryMapper final : public MemoryMapper { +public: + InProcessMemoryMapper() {} + + using MemoryMapper::reserve; + void reserve(size_t NumBytes, OnReservedFunction OnReserved) override; + + using MemoryMapper::initialize; + void initialize(Mapping &M, std::vector &Segments, + OnInitializedFunction OnInitialized) override; + + using MemoryMapper::deinitialize; + void deinitialize(Mapping &M, std::vector &Bases, + OnDeinitializedFunction OnDeinitialized) override; + + using MemoryMapper::release; + void release(std::vector &Allocs, + OnReleasedFunction OnRelease) override; + + ~InProcessMemoryMapper() override; + +private: + struct Allocation { + size_t Size; + SmallVector DeinitializeActions; + }; + using AllocationMap = DenseMap; + + struct Reservation { + size_t Size; + AllocationMap Allocations; + }; + using ReservationMap = DenseMap; + + std::mutex Mutex; + ReservationMap Reservations; +}; + +} // 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 --- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt +++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt @@ -23,6 +23,7 @@ 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 --- /dev/null +++ b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp @@ -0,0 +1,130 @@ +//===- 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" +#include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.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(); + } + + Mapping M; + M.LocalAddr = MB.base(); + M.RemoteAddr = ExecutorAddr::fromPtr(MB.base()); + M.Size = MB.allocatedSize(); + OnReserved(M); +} + +void InProcessMemoryMapper::initialize(Mapping &M, + std::vector &Segments, + OnInitializedFunction OnInitialized) { + std::lock_guard Lock(Mutex); + + auto &R = Reservations[M.LocalAddr]; + + for (auto &Segment : Segments) { + auto *Base = (M.RemoteAddr + Segment.Offset).toPtr(); + auto Size = Segment.ContentSize + Segment.ZeroFillSize; + + std::memset( + (M.RemoteAddr + Segment.Offset + Segment.ContentSize).toPtr(), + 0, Segment.ZeroFillSize); + + if (auto EC = + sys::Memory::protectMappedMemory({Base, Size}, Segment.Prot)) { + return OnInitialized(errorCodeToError(EC)); + } + + if (Segment.Prot & sys::Memory::MF_EXEC) + sys::Memory::InvalidateInstructionCache(Base, Size); + + auto &A = R.Allocations[Base]; + A.Size = Size; + for (auto &ActionPair : Segment.Actions) { + if (auto Err = ActionPair.Finalize.runWithSPSRetErrorMerged()) + return OnInitialized(std::move(Err)); + + A.DeinitializeActions.push_back(ActionPair.Dealloc); + } + } + + OnInitialized(Error::success()); +} + +void InProcessMemoryMapper::deinitialize( + MemoryMapper::Mapping &M, std::vector &Bases, + MemoryMapper::OnDeinitializedFunction OnDeinitialized) { + std::lock_guard Lock(Mutex); + auto &R = Reservations[M.LocalAddr]; + + Error Err = Error::success(); + for (auto Base : Bases) { + auto &A = R.Allocations[Base.toPtr()]; + while (!A.DeinitializeActions.empty()) { + auto &Action = A.DeinitializeActions.back(); + + Err = joinErrors(std::move(Err), Action.runWithSPSRetErrorMerged()); + A.DeinitializeActions.pop_back(); + } + } + + OnDeinitialized(std::move(Err)); +} + +void InProcessMemoryMapper::release(std::vector &M, + OnReleasedFunction OnReleased) { + Error Err = Error::success(); + + for (auto Mapping : M) { + auto MB = sys::MemoryBlock(Mapping.LocalAddr, Mapping.Size); + + auto EC = sys::Memory::releaseMappedMemory(MB); + if (EC) { + Err = joinErrors(std::move(Err), errorCodeToError(EC)); + } + + std::lock_guard Lock(Mutex); + Reservations.erase(Mapping.LocalAddr); + } + + OnReleased(std::move(Err)); +} + +InProcessMemoryMapper::~InProcessMemoryMapper() { + std::lock_guard Lock(Mutex); + + for (auto &R : Reservations) { + for (auto &A : R.getSecond().Allocations) { + while (!A.second.DeinitializeActions.empty()) { + auto &Action = A.second.DeinitializeActions.back(); + cantFail(Action.runWithSPSRetErrorMerged()); + A.second.DeinitializeActions.pop_back(); + } + } + } +} + +} // namespace orc + +} // namespace llvm diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt --- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt +++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt @@ -24,6 +24,7 @@ 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 --- /dev/null +++ b/llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp @@ -0,0 +1,97 @@ +//===------------------------ 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 { + +#ifdef LLVM_ON_UNIX +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(ExecutorSharedMemoryManagerTest, AllocFinalizeFree) { + int InitializeCounter = 0; + int DeinitializeCounter = 0; + { + std::unique_ptr Mapper = + std::make_unique(); + auto PageSize = sys::Process::getPageSize(); + cantFail(PageSize.takeError()); + auto AllocSize = *PageSize * 2; + + Expected Mem = Mapper->reserve(AllocSize); + EXPECT_THAT_ERROR(Mem.takeError(), Succeeded()); + + std::string HW = "Hello, world!"; + std::strcpy(static_cast(Mem->LocalAddr), HW.c_str()); + + MemoryMapper::SegInfo Segment1; + Segment1.Offset = 0; + Segment1.ContentSize = HW.size(); + Segment1.ZeroFillSize = *PageSize - Segment1.ContentSize; + Segment1.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE; + Segment1.Actions.push_back( + {cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&InitializeCounter))), + cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&DeinitializeCounter)))}); + + MemoryMapper::SegInfo Segment2; + Segment2.Offset = *PageSize; + Segment2.ContentSize = HW.size(); + Segment2.ZeroFillSize = *PageSize - Segment2.ContentSize; + Segment2.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE; + Segment2.Actions.push_back( + {cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&InitializeCounter))), + cantFail(WrapperFunctionCall::Create>( + ExecutorAddr::fromPtr(incrementWrapper), + ExecutorAddr::fromPtr(&DeinitializeCounter)))}); + + std::vector Segments{Segment1, Segment2}; + + EXPECT_EQ(InitializeCounter, 0); + EXPECT_EQ(DeinitializeCounter, 0); + + EXPECT_THAT_ERROR(Mapper->initialize(*Mem, Segments), Succeeded()); + + EXPECT_EQ(InitializeCounter, 2); + EXPECT_EQ(DeinitializeCounter, 0); + + EXPECT_EQ(HW, std::string(static_cast(Mem->LocalAddr))); + + std::vector SegAddr = {Mem->RemoteAddr}; + EXPECT_THAT_ERROR(Mapper->deinitialize(*Mem, SegAddr), Succeeded()); + + EXPECT_EQ(InitializeCounter, 2); + EXPECT_EQ(DeinitializeCounter, 1); + } + EXPECT_EQ(InitializeCounter, 2); + EXPECT_EQ(DeinitializeCounter, 2); +} +#endif + +} // namespace