From b8e29f8fd5f168e453cc2c312a26f772448625e6 Mon Sep 17 00:00:00 2001 From: Lang Hames Date: Thu, 9 Oct 2025 12:33:12 +1100 Subject: [PATCH] [orc-rt] Add SimpleNativeMemoryMap. SimpleNativeMemoryMap is a memory allocation backend for use with ORC. It can... 1. Reserve slabs of address space. 2. Finalize regions of memory within a reserved slab: copying content into requested addresses, applying memory protections, running finalization actions, and storing deallocation actions to be run by deallocate. 3. Deallocate finalized memory regions: running deallocate actions and, if possible, making memory in these regions available for use by future finalization operations. (Some systems prohibit reuse of executable memory. On these systems deallocated memory is no longer usable within the process). 4. Release reserved slabs. This runs deallocate for any not-yet-deallocated finalized regions, and then (if possible) returns the address space to system. (On systems that prohibit reuse of executable memory parts of the released address space may be permanently unusable by the process). SimpleNativeMemoryMap is intended primarily for use by llvm::orc::JITLinkMemoryManager implementations to allocate JIT'd code and data. --- orc-rt/include/CMakeLists.txt | 1 + orc-rt/include/orc-rt/SPSWrapperFunction.h | 3 + orc-rt/include/orc-rt/SimpleNativeMemoryMap.h | 123 +++++++ orc-rt/lib/executor/CMakeLists.txt | 1 + orc-rt/lib/executor/SimpleNativeMemoryMap.cpp | 330 ++++++++++++++++++ orc-rt/lib/executor/Unix/NativeMemoryAPIs.inc | 103 ++++++ orc-rt/unittests/CMakeLists.txt | 1 + .../unittests/SimpleNativeMemoryMapTest.cpp | 318 +++++++++++++++++ 8 files changed, 880 insertions(+) create mode 100644 orc-rt/include/orc-rt/SimpleNativeMemoryMap.h create mode 100644 orc-rt/lib/executor/SimpleNativeMemoryMap.cpp create mode 100644 orc-rt/lib/executor/Unix/NativeMemoryAPIs.inc create mode 100644 orc-rt/unittests/SimpleNativeMemoryMapTest.cpp 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); +}