diff --git a/orc-rt/include/CMakeLists.txt b/orc-rt/include/CMakeLists.txt index 89a92bb895a50..7f1438eea4409 100644 --- a/orc-rt/include/CMakeLists.txt +++ b/orc-rt/include/CMakeLists.txt @@ -15,6 +15,7 @@ set(ORC_RT_HEADERS orc-rt/ResourceManager.h orc-rt/RTTI.h orc-rt/ScopeExit.h + orc-rt/SimpleNativeMemoryMap.h orc-rt/SimplePackedSerialization.h orc-rt/SPSAllocAction.h orc-rt/SPSMemoryFlags.h diff --git a/orc-rt/include/orc-rt/SPSWrapperFunction.h b/orc-rt/include/orc-rt/SPSWrapperFunction.h index c48694c6d2bfc..824e009306e2a 100644 --- a/orc-rt/include/orc-rt/SPSWrapperFunction.h +++ b/orc-rt/include/orc-rt/SPSWrapperFunction.h @@ -14,9 +14,12 @@ #ifndef ORC_RT_SPSWRAPPERFUNCTION_H #define ORC_RT_SPSWRAPPERFUNCTION_H +#include "orc-rt/Compiler.h" #include "orc-rt/SimplePackedSerialization.h" #include "orc-rt/WrapperFunction.h" +#define ORC_RT_SPS_INTERFACE ORC_RT_INTERFACE + namespace orc_rt { namespace detail { diff --git a/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h b/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h new file mode 100644 index 0000000000000..77d3339e2ddbd --- /dev/null +++ b/orc-rt/include/orc-rt/SimpleNativeMemoryMap.h @@ -0,0 +1,123 @@ +//===- SimpleNativeMemoryMap.h -- Mem via standard host OS APIs -*- 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 +// +//===----------------------------------------------------------------------===// +// +// SimpleNativeMemoryMap and related APIs. +// +//===----------------------------------------------------------------------===// + +#ifndef ORC_RT_SIMPLENATIVEMEMORYMAP_H +#define ORC_RT_SIMPLENATIVEMEMORYMAP_H + +#include "orc-rt/AllocAction.h" +#include "orc-rt/Error.h" +#include "orc-rt/MemoryFlags.h" +#include "orc-rt/ResourceManager.h" +#include "orc-rt/SPSWrapperFunction.h" +#include "orc-rt/move_only_function.h" + +#include +#include +#include +#include + +namespace orc_rt { + +/// JIT'd memory management backend. +/// +/// Intances can: +/// 1. Reserve address space. +/// 2. Finalize memory regions within reserved memory (copying content, +/// applying permissions, running finalize actions, and recording +/// deallocate actions). +/// 3. Deallocate memory regions within reserved memory (running +/// deallocate actions and making memory available for future +/// finalize calls (if the system permits this). +/// 4. Release address space, deallocating any not-yet-deallocated finalized +/// regions, and returning the address space to the system for reuse (if +/// the system permits). +class SimpleNativeMemoryMap : public ResourceManager { +public: + /// Reserves a slab of contiguous address space for allocation. + /// + /// Returns the base address of the allocated memory. + using OnReserveCompleteFn = move_only_function)>; + void reserve(OnReserveCompleteFn &&OnComplete, size_t Size); + + /// Release a slab of contiguous address space back to the system. + using OnReleaseCompleteFn = move_only_function; + void release(OnReleaseCompleteFn &&OnComplete, void *Addr); + + struct FinalizeRequest { + struct Segment { + enum class ContentType : uint8_t { Uninitialized, ZeroFill, Regular }; + + Segment() = default; + Segment(void *Address, size_t Size, AllocGroup G, ContentType C) + : Address(Address), Size(Size), G(G), C(C) {} + + void *Address = nullptr; + size_t Size = 0; + AllocGroup G; + ContentType C = ContentType::Uninitialized; + char *data() { return reinterpret_cast(Address); } + size_t size() const { return Size; } + }; + + std::vector Segments; + std::vector AAPs; + }; + + /// Writes content into the requested ranges, applies permissions, and + /// performs allocation actions. + using OnFinalizeCompleteFn = move_only_function)>; + void finalize(OnFinalizeCompleteFn &&OnComplete, FinalizeRequest FR); + + /// Runs deallocation actions and resets memory permissions for the requested + /// memory. + using OnDeallocateCompleteFn = move_only_function; + void deallocate(OnDeallocateCompleteFn &&OnComplete, void *Base); + + void detach(ResourceManager::OnCompleteFn OnComplete) override; + void shutdown(ResourceManager::OnCompleteFn OnComplete) override; + +private: + struct SlabInfo { + SlabInfo(size_t Size) : Size(Size) {} + size_t Size; + std::unordered_map> DeallocActions; + }; + + void shutdownNext(OnCompleteFn OnComplete, std::vector Bases); + Error makeBadSlabError(void *Base, const char *Op); + SlabInfo *findSlabInfoFor(void *Base); + Error recordDeallocActions(void *Base, + std::vector DeallocActions); + + std::mutex M; + std::map Slabs; +}; + +} // namespace orc_rt + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes); + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_release_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes); + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes); + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, orc_rt_WrapperFunctionBuffer ArgBytes); + +#endif // ORC_RT_SIMPLENATIVEMEMORYMAP_H diff --git a/orc-rt/lib/executor/CMakeLists.txt b/orc-rt/lib/executor/CMakeLists.txt index 08103251fa199..2c709e2b2ca5d 100644 --- a/orc-rt/lib/executor/CMakeLists.txt +++ b/orc-rt/lib/executor/CMakeLists.txt @@ -3,6 +3,7 @@ set(files AllocAction.cpp ResourceManager.cpp RTTI.cpp + SimpleNativeMemoryMap.cpp ) add_library(orc-rt-executor STATIC ${files}) diff --git a/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp b/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp new file mode 100644 index 0000000000000..603ef8bef9848 --- /dev/null +++ b/orc-rt/lib/executor/SimpleNativeMemoryMap.cpp @@ -0,0 +1,330 @@ +//===- SimpleNativeMemoryMap.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 +// +//===----------------------------------------------------------------------===// +// +// SimpleNativeMemoryMap and related APIs. +// +// TODO: We don't reset / uncommit pages on deallocate, or on failure during +// finalize. We should do that to reduce memory pressure. +// +//===----------------------------------------------------------------------===// + +#include "orc-rt/SimpleNativeMemoryMap.h" +#include "orc-rt/SPSAllocAction.h" +#include "orc-rt/SPSMemoryFlags.h" +#include + +#if defined(__APPLE__) || defined(__linux__) +#include "Unix/NativeMemoryAPIs.inc" +#else +#error "Target OS memory APIs unsupported" +#endif + +namespace orc_rt { + +struct SPSSimpleNativeMemoryMapSegment; + +template <> +class SPSSerializationTraits { + using SPSType = SPSTuple; + +public: + static bool deserialize(SPSInputBuffer &IB, + SimpleNativeMemoryMap::FinalizeRequest::Segment &S) { + using ContentType = + SimpleNativeMemoryMap::FinalizeRequest::Segment::ContentType; + + ExecutorAddr Address; + uint64_t Size; + AllocGroup G; + uint8_t C; + if (!SPSType::AsArgList::deserialize(IB, Address, Size, G, C)) + return false; + if (Size >= std::numeric_limits::max()) + return false; + S.Address = Address.toPtr(); + S.Size = Size; + S.G = G; + S.C = static_cast(C); + switch (S.C) { + case ContentType::Uninitialized: + return true; + case ContentType::ZeroFill: + memset(reinterpret_cast(S.Address), 0, S.Size); + return true; + case ContentType::Regular: + // Read content directly into target address. + return IB.read(reinterpret_cast(S.Address), S.Size); + } + } +}; + +struct SPSSimpleNativeMemoryMapFinalizeRequest; + +template <> +class SPSSerializationTraits { + using SPSType = SPSTuple, + SPSSequence>; + +public: + static bool deserialize(SPSInputBuffer &IB, + SimpleNativeMemoryMap::FinalizeRequest &FR) { + return SPSType::AsArgList::deserialize(IB, FR.Segments, FR.AAPs); + } +}; + +void SimpleNativeMemoryMap::reserve(OnReserveCompleteFn &&OnComplete, + size_t Size) { + // FIXME: Get page size from session object. + if (Size % (64 * 1024)) { + return OnComplete(make_error( + (std::ostringstream() + << "SimpleNativeMemoryMap error: reserved size " << std::hex << Size + << " is not a page-size multiple") + .str())); + } + + auto Addr = hostOSMemoryReserve(Size); + if (!Addr) + return OnComplete(Addr.takeError()); + + { + std::scoped_lock Lock(M); + assert(!Slabs.count(*Addr) && + "hostOSMemoryReserve returned duplicate addresses"); + Slabs.emplace(std::make_pair(*Addr, SlabInfo(Size))); + } + + OnComplete(*Addr); +} + +void SimpleNativeMemoryMap::release(OnReleaseCompleteFn &&OnComplete, + void *Addr) { + std::optional SI; + { + std::scoped_lock Lock(M); + auto I = Slabs.find(Addr); + if (I != Slabs.end()) { + SI = std::move(I->second); + Slabs.erase(I); + } + } + + if (!SI) { + std::ostringstream ErrMsg; + ErrMsg << "SimpleNativeMemoryMap error: release called on unrecognized " + "address " + << Addr; + return OnComplete(make_error(ErrMsg.str())); + } + + for (auto &[Addr, DAAs] : SI->DeallocActions) + runDeallocActions(std::move(DAAs)); + + OnComplete(hostOSMemoryRelease(Addr, SI->Size)); +} + +void SimpleNativeMemoryMap::finalize(OnFinalizeCompleteFn &&OnComplete, + FinalizeRequest FR) { + + void *Base = nullptr; + + // TODO: Record finalize segments for release. + // std::vector> FinalizeSegments; + + for (auto &S : FR.Segments) { + if (auto Err = hostOSMemoryProtect(S.Address, S.Size, S.G.getMemProt())) + return OnComplete(std::move(Err)); + switch (S.G.getMemLifetime()) { + case MemLifetime::Standard: + if (!Base || S.Address < Base) + Base = S.Address; + break; + case MemLifetime::Finalize: + // TODO: Record finalize segment for release. + // FinalizeSegments.push_back({S.Address, S.Size}); + break; + } + } + + if (!Base) + return OnComplete(make_error( + "SimpleNativeMemoryMap finalize error: finalization requires at least " + "one standard-lifetime segment")); + + auto DeallocActions = runFinalizeActions(std::move(FR.AAPs)); + if (!DeallocActions) + return OnComplete(DeallocActions.takeError()); + + if (auto Err = recordDeallocActions(Base, std::move(*DeallocActions))) { + runDeallocActions(std::move(*DeallocActions)); + return OnComplete(std::move(Err)); + } + + OnComplete(Base); +} + +void SimpleNativeMemoryMap::deallocate(OnDeallocateCompleteFn &&OnComplete, + void *Base) { + std::vector DAAs; + + { + std::unique_lock Lock(M); + auto *SI = findSlabInfoFor(Base); + if (!SI) { + Lock.unlock(); + return OnComplete(makeBadSlabError(Base, "finalize")); + } + + auto I = SI->DeallocActions.find(Base); + if (I == SI->DeallocActions.end()) { + Lock.unlock(); + std::ostringstream ErrMsg; + ErrMsg << "SimpleNativeMemoryMap deallocate error: no deallocate actions " + "registered for segment base address " + << Base; + return OnComplete(make_error(ErrMsg.str())); + } + + DAAs = std::move(I->second); + SI->DeallocActions.erase(I); + } + + runDeallocActions(std::move(DAAs)); + OnComplete(Error::success()); +} + +void SimpleNativeMemoryMap::detach(ResourceManager::OnCompleteFn OnComplete) { + // Detach is a noop for now: we just retain all actions to run at shutdown + // time. + OnComplete(Error::success()); +} + +void SimpleNativeMemoryMap::shutdown(ResourceManager::OnCompleteFn OnComplete) { + // TODO: Establish a clear order to run dealloca actions across slabs, + // object boundaries. + + // Collect slab base addresses for removal. + std::vector Bases; + { + std::scoped_lock Lock(M); + for (auto &[Base, _] : Slabs) + Bases.push_back(Base); + } + + shutdownNext(std::move(OnComplete), std::move(Bases)); +} + +void SimpleNativeMemoryMap::shutdownNext( + ResourceManager::OnCompleteFn OnComplete, std::vector Bases) { + if (Bases.empty()) + return OnComplete(Error::success()); + + auto *Base = Bases.back(); + Bases.pop_back(); + + release( + [this, Bases = std::move(Bases), + OnComplete = std::move(OnComplete)](Error Err) mutable { + if (Err) { + // TODO: Log release error? + consumeError(std::move(Err)); + } + shutdownNext(std::move(OnComplete), std::move(Bases)); + }, + Base); +} + +Error SimpleNativeMemoryMap::makeBadSlabError(void *Base, const char *Op) { + std::ostringstream ErrMsg; + ErrMsg << "SimpleNativeMemoryMap " << Op << " error: segment base address " + << Base << " does not fall within an allocated slab"; + return make_error(ErrMsg.str()); +} + +SimpleNativeMemoryMap::SlabInfo * +SimpleNativeMemoryMap::findSlabInfoFor(void *Base) { + // NOTE: We assume that the caller is holding a lock for M. + auto I = Slabs.upper_bound(Base); + if (I == Slabs.begin()) + return nullptr; + + --I; + if (reinterpret_cast(I->first) + I->second.Size <= + reinterpret_cast(Base)) + return nullptr; + + return &I->second; +} + +Error SimpleNativeMemoryMap::recordDeallocActions( + void *Base, std::vector DeallocActions) { + + std::unique_lock Lock(M); + auto *SI = findSlabInfoFor(Base); + if (!SI) { + Lock.unlock(); + return makeBadSlabError(Base, "deallocate"); + } + + auto I = SI->DeallocActions.find(Base); + if (I != SI->DeallocActions.end()) { + Lock.unlock(); + std::ostringstream ErrMsg; + ErrMsg << "SimpleNativeMemoryMap finalize error: segment base address " + "reused in subsequent finalize call"; + return make_error(ErrMsg.str()); + } + + SI->DeallocActions[Base] = std::move(DeallocActions); + return Error::success(); +} + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, + orc_rt_WrapperFunctionBuffer ArgBytes) { + using Sig = SPSExpected(SPSExecutorAddr, SPSSize); + SPSWrapperFunction::handle( + Session, CallCtx, Return, ArgBytes, + WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::reserve)); +} + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_release_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, + orc_rt_WrapperFunctionBuffer ArgBytes) { + using Sig = SPSError(SPSExecutorAddr, SPSExecutorAddr); + SPSWrapperFunction::handle( + Session, CallCtx, Return, ArgBytes, + WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::release)); +} + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, + orc_rt_WrapperFunctionBuffer ArgBytes) { + using Sig = SPSExpected( + SPSExecutorAddr, SPSSimpleNativeMemoryMapFinalizeRequest); + SPSWrapperFunction::handle( + Session, CallCtx, Return, ArgBytes, + WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::finalize)); +} + +ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper( + orc_rt_SessionRef Session, void *CallCtx, + orc_rt_WrapperFunctionReturn Return, + orc_rt_WrapperFunctionBuffer ArgBytes) { + using Sig = SPSError(SPSExecutorAddr, SPSExecutorAddr); + SPSWrapperFunction::handle(Session, CallCtx, Return, ArgBytes, + WrapperFunction::handleWithAsyncMethod( + &SimpleNativeMemoryMap::deallocate)); +} + +} // namespace orc_rt diff --git a/orc-rt/lib/executor/Unix/NativeMemoryAPIs.inc b/orc-rt/lib/executor/Unix/NativeMemoryAPIs.inc new file mode 100644 index 0000000000000..41c6632f40dd1 --- /dev/null +++ b/orc-rt/lib/executor/Unix/NativeMemoryAPIs.inc @@ -0,0 +1,103 @@ +//===- NativeMemoryAPIs.inc -------------------------------------*- 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 +// +//===----------------------------------------------------------------------===// +// +// Generic wrappers for unix-style memory APIs (mmap, mprotect, etc.). +// +//===----------------------------------------------------------------------===// + +#include "orc-rt/Error.h" +#include "orc-rt/MemoryFlags.h" + +#include +#include +#include +#include + +namespace { + +int toNativeProtFlags(orc_rt::MemProt MP) { + int Prot = PROT_NONE; + if ((MP & orc_rt::MemProt::Read) != orc_rt::MemProt::None) + Prot |= PROT_READ; + if ((MP & orc_rt::MemProt::Write) != orc_rt::MemProt::None) + Prot |= PROT_WRITE; + if ((MP & orc_rt::MemProt::Exec) != orc_rt::MemProt::None) + Prot |= PROT_EXEC; + return Prot; +} + +#if defined(__APPLE__) +extern "C" void sys_icache_invalidate(const void *Addr, size_t Size); +#else +extern "C" void __clear_cache(void *Start, void *End); +#endif + +static void invalidateInstructionCache(void *Addr, size_t Size) { +#if defined(__APPLE__) + sys_icache_invalidate(Addr, Size); +#else + __clear_cache(Addr, reinterpret_cast(Addr) + Size); +#endif +} + +orc_rt::Expected hostOSMemoryReserve(size_t Size) { + if (Size == 0) + return nullptr; + + int FD = 0; + int MapFlags = MAP_PRIVATE; + +#if defined(MAP_ANON) + // If MAP_ANON is available then use it. + FD = -1; + MapFlags |= MAP_ANON; +#else // !defined(MAP_ANON) + // Fall back to /dev/zero for strict POSIX. + fd = open("/dev/zero", O_RDWR); + if (fd == -1) { + auto ErrNum = errno; + return make_error( + std::string("Could not open /dev/zero for memory reserve: ") + + strerror(ErrNum)); + } +#endif + + void *Addr = mmap(nullptr, Size, PROT_READ | PROT_WRITE, MapFlags, FD, 0); + if (Addr == MAP_FAILED) { + auto ErrNum = errno; + return orc_rt::make_error( + std::string("mmap for memory reserve failed: ") + strerror(ErrNum)); + } + + return Addr; +} + +orc_rt::Error hostOSMemoryRelease(void *Base, size_t Size) { + if (munmap(Base, Size) != 0) { + auto ErrNum = errno; + return orc_rt::make_error( + std::string("munmap for memory release failed: ") + strerror(ErrNum)); + } + return orc_rt::Error::success(); +} + +orc_rt::Error hostOSMemoryProtect(void *Base, size_t Size, orc_rt::MemProt MP) { + if (mprotect(Base, Size, toNativeProtFlags(MP)) != 0) { + auto ErrNum = errno; + return orc_rt::make_error( + std::string("mprotect for memory finalize failed: ") + + strerror(ErrNum)); + } + + if ((MP & orc_rt::MemProt::Exec) != orc_rt::MemProt::None) + invalidateInstructionCache(Base, Size); + + return orc_rt::Error::success(); +} + +} // anonymous namespace diff --git a/orc-rt/unittests/CMakeLists.txt b/orc-rt/unittests/CMakeLists.txt index 4d3da682d4222..639eaecf16366 100644 --- a/orc-rt/unittests/CMakeLists.txt +++ b/orc-rt/unittests/CMakeLists.txt @@ -23,6 +23,7 @@ add_orc_rt_unittest(CoreTests MemoryFlagsTest.cpp RTTITest.cpp ScopeExitTest.cpp + SimpleNativeMemoryMapTest.cpp SimplePackedSerializationTest.cpp SPSAllocActionTest.cpp SPSMemoryFlagsTest.cpp diff --git a/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp new file mode 100644 index 0000000000000..ebd9bc38905fa --- /dev/null +++ b/orc-rt/unittests/SimpleNativeMemoryMapTest.cpp @@ -0,0 +1,318 @@ +//===-- SPSNativeMemoryMapTest.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 +// +//===----------------------------------------------------------------------===// +// +// Test SPS serialization for MemoryFlags APIs. +// +//===----------------------------------------------------------------------===// + +#include "orc-rt/SimpleNativeMemoryMap.h" +#include "orc-rt/SPSAllocAction.h" +#include "orc-rt/SPSMemoryFlags.h" + +#include "AllocActionTestUtils.h" +#include "DirectCaller.h" +#include "gtest/gtest.h" + +#include + +using namespace orc_rt; + +namespace orc_rt { + +struct SPSSimpleNativeMemoryMapSegment; + +/// A SimpleNativeMemoryMap::FinalizeRequest::Segment plus segment content (if +/// segment content type is regular). +struct TestSNMMSegment + : public SimpleNativeMemoryMap::FinalizeRequest::Segment { + + enum TestSNMMSegmentContent { Uninitialized, ZeroFill }; + + TestSNMMSegment(void *Address, AllocGroup G, std::string Content) + : SimpleNativeMemoryMap::FinalizeRequest::Segment( + Address, Content.size(), G, ContentType::Regular), + Content(std::move(Content)) {} + + TestSNMMSegment(void *Address, size_t Size, AllocGroup G, + TestSNMMSegmentContent Content) + : SimpleNativeMemoryMap::FinalizeRequest::Segment( + Address, Size, G, + Content == ZeroFill ? ContentType::ZeroFill + : ContentType::Uninitialized) {} + + std::string Content; +}; + +template <> +class SPSSerializationTraits { + using SPSType = SPSTuple; + +public: + static size_t size(const TestSNMMSegment &S) { + using ContentType = + SimpleNativeMemoryMap::FinalizeRequest::Segment::ContentType; + assert((S.C != ContentType::Regular || S.Size == S.Content.size())); + return SPSType::AsArgList::size(ExecutorAddr::fromPtr(S.Address), + static_cast(S.Size), S.G, + static_cast(S.C)) + + (S.C == ContentType::Regular ? S.Size : 0); + } + + static bool serialize(SPSOutputBuffer &OB, const TestSNMMSegment &S) { + using ContentType = + SimpleNativeMemoryMap::FinalizeRequest::Segment::ContentType; + assert((S.C != ContentType::Regular || S.Size == S.Content.size())); + if (!SPSType::AsArgList::serialize(OB, ExecutorAddr::fromPtr(S.Address), + static_cast(S.Size), S.G, + static_cast(S.C))) + return false; + if (S.C == ContentType::Regular) + return OB.write(S.Content.data(), S.Content.size()); + return true; + } +}; + +struct SPSSimpleNativeMemoryMapFinalizeRequest; + +struct TestSNMMFinalizeRequest { + std::vector Segments; + std::vector AAPs; +}; + +template <> +class SPSSerializationTraits { + using SPSType = SPSTuple, + SPSSequence>; + +public: + static size_t size(const TestSNMMFinalizeRequest &FR) { + return SPSType::AsArgList::size(FR.Segments, FR.AAPs); + } + static bool serialize(SPSOutputBuffer &OB, + const TestSNMMFinalizeRequest &FR) { + return SPSType::AsArgList::serialize(OB, FR.Segments, FR.AAPs); + } +}; + +} // namespace orc_rt + +template move_only_function waitFor(std::future &F) { + std::promise P; + F = P.get_future(); + return [P = std::move(P)](T Val) mutable { P.set_value(std::move(Val)); }; +} + +TEST(SimpleNativeMemoryMapTest, CreateAndDestroy) { + // Test that we can create and destroy a SimpleNativeMemoryMap instance as + // expected. + auto SNMM = std::make_unique(); +} + +template +static void snmm_reserve(OnCompleteFn &&OnComplete, + SimpleNativeMemoryMap *Instance, size_t Size) { + using SPSSig = SPSExpected(SPSExecutorAddr, SPSSize); + SPSWrapperFunction::call( + DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper), + std::forward(OnComplete), Instance, Size); +} + +template +static void snmm_finalize(OnCompleteFn &&OnComplete, + SimpleNativeMemoryMap *Instance, + TestSNMMFinalizeRequest FR) { + using SPSSig = SPSExpected( + SPSExecutorAddr, SPSSimpleNativeMemoryMapFinalizeRequest); + SPSWrapperFunction::call( + DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_finalize_sps_wrapper), + std::forward(OnComplete), Instance, std::move(FR)); +} + +template +static void snmm_deallocate(OnCompleteFn &&OnComplete, + SimpleNativeMemoryMap *Instance, void *Base) { + using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr); + SPSWrapperFunction::call( + DirectCaller(nullptr, + orc_rt_SimpleNativeMemoryMap_deallocate_sps_wrapper), + std::forward(OnComplete), Instance, Base); +} + +template +static void snmm_release(OnCompleteFn &&OnComplete, + SimpleNativeMemoryMap *Instance, void *Addr) { + using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr); + SPSWrapperFunction::call( + DirectCaller(nullptr, orc_rt_SimpleNativeMemoryMap_release_sps_wrapper), + std::forward(OnComplete), Instance, Addr); +} + +TEST(SimpleNativeMemoryMapTest, ReserveAndRelease) { + // Test that we can reserve and release a slab of address space as expected, + // without finalizing any memory within it. + auto SNMM = std::make_unique(); + std::future>> ReserveAddr; + snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024); + auto Addr = cantFail(cantFail(ReserveAddr.get())); + + std::future> ReleaseResult; + snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr); + cantFail(cantFail(ReleaseResult.get())); +} + +// Write the given value to the address pointed to by P. +static orc_rt_WrapperFunctionBuffer +write_value_sps_allocaction(const char *ArgData, size_t ArgSize) { + return SPSAllocActionFunction::handle( + ArgData, ArgSize, + [](ExecutorAddr P, uint64_t Val) { + *P.toPtr() = Val; + return WrapperFunctionBuffer(); + }) + .release(); +} + +// Read the uint64_t value at Src and write it to Dst. +// Increments int via pointer. +static orc_rt_WrapperFunctionBuffer +read_value_sps_allocaction(const char *ArgData, size_t ArgSize) { + return SPSAllocActionFunction::handle( + ArgData, ArgSize, + [](ExecutorAddr Dst, ExecutorAddr Src) { + *Dst.toPtr() = *Src.toPtr(); + return WrapperFunctionBuffer(); + }) + .release(); +} + +TEST(SimpleNativeMemoryMap, FullPipelineForOneRWSegment) { + // Test that we can: + // 1. reserve some address space. + // 2. finalize a range within it as read/write, and that finalize actions + // are applied as expected. + // 3. deallocate the finalized range, with deallocation actions applied as + // expected. + // 4. release the address range. + + auto SNMM = std::make_unique(); + std::future>> ReserveAddr; + snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024); + void *Addr = cantFail(cantFail(ReserveAddr.get())); + + std::future>> FinalizeKey; + TestSNMMFinalizeRequest FR; + void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base. + reinterpret_cast(reinterpret_cast(Addr) + 64 * 1024); + uint64_t SentinelValue = 0; + + FR.Segments.push_back({FinalizeBase, 64 * 1024, + MemProt::Read | MemProt::Write, + TestSNMMSegment::ZeroFill}); + FR.AAPs.push_back( + {*MakeAllocAction::from( + write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase), + uint64_t(42)), + *MakeAllocAction::from( + read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue), + ExecutorAddr::fromPtr(FinalizeBase))}); + snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR)); + void *FinalizeKeyAddr = cantFail(cantFail(FinalizeKey.get())); + + EXPECT_EQ(SentinelValue, 0U); + + std::future> DeallocResult; + snmm_deallocate(waitFor(DeallocResult), SNMM.get(), FinalizeKeyAddr); + cantFail(cantFail(DeallocResult.get())); + + EXPECT_EQ(SentinelValue, 42); + + std::future> ReleaseResult; + snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr); + cantFail(cantFail(ReleaseResult.get())); +} + +TEST(SimpleNativeMemoryMap, ReserveFinalizeShutdown) { + // Test that memory is deallocated in the case where we reserve and finalize + // some memory, then just shut down the memory manager. + + auto SNMM = std::make_unique(); + std::future>> ReserveAddr; + snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024); + void *Addr = cantFail(cantFail(ReserveAddr.get())); + + std::future>> FinalizeKey; + TestSNMMFinalizeRequest FR; + void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base. + reinterpret_cast(reinterpret_cast(Addr) + 64 * 1024); + uint64_t SentinelValue = 0; + + FR.Segments.push_back({FinalizeBase, 64 * 1024, + MemProt::Read | MemProt::Write, + TestSNMMSegment::ZeroFill}); + FR.AAPs.push_back( + {*MakeAllocAction::from( + write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase), + uint64_t(42)), + *MakeAllocAction::from( + read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue), + ExecutorAddr::fromPtr(FinalizeBase))}); + snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR)); + cantFail(cantFail(FinalizeKey.get())); + + EXPECT_EQ(SentinelValue, 0U); + + std::future ShutdownResult; + SNMM->shutdown(waitFor(ShutdownResult)); + cantFail(ShutdownResult.get()); + + EXPECT_EQ(SentinelValue, 42); +} + +TEST(SimpleNativeMemoryMap, ReserveFinalizeDetachShutdown) { + // Test that memory is deallocated in the case where we reserve and finalize + // some memory, then just shut down the memory manager. + + auto SNMM = std::make_unique(); + std::future>> ReserveAddr; + snmm_reserve(waitFor(ReserveAddr), SNMM.get(), 1024 * 1024 * 1024); + void *Addr = cantFail(cantFail(ReserveAddr.get())); + + std::future>> FinalizeKey; + TestSNMMFinalizeRequest FR; + void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base. + reinterpret_cast(reinterpret_cast(Addr) + 64 * 1024); + uint64_t SentinelValue = 0; + + FR.Segments.push_back({FinalizeBase, 64 * 1024, + MemProt::Read | MemProt::Write, + TestSNMMSegment::ZeroFill}); + FR.AAPs.push_back( + {*MakeAllocAction::from( + write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase), + uint64_t(42)), + *MakeAllocAction::from( + read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue), + ExecutorAddr::fromPtr(FinalizeBase))}); + snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR)); + cantFail(cantFail(FinalizeKey.get())); + + EXPECT_EQ(SentinelValue, 0U); + + std::future DetachResult; + SNMM->detach(waitFor(DetachResult)); + cantFail(DetachResult.get()); + + EXPECT_EQ(SentinelValue, 0); + + std::future ShutdownResult; + SNMM->shutdown(waitFor(ShutdownResult)); + cantFail(ShutdownResult.get()); + + EXPECT_EQ(SentinelValue, 42); +}