[orc-rt] Align SimpleNativeMemoryMap Segment with LLVM type. (#162823)

This commit aims to align SimpleNativeMemoryMap::FinalizeRequest::Segment with
llvm::orc::tpctypes::SegFinalizeRequest. This will simplify construction of a
new LLVM JITLinkMemoryManager that's capable of using SimpleNativeMemoryMap as
a backend.
This commit is contained in:
Lang Hames
2025-10-11 10:33:28 +11:00
committed by GitHub
parent 0aef9ebc14
commit 27e2d5c46f
3 changed files with 100 additions and 87 deletions

View File

@@ -54,18 +54,10 @@ public:
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;
AllocGroup AG;
char *Address = nullptr;
size_t Size = 0;
AllocGroup G;
ContentType C = ContentType::Uninitialized;
char *data() { return reinterpret_cast<char *>(Address); }
size_t size() const { return Size; }
span<const char> Content;
};
std::vector<Segment> Segments;

View File

@@ -31,36 +31,22 @@ struct SPSSimpleNativeMemoryMapSegment;
template <>
class SPSSerializationTraits<SPSSimpleNativeMemoryMapSegment,
SimpleNativeMemoryMap::FinalizeRequest::Segment> {
using SPSType = SPSTuple<SPSExecutorAddr, uint64_t, SPSAllocGroup, uint8_t>;
using SPSType =
SPSTuple<SPSAllocGroup, SPSExecutorAddr, uint64_t, SPSSequence<char>>;
public:
static bool deserialize(SPSInputBuffer &IB,
SimpleNativeMemoryMap::FinalizeRequest::Segment &S) {
using ContentType =
SimpleNativeMemoryMap::FinalizeRequest::Segment::ContentType;
AllocGroup AG;
ExecutorAddr Address;
uint64_t Size;
AllocGroup G;
uint8_t C;
if (!SPSType::AsArgList::deserialize(IB, Address, Size, G, C))
span<const char> Content;
if (!SPSType::AsArgList::deserialize(IB, AG, Address, Size, Content))
return false;
if (Size >= std::numeric_limits<size_t>::max())
if (Size > std::numeric_limits<size_t>::max())
return false;
S.Address = Address.toPtr<void *>();
S.Size = Size;
S.G = G;
S.C = static_cast<ContentType>(C);
switch (S.C) {
case ContentType::Uninitialized:
return true;
case ContentType::ZeroFill:
memset(reinterpret_cast<char *>(S.Address), 0, S.Size);
return true;
case ContentType::Regular:
// Read content directly into target address.
return IB.read(reinterpret_cast<char *>(S.Address), S.Size);
}
S = {AG, Address.toPtr<char *>(), static_cast<size_t>(Size), Content};
return true;
}
};
@@ -138,10 +124,31 @@ void SimpleNativeMemoryMap::finalize(OnFinalizeCompleteFn &&OnComplete,
// TODO: Record finalize segments for release.
// std::vector<std::pair<void*, size_t>> FinalizeSegments;
// Check segment validity before proceeding.
for (auto &S : FR.Segments) {
if (auto Err = hostOSMemoryProtect(S.Address, S.Size, S.G.getMemProt()))
if (S.Content.size() > S.Size) {
return OnComplete(make_error<StringError>(
(std::ostringstream()
<< "For segment [" << (void *)S.Address << ".."
<< (void *)(S.Address + S.Size) << "), "
<< " content size (" << std::hex << S.Content.size()
<< ") exceeds segment size (" << S.Size << ")")
.str()));
}
// Copy any requested content.
if (!S.Content.empty())
memcpy(S.Address, S.Content.data(), S.Content.size());
// Zero-fill the rest of the section.
if (size_t ZeroFillSize = S.Size - S.Content.size())
memset(S.Address + S.Content.size(), 0, ZeroFillSize);
if (auto Err = hostOSMemoryProtect(S.Address, S.Size, S.AG.getMemProt()))
return OnComplete(std::move(Err));
switch (S.G.getMemLifetime()) {
switch (S.AG.getMemLifetime()) {
case MemLifetime::Standard:
if (!Base || S.Address < Base)
Base = S.Address;

View File

@@ -31,49 +31,32 @@ struct SPSSimpleNativeMemoryMapSegment;
struct TestSNMMSegment
: public SimpleNativeMemoryMap::FinalizeRequest::Segment {
enum TestSNMMSegmentContent { Uninitialized, ZeroFill };
TestSNMMSegment(void *Address, AllocGroup G, std::string Content)
TestSNMMSegment(AllocGroup AG, char *Address, size_t Size,
std::vector<char> C = {})
: SimpleNativeMemoryMap::FinalizeRequest::Segment(
Address, Content.size(), G, ContentType::Regular),
Content(std::move(Content)) {}
{AG, Address, Size, {}}),
OwnedContent(std::move(C)) {
this->Content = {OwnedContent.data(), OwnedContent.size()};
}
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;
std::vector<char> OwnedContent;
};
template <>
class SPSSerializationTraits<SPSSimpleNativeMemoryMapSegment, TestSNMMSegment> {
using SPSType = SPSTuple<SPSExecutorAddr, uint64_t, SPSAllocGroup, uint8_t>;
using SPSType =
SPSTuple<SPSAllocGroup, SPSExecutorAddr, uint64_t, SPSSequence<char>>;
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<uint64_t>(S.Size), S.G,
static_cast<uint8_t>(S.C)) +
(S.C == ContentType::Regular ? S.Size : 0);
return SPSType::AsArgList::size(S.AG, ExecutorAddr::fromPtr(S.Address),
static_cast<uint64_t>(S.Size), S.Content);
}
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<uint64_t>(S.Size), S.G,
static_cast<uint8_t>(S.C)))
return false;
if (S.C == ContentType::Regular)
return OB.write(S.Content.data(), S.Content.size());
return true;
return SPSType::AsArgList::serialize(
OB, S.AG, ExecutorAddr::fromPtr(S.Address),
static_cast<uint64_t>(S.Size), S.Content);
}
};
@@ -207,30 +190,61 @@ TEST(SimpleNativeMemoryMap, FullPipelineForOneRWSegment) {
std::future<Expected<Expected<void *>>> FinalizeKey;
TestSNMMFinalizeRequest FR;
void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(Addr) + 64 * 1024);
uint64_t SentinelValue = 0;
char *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
reinterpret_cast<char *>(Addr) + 64 * 1024;
uint64_t SentinelValue1 = 0; // Read from pre-filled content
uint64_t SentinelValue2 = 0; // Written in finalize, read back during dealloc.
uint64_t SentinelValue3 = 42; // Read from zero-filled region.
FR.Segments.push_back({FinalizeBase, 64 * 1024,
MemProt::Read | MemProt::Write,
TestSNMMSegment::ZeroFill});
// Build initial content vector.
std::vector<char> Content;
Content.resize(sizeof(uint64_t) * 2);
memcpy(Content.data(), &SentinelValue3, sizeof(uint64_t));
memcpy(Content.data() + sizeof(uint64_t), &SentinelValue1, sizeof(uint64_t));
FR.Segments.push_back({MemProt::Read | MemProt::Write, FinalizeBase,
64 * 1024, std::move(Content)});
// Read initial content into Sentinel 1.
FR.AAPs.push_back({
*MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue1),
ExecutorAddr::fromPtr(FinalizeBase)),
{} // No dealloc action.
});
// Write value in finalize action, then read back into Sentinel 2.
FR.AAPs.push_back(
{*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),
write_value_sps_allocaction,
ExecutorAddr::fromPtr(FinalizeBase) + sizeof(uint64_t),
uint64_t(42)),
*MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue),
ExecutorAddr::fromPtr(FinalizeBase))});
read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue2),
ExecutorAddr::fromPtr(FinalizeBase) + sizeof(uint64_t))});
// Read first 64 bits of the zero-fill region.
FR.AAPs.push_back({
*MakeAllocAction<SPSExecutorAddr, SPSExecutorAddr>::from(
read_value_sps_allocaction, ExecutorAddr::fromPtr(&SentinelValue3),
ExecutorAddr::fromPtr(FinalizeBase) + sizeof(uint64_t) * 2),
{} // No dealloc action.
});
snmm_finalize(waitFor(FinalizeKey), SNMM.get(), std::move(FR));
void *FinalizeKeyAddr = cantFail(cantFail(FinalizeKey.get()));
EXPECT_EQ(SentinelValue, 0U);
EXPECT_EQ(SentinelValue1, 42U);
EXPECT_EQ(SentinelValue2, 0U);
EXPECT_EQ(SentinelValue3, 0U);
std::future<Expected<Error>> DeallocResult;
snmm_deallocate(waitFor(DeallocResult), SNMM.get(), FinalizeKeyAddr);
cantFail(cantFail(DeallocResult.get()));
EXPECT_EQ(SentinelValue, 42);
EXPECT_EQ(SentinelValue1, 42U);
EXPECT_EQ(SentinelValue2, 42U);
EXPECT_EQ(SentinelValue3, 0U);
std::future<Expected<Error>> ReleaseResult;
snmm_release(waitFor(ReleaseResult), SNMM.get(), Addr);
@@ -248,13 +262,13 @@ TEST(SimpleNativeMemoryMap, ReserveFinalizeShutdown) {
std::future<Expected<Expected<void *>>> FinalizeKey;
TestSNMMFinalizeRequest FR;
void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(Addr) + 64 * 1024);
char *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
reinterpret_cast<char *>(Addr) + 64 * 1024;
uint64_t SentinelValue = 0;
FR.Segments.push_back({FinalizeBase, 64 * 1024,
MemProt::Read | MemProt::Write,
TestSNMMSegment::ZeroFill});
FR.Segments.push_back(
{MemProt::Read | MemProt::Write, FinalizeBase, 64 * 1024});
FR.AAPs.push_back(
{*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),
@@ -285,13 +299,13 @@ TEST(SimpleNativeMemoryMap, ReserveFinalizeDetachShutdown) {
std::future<Expected<Expected<void *>>> FinalizeKey;
TestSNMMFinalizeRequest FR;
void *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(Addr) + 64 * 1024);
char *FinalizeBase = // Finalize addr at non-zero (64kb) offset from base.
reinterpret_cast<char *>(Addr) + 64 * 1024;
uint64_t SentinelValue = 0;
FR.Segments.push_back({FinalizeBase, 64 * 1024,
MemProt::Read | MemProt::Write,
TestSNMMSegment::ZeroFill});
FR.Segments.push_back(
{MemProt::Read | MemProt::Write, FinalizeBase, 64 * 1024});
FR.AAPs.push_back(
{*MakeAllocAction<SPSExecutorAddr, uint64_t>::from(
write_value_sps_allocaction, ExecutorAddr::fromPtr(FinalizeBase),