Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 3 additions & 11 deletions orc-rt/include/orc-rt/SimpleNativeMemoryMap.h
Original file line number Diff line number Diff line change
Expand Up @@ -54,18 +54,10 @@ class SimpleNativeMemoryMap : public ResourceManager {

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;
Expand Down
55 changes: 31 additions & 24 deletions orc-rt/lib/executor/SimpleNativeMemoryMap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
};

Expand Down Expand Up @@ -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;
Expand Down
120 changes: 67 additions & 53 deletions orc-rt/unittests/SimpleNativeMemoryMapTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -31,49 +31,32 @@ struct SPSSimpleNativeMemoryMapSegment;
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)
TestSNMMSegment(AllocGroup AG, char *Address, size_t Size,
std::vector<char> C = {})
: SimpleNativeMemoryMap::FinalizeRequest::Segment(
Address, Size, G,
Content == ZeroFill ? ContentType::ZeroFill
: ContentType::Uninitialized) {}
{AG, Address, Size, {}}),
OwnedContent(std::move(C)) {
this->Content = {OwnedContent.data(), OwnedContent.size()};
}

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);
}
};

Expand Down Expand Up @@ -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;

FR.Segments.push_back({FinalizeBase, 64 * 1024,
MemProt::Read | MemProt::Write,
TestSNMMSegment::ZeroFill});
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.

// 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);
Expand All @@ -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),
Expand Down Expand Up @@ -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),
Expand Down
Loading