From aa635cd136232f7ee7b2c6994fbdec36b733eb32 Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Wed, 10 Jun 2026 16:25:05 -0700 Subject: [PATCH] Add diff-based replay reset Restore prologue memory between replay iterations by scattering only the byte ranges recorded in the diff epilogue snapshot, instead of recopying the full prologue state. Diff snapshots are parsed once into a reusable DiffPlan; changed ranges are coalesced, packed into device buffers, and applied by a device scatter kernel. Selectable via --reset-mode / MNEME_REPLAY_RESET_MODE with tuning knobs for range coalescing and per-task chunk size. --- docs/usage/cli.md | 5 +- docs/usage/replay.md | 24 ++ include/mneme/MnemeDeviceKernels.hpp | 21 ++ include/mneme/MnemeReplay.hpp | 310 ++++++++++++++++++++++++- include/mneme/MnemeSnapshot.hpp | 201 ++++++++++++---- python/mneme/async_executor.py | 6 + python/mneme/commands.py | 8 + python/mneme/device.py | 9 + python/mneme/replay_executor.py | 58 ++++- python/tests/test_device.py | 3 +- python/tests/test_replay_executor.py | 49 +++- src/MnemeDeviceCode.cpp | 35 +++ src/python/device.cpp | 88 ++++++- tests/unit_tests/CMakeLists.txt | 2 +- tests/unit_tests/ReadWriteSnapshot.cpp | 112 ++++++++- 15 files changed, 847 insertions(+), 84 deletions(-) create mode 100644 include/mneme/MnemeDeviceKernels.hpp diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 1d66e005..9f4e0763 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -196,9 +196,10 @@ When omitted, all values default to those recorded during execution. ### Execution Options -| Option | Description | -| --------------------- | --------------------------------------------------------------------- | +| Option | Description | +| -------------- | --------------------------------------------------------------------- | | `--iterations`, `-it` | Number of iterations to execute the kernel for statistical evaluation | +| `--reset-mode` | Memory reset mode between iterations: `bytes` or `diff` | ### Examples diff --git a/docs/usage/replay.md b/docs/usage/replay.md index 6f6d4850..7a50a03e 100644 --- a/docs/usage/replay.md +++ b/docs/usage/replay.md @@ -21,3 +21,27 @@ optimization pipelines or launch configurations. For a detailed explanation of the annotation API, metadata fields, threshold semantics, and a complete example, see **[Usage → Verification](verification.md)**. + +## Reset mode + +Replay restores the recorded prologue before each measured kernel run. +When the recording has a diff epilogue, Mneme can avoid copying the +entire prologue by restoring only the prologue byte ranges that differ +in the epilogue: + +```bash +mneme replay ... --reset-mode diff "default" +``` + +Supported modes are: + +| Mode | Behavior | +| ---- | -------- | +| `bytes` | Always use the original full prologue reset. | +| `diff` | Require diff reset and fail if the epilogue is not a valid diff snapshot. Diff reset restores ranges with one device scatter kernel per warm reset. | + +When no reset mode is supplied, Mneme selects `diff` for diff epilogue +snapshots and `bytes` otherwise. Diff reset uses raw ranges by default; +tune optional coalescing with +`MNEME_REPLAY_DIFF_SCATTER_MAX_GAP_BYTES` and chunk size with +`MNEME_REPLAY_DIFF_SCATTER_TASK_BYTES`. diff --git a/include/mneme/MnemeDeviceKernels.hpp b/include/mneme/MnemeDeviceKernels.hpp new file mode 100644 index 00000000..521a5f33 --- /dev/null +++ b/include/mneme/MnemeDeviceKernels.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include +#include + +#include "mneme/DeviceTraits.hpp" + +namespace mneme { + +struct DiffResetScatterTask { + uint8_t *Dst = nullptr; + const uint8_t *Src = nullptr; + size_t Size = 0; +}; + +template +typename DeviceTraits::DeviceError_t launchDiffResetScatterKernel( + const DiffResetScatterTask *Tasks, size_t NumTasks, + typename DeviceTraits::DeviceStream_t Stream); + +} // namespace mneme diff --git a/include/mneme/MnemeReplay.hpp b/include/mneme/MnemeReplay.hpp index d62d0885..d07762ee 100644 --- a/include/mneme/MnemeReplay.hpp +++ b/include/mneme/MnemeReplay.hpp @@ -1,4 +1,8 @@ +#include +#include #include +#include +#include #include #include #include @@ -6,9 +10,12 @@ #include #include #include +#include #include +#include #include "mneme/DeviceTraits.hpp" +#include "mneme/MnemeDeviceKernels.hpp" #include "mneme/MnemeLogger.hpp" #include "mneme/MnemeMemory.hpp" #include "mneme/MnemePageManager.hpp" @@ -17,6 +24,33 @@ namespace mneme { +enum class ReplayResetMode { Bytes, Diff }; + +inline ReplayResetMode parseReplayResetMode(std::string Mode) { + std::transform(Mode.begin(), Mode.end(), Mode.begin(), + [](unsigned char C) { return std::tolower(C); }); + if (Mode == "bytes" || Mode == "full") + return ReplayResetMode::Bytes; + if (Mode == "diff") + return ReplayResetMode::Diff; + LOG_FATAL("Unknown Mneme replay reset mode: " + Mode + + ". Expected 'bytes' or 'diff'."); + return ReplayResetMode::Bytes; +} + +inline size_t getEnvSizeOrDefault(const char *Name, size_t Default) { + const char *Value = std::getenv(Name); + if (!Value || !*Value) + return Default; + try { + return std::stoull(Value); + } catch (...) { + LOG_WARN("Invalid value '{}' for {}, using default {}", Value, Name, + Default); + return Default; + } +} + template class PrologueState; template class EpilogueState; @@ -156,15 +190,24 @@ template class ReplayMemState { } }; -// Replay state for the recorded kernel input. +// Replay state for the recorded kernel input. Besides the full byte reset +// inherited from the base, a prologue can carry a diff reset plan: the byte +// ranges the kernel mutated (taken from a diff epilogue snapshot), coalesced, +// packed into device buffers, and restored in-place by a device scatter +// kernel between replay iterations. template class PrologueState : public ReplayMemState { public: + using MnemeDeviceRT = DeviceTraits; + using DeviceStream_t = typename MnemeDeviceRT::DeviceStream_t; + PrologueState(const std::string &KernelName, const std::string &SnapshotFile) : ReplayMemState( MnemeSnapshot::readBytesSnapshot(KernelName, SnapshotFile)) {} + ~PrologueState() override { releaseScatterPlan(); } + void load() override { for (auto &[DevAddr, MemBlob] : this->DeviceMemoryState) { auto EC = DeviceTraits::DeviceErrorCheck( @@ -186,8 +229,257 @@ class PrologueState : public ReplayMemState { PrologueState *asPrologue() override { return this; } + // Builds (or clears) the diff reset plan for this prologue from the + // epilogue's diff snapshot. With ReplayResetMode::Bytes the plan is simply + // released and reset() falls back to the full byte copy. + void prepareResetPlan(const EpilogueState &Epilogue, + ReplayResetMode Mode) { + buildDiffResetPlan(Epilogue, Mode); + } + + using ReplayMemState::reset; + + void reset(ReplayResetMode Mode, DeviceStream_t Stream = 0) { + if (Mode == ReplayResetMode::Diff) { + resetFromDiffPlan(Stream); + return; + } + reset(); + } + protected: bool isPrologue() const override { return true; } + +private: + struct ResetSpan { + void *Dst = nullptr; + const uint8_t *Src = nullptr; + size_t Size = 0; + }; + + std::vector DiffResetSpans; + std::unique_ptr HostScatterData; + DiffResetScatterTask *DeviceScatterTasks = nullptr; + uint8_t *DeviceScatterData = nullptr; + size_t ScatterTaskCount = 0; + size_t ScatterDataSize = 0; + + void releaseScatterPlan() { + if (DeviceScatterTasks) { + auto EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceFree(DeviceScatterTasks)); + if (EC) + LOG_FATAL("Could not release Mneme diff scatter task buffer EC: " + + EC.value()); + DeviceScatterTasks = nullptr; + } + if (DeviceScatterData) { + auto EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceFree(DeviceScatterData)); + if (EC) + LOG_FATAL("Could not release Mneme diff scatter data buffer EC: " + + EC.value()); + DeviceScatterData = nullptr; + } + HostScatterData.reset(); + ScatterTaskCount = 0; + ScatterDataSize = 0; + } + + static std::vector::DiffRange> + coalesceRanges( + std::vector::DiffRange> Ranges, + size_t MaxGap) { + if (Ranges.empty()) + return Ranges; + std::sort(Ranges.begin(), Ranges.end(), + [](auto &LHS, auto &RHS) { return LHS.Offset < RHS.Offset; }); + + std::vector::DiffRange> Coalesced; + auto Current = Ranges.front(); + for (size_t I = 1; I < Ranges.size(); ++I) { + auto &Next = Ranges[I]; + size_t CurrentEnd = Current.Offset + Current.Size; + if (Next.Offset <= CurrentEnd + MaxGap) { + size_t NextEnd = Next.Offset + Next.Size; + Current.Size = std::max(CurrentEnd, NextEnd) - Current.Offset; + } else { + Coalesced.push_back(Current); + Current = Next; + } + } + Coalesced.push_back(Current); + return Coalesced; + } + + void addScatterSpans( + const std::vector::DiffRange> &Ranges, + const uint8_t *HostBase, uint8_t *DeviceBase, size_t MaxGap) { + auto Spans = coalesceRanges(Ranges, MaxGap); + for (auto &Range : Spans) { + DiffResetSpans.push_back(ResetSpan{DeviceBase + Range.Offset, + HostBase + Range.Offset, Range.Size}); + } + } + + void prepareDeviceScatterPlan(size_t TaskBytes) { + releaseScatterPlan(); + if (DiffResetSpans.empty()) + return; + if (TaskBytes == 0) + TaskBytes = 4096; + + std::vector HostTasks; + HostTasks.reserve(DiffResetSpans.size()); + for (auto &Span : DiffResetSpans) + ScatterDataSize += Span.Size; + HostScatterData = std::make_unique(ScatterDataSize); + + size_t PackedOffset = 0; + for (auto &Span : DiffResetSpans) { + size_t Offset = 0; + while (Offset < Span.Size) { + size_t ChunkSize = std::min(TaskBytes, Span.Size - Offset); + std::memcpy(HostScatterData.get() + PackedOffset, Span.Src + Offset, + ChunkSize); + HostTasks.push_back(DiffResetScatterTask{ + static_cast(Span.Dst) + Offset, + reinterpret_cast(PackedOffset), ChunkSize}); + PackedOffset += ChunkSize; + Offset += ChunkSize; + } + } + + if (PackedOffset != ScatterDataSize) + LOG_FATAL("Internal Mneme diff scatter packing size mismatch"); + + auto EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceMalloc( + reinterpret_cast(&DeviceScatterData), ScatterDataSize)); + if (EC) + LOG_FATAL("Could not allocate Mneme diff scatter data buffer EC: " + + EC.value()); + + EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceCopy( + DeviceScatterData, HostScatterData.get(), ScatterDataSize, + MnemeDeviceRT::MemcpyHostToDeviceKind())); + if (EC) + LOG_FATAL("Could not copy Mneme diff scatter data buffer EC: " + + EC.value()); + HostScatterData.reset(); + + for (auto &Task : HostTasks) + Task.Src = DeviceScatterData + reinterpret_cast(Task.Src); + + ScatterTaskCount = HostTasks.size(); + EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceMalloc( + reinterpret_cast(&DeviceScatterTasks), + ScatterTaskCount * sizeof(DiffResetScatterTask))); + if (EC) + LOG_FATAL("Could not allocate Mneme diff scatter task buffer EC: " + + EC.value()); + + EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceCopy( + DeviceScatterTasks, HostTasks.data(), + ScatterTaskCount * sizeof(DiffResetScatterTask), + MnemeDeviceRT::MemcpyHostToDeviceKind())); + if (EC) + LOG_FATAL("Could not copy Mneme diff scatter task buffer EC: " + + EC.value()); + } + + bool buildDiffResetPlan(const EpilogueState &Epilogue, + ReplayResetMode Mode) { + releaseScatterPlan(); + DiffResetSpans.clear(); + + if (Mode == ReplayResetMode::Bytes) + return false; + + const std::string &EpilogueSnapshot = Epilogue.getSnapshotPath(); + std::string Error; + auto DiffPlan = + MnemeSnapshot::readDiffPlan(EpilogueSnapshot, &Error); + if (!DiffPlan) { + LOG_FATAL("Could not use diff reset for " + EpilogueSnapshot + ": " + + (Error.empty() ? "epilogue is not a diff snapshot" : Error)); + return false; + } + + if (DiffPlan->Globals.size() != this->GlobalVars.size()) { + Error = "diff global count does not match prologue"; + } else if (DiffPlan->Blobs.size() != this->DeviceMemoryState.size()) { + Error = "diff blob count does not match prologue"; + } + + if (!Error.empty()) { + LOG_FATAL("Could not use diff reset: " + Error); + return false; + } + + size_t MaxGap = + getEnvSizeOrDefault("MNEME_REPLAY_DIFF_SCATTER_MAX_GAP_BYTES", 0); + size_t ScatterTaskBytes = + getEnvSizeOrDefault("MNEME_REPLAY_DIFF_SCATTER_TASK_BYTES", 4096); + + for (auto &Global : DiffPlan->Globals) { + auto It = this->GlobalVars.find(Global.Name); + if (It == this->GlobalVars.end()) { + Error = "diff references global missing from prologue: " + Global.Name; + break; + } + auto &GVI = It->second; + if (GVI.VarSize != Global.VarSize) { + Error = "diff global size mismatch for: " + Global.Name; + break; + } + addScatterSpans(Global.Ranges, static_cast(GVI.HostAddr), + static_cast(GVI.DevAddr), MaxGap); + } + + if (Error.empty()) { + for (auto &BlobPlan : DiffPlan->Blobs) { + auto It = this->DeviceMemoryState.find(BlobPlan.DevAddr); + if (It == this->DeviceMemoryState.end()) { + Error = "diff references device allocation missing from prologue"; + break; + } + auto &Blob = It->second; + if (Blob.getActualSize() != BlobPlan.ActualSize || + Blob.getSize() != BlobPlan.Size) { + Error = "diff memory blob size mismatch"; + break; + } + addScatterSpans(BlobPlan.Ranges, Blob.getHostData().get(), + static_cast(Blob.getBlobAddr()), MaxGap); + } + } + + if (!Error.empty()) { + LOG_FATAL("Could not use diff reset: " + Error); + return false; + } + + prepareDeviceScatterPlan(ScatterTaskBytes); + + LOG_INFO("Prepared Mneme diff reset plan"); + return true; + } + + void resetFromDiffPlan(DeviceStream_t Stream) { + if (ScatterTaskCount == 0) + return; + auto EC = MnemeDeviceRT::DeviceErrorCheck( + launchDiffResetScatterKernel(DeviceScatterTasks, + ScatterTaskCount, Stream)); + if (EC) + LOG_FATAL("Could not launch Mneme diff scatter reset kernel EC: " + + EC.value()); + EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceStreamSynchronize(Stream)); + if (EC) + LOG_FATAL("Could not synchronize Mneme diff scatter reset kernel EC: " + + EC.value()); + } }; // Replay state for the expected kernel output. Unlike the prologue, load() @@ -195,8 +487,10 @@ class PrologueState : public ReplayMemState { template class EpilogueState : public ReplayMemState { public: - explicit EpilogueState(Snapshot SnapshotState) - : ReplayMemState(std::move(SnapshotState)) {} + explicit EpilogueState(Snapshot SnapshotState, + std::string SnapshotPath = "") + : ReplayMemState(std::move(SnapshotState)), + SnapshotPath(std::move(SnapshotPath)) {} void load() override { for (auto &[DevAddr, MemBlob] : this->DeviceMemoryState) { @@ -210,6 +504,10 @@ class EpilogueState : public ReplayMemState { EpilogueState *asEpilogue() override { return this; } + // Path of the snapshot file this state was reconstructed from. A diff reset + // plan re-reads this file to learn which byte ranges the kernel mutated. + const std::string &getSnapshotPath() const { return SnapshotPath; } + // Verifies a replayed prologue against this expected-output epilogue. At call // time the prologue's device buffers hold the kernel's actual output. virtual bool matches(const PrologueState &Prologue) const { @@ -265,6 +563,9 @@ class EpilogueState : public ReplayMemState { protected: bool isPrologue() const override { return false; } + +private: + std::string SnapshotPath; }; template @@ -282,7 +583,8 @@ makeReplayEpilogueState(const std::string &KernelName, Snapshot Snap = MnemeSnapshot::openSnapshot(SnapshotFile) ->reconstruct(KernelName, BasePrologueFile); - return std::make_unique>(std::move(Snap)); + return std::make_unique>(std::move(Snap), + SnapshotFile); } } // namespace mneme diff --git a/include/mneme/MnemeSnapshot.hpp b/include/mneme/MnemeSnapshot.hpp index 4b6a4062..a3d51273 100644 --- a/include/mneme/MnemeSnapshot.hpp +++ b/include/mneme/MnemeSnapshot.hpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -128,6 +129,40 @@ template class MnemeSnapshot { using DeviceError_t = typename MnemeDeviceRT::DeviceError_t; using DeviceStream_t = typename MnemeDeviceRT::DeviceStream_t; using KernelFunction_t = typename MnemeDeviceRT::KernelFunction_t; + +public: + using GlobalSnapshotData = std::unordered_map>; + + struct DiffRange { + size_t Offset = 0; + size_t Size = 0; + llvm::StringRef Bytes; + }; + + struct DiffGlobal { + std::string Name; + size_t VarSize = 0; + void *DevAddr = nullptr; + std::vector Ranges; + }; + + struct DiffBlob { + size_t ActualSize = 0; + size_t Size = 0; + void *DevAddr = nullptr; + Metadata Md; + std::vector Ranges; + }; + + struct DiffPlan { + std::unique_ptr Storage; + std::vector Globals; + std::vector Blobs; + size_t RawRangeCount = 0; + size_t ChangedBytes = 0; + }; + +private: static constexpr const char DiffMagic[] = "MNEME_DIFF_V1"; static constexpr size_t DiffMagicSize = sizeof(DiffMagic) - 1; static constexpr size_t DiffChunkSize = 1 << 20; @@ -249,18 +284,64 @@ template class MnemeSnapshot { util::writeBytes(OS, llvm::StringRef(DiffBytes.data(), DiffBytes.size())); } - static void applyDiffRanges(const char *&Buffer, - llvm::MutableArrayRef Target, - size_t NumRanges) { + static void readDiffRanges(const char *&Buffer, size_t NumRanges, + std::vector &Ranges, + DiffPlan &Plan) { + Ranges.reserve(NumRanges); for (size_t R = 0; R < NumRanges; ++R) { size_t Offset = util::extractScalar(Buffer); size_t Size = util::extractScalar(Buffer); - if (Offset > Target.size() || Size > Target.size() - Offset) - LOG_FATAL("Malformed Mneme diff range: offset " + - std::to_string(Offset) + " size " + std::to_string(Size) + - " exceeds target size " + std::to_string(Target.size())); - std::memcpy(Target.data() + Offset, Buffer, Size); + llvm::StringRef Bytes(Buffer, Size); Buffer += Size; + Ranges.push_back({Offset, Size, Bytes}); + Plan.RawRangeCount++; + Plan.ChangedBytes += Size; + } + } + + static DiffPlan parseDiffPayload(llvm::StringRef Bytes) { + const char *Buffer = Bytes.begin() + DiffMagicSize; + DiffPlan Plan; + + size_t TotalGlobals = util::extractScalar(Buffer); + Plan.Globals.reserve(TotalGlobals); + for (size_t I = 0; I < TotalGlobals; ++I) { + DiffGlobal Global; + Global.Name = util::readSizedString(Buffer); + Global.VarSize = util::extractScalar(Buffer); + Global.DevAddr = util::extractScalar(Buffer); + size_t NumRanges = util::extractScalar(Buffer); + readDiffRanges(Buffer, NumRanges, Global.Ranges, Plan); + Plan.Globals.push_back(std::move(Global)); + } + + size_t TotalMemBlobs = util::extractScalar(Buffer); + Plan.Blobs.reserve(TotalMemBlobs); + for (size_t I = 0; I < TotalMemBlobs; ++I) { + DiffBlob Blob; + Blob.ActualSize = util::extractScalar(Buffer); + Blob.Size = util::extractScalar(Buffer); + Blob.DevAddr = util::extractScalar(Buffer); + Blob.Md = metadata::fromBuffer(Buffer); + size_t NumRanges = util::extractScalar(Buffer); + readDiffRanges(Buffer, NumRanges, Blob.Ranges, Plan); + Plan.Blobs.push_back(std::move(Blob)); + } + + return Plan; + } + + static void applyDiffRanges(llvm::ArrayRef Ranges, + llvm::MutableArrayRef Target) { + for (auto &Range : Ranges) { + if (Range.Offset > Target.size() || + Range.Size > Target.size() - Range.Offset) + LOG_FATAL("Malformed Mneme diff range: offset " + + std::to_string(Range.Offset) + " size " + + std::to_string(Range.Size) + " exceeds target size " + + std::to_string(Target.size())); + std::memcpy(Target.data() + Range.Offset, Range.Bytes.data(), + Range.Size); } } @@ -307,63 +388,99 @@ template class MnemeSnapshot { std::unordered_map &GlobalVars, llvm::DenseMap> &DeviceMemory, llvm::MemoryBuffer *DiffBuffer) { - auto *Start = DiffBuffer->getBufferStart(); - auto *CurrentPtr = Start + DiffMagicSize; - size_t TotalGlobals = util::extractScalar(CurrentPtr); - if (TotalGlobals != GlobalVars.size()) + DiffPlan Diff = parseDiffPayload(DiffBuffer->getBuffer()); + + if (Diff.Globals.size() != GlobalVars.size()) LOG_FATAL("Mneme diff " + Filename + " does not match prologue global count"); - for (size_t I = 0; I < TotalGlobals; ++I) { - std::string Name = util::readSizedString(CurrentPtr); - size_t VarSize = util::extractScalar(CurrentPtr); - void *DevAddr = util::extractScalar(CurrentPtr); - size_t NumRanges = util::extractScalar(CurrentPtr); - - auto It = GlobalVars.find(Name); + for (auto &Global : Diff.Globals) { + auto It = GlobalVars.find(Global.Name); if (It == GlobalVars.end()) LOG_FATAL("Mneme diff references global missing from prologue: " + - Name); - if (It->second.VarSize != VarSize) - LOG_FATAL("Mneme diff global size mismatch for: " + Name); - It->second.DevAddr = DevAddr; + Global.Name); + if (It->second.VarSize != Global.VarSize) + LOG_FATAL("Mneme diff global size mismatch for: " + Global.Name); + It->second.DevAddr = Global.DevAddr; applyDiffRanges( - CurrentPtr, + Global.Ranges, llvm::MutableArrayRef( - static_cast(It->second.HostAddr), It->second.VarSize), - NumRanges); + static_cast(It->second.HostAddr), It->second.VarSize)); } - size_t TotalMemBlobs = util::extractScalar(CurrentPtr); - if (TotalMemBlobs != DeviceMemory.size()) + if (Diff.Blobs.size() != DeviceMemory.size()) LOG_FATAL("Mneme diff " + Filename + " does not match prologue memory blob count"); - for (size_t I = 0; I < TotalMemBlobs; ++I) { - size_t ActualSize = util::extractScalar(CurrentPtr); - size_t Size = util::extractScalar(CurrentPtr); - void *DeviceAddr = util::extractScalar(CurrentPtr); - auto MD = metadata::fromBuffer(CurrentPtr); - size_t NumRanges = util::extractScalar(CurrentPtr); - - auto It = DeviceMemory.find(DeviceAddr); + for (auto &BlobPlan : Diff.Blobs) { + auto It = DeviceMemory.find(BlobPlan.DevAddr); if (It == DeviceMemory.end()) LOG_FATAL("Mneme diff references device allocation missing from " "prologue"); auto &Blob = It->second; - if (Blob.getActualSize() != ActualSize || Blob.getSize() != Size) + if (Blob.getActualSize() != BlobPlan.ActualSize || + Blob.getSize() != BlobPlan.Size) LOG_FATAL("Mneme diff memory blob size mismatch"); - Blob.setMetadata(MD); + Blob.setMetadata(BlobPlan.Md); applyDiffRanges( - CurrentPtr, + BlobPlan.Ranges, llvm::MutableArrayRef(Blob.getHostData().get(), - Blob.getSize()), - NumRanges); + Blob.getSize())); } } public: - using GlobalSnapshotData = std::unordered_map>; + static bool isDiffSnapshotFile(const std::string &Filename, + std::string *Error = nullptr) { + if (!std::filesystem::exists(Filename)) { + if (Error) + *Error = "snapshot file does not exist"; + return false; + } + + std::ifstream Input(Filename, std::ios::binary); + if (!Input) { + if (Error) + *Error = "error when opening file"; + return false; + } + + char Magic[DiffMagicSize]; + Input.read(Magic, DiffMagicSize); + if (Input.gcount() != static_cast(DiffMagicSize)) + return false; + return llvm::StringRef(Magic, DiffMagicSize) == + llvm::StringRef(DiffMagic, DiffMagicSize); + } + + static std::optional + readDiffPlan(const std::string &Filename, std::string *Error = nullptr) { + auto SetError = [&](const std::string &Msg) -> std::optional { + if (Error) + *Error = Msg; + return std::nullopt; + }; + + if (!std::filesystem::exists(Filename)) + return SetError("snapshot file does not exist"); + + if (!isDiffSnapshotFile(Filename, Error)) + return SetError(Error && !Error->empty() + ? *Error + : "snapshot is not a Mneme diff snapshot"); + + llvm::ErrorOr> BufferOrErr = + llvm::MemoryBuffer::getFile(Filename, /*IsText=*/false, + /*RequiresNullTerminator=*/false); + if (std::error_code EC = BufferOrErr.getError()) + return SetError("error when opening file " + EC.message()); + + auto Storage = std::move(BufferOrErr.get()); + DiffPlan Plan = parseDiffPayload(Storage->getBuffer()); + Plan.Storage = std::move(Storage); + + return Plan; + } static std::pair fromBuffer(const char *&Buffer) { diff --git a/python/mneme/async_executor.py b/python/mneme/async_executor.py index 19c94174..91fe6117 100644 --- a/python/mneme/async_executor.py +++ b/python/mneme/async_executor.py @@ -92,6 +92,7 @@ def __init__( on_startup_failure_limit: Optional[ Callable[["TuneWorkerHandle", str], None] ] = None, + reset_mode: str = None, ): """ Construct a worker handle and start the worker process + monitor thread. @@ -143,6 +144,7 @@ def __init__( self.max_startup_failures = max_startup_failures self._startup_failures = 0 self._on_startup_failure_limit = on_startup_failure_limit + self.reset_mode = reset_mode self._state = None # ProcessEvent self._process = None # Process @@ -188,6 +190,7 @@ def _spawn_process(self): self.results_db_dir, self._state, self.warmup, + self.reset_mode, ), daemon=False, ) @@ -401,6 +404,7 @@ def __init__( num_workers: int, warmup: int = 2, max_startup_failures: int = 3, + reset_mode: str = None, ): """ Construct an asynchronous executor with a fixed-size worker pool. @@ -430,6 +434,7 @@ def __init__( self.iterations = iterations self.warmup = warmup self.max_startup_failures = max_startup_failures + self.reset_mode = reset_mode self._num_workers = num_workers self._failed_workers = set() self._broken_error = None @@ -448,6 +453,7 @@ def __init__( warmup=warmup, max_startup_failures=max_startup_failures, on_startup_failure_limit=self._handle_startup_failure_limit, + reset_mode=reset_mode, ) ) diff --git a/python/mneme/commands.py b/python/mneme/commands.py index 4545c8a1..f23109bc 100644 --- a/python/mneme/commands.py +++ b/python/mneme/commands.py @@ -542,6 +542,13 @@ def set_cli_args(parser): default=3, ) + parser.add_argument( + "--reset-mode", + choices=("bytes", "diff"), + default=None, + help="Memory reset mode used between replay iterations (default: choose from epilogue snapshot, or MNEME_REPLAY_RESET_MODE)", + ) + parser.add_argument( "--output-ll", "-ol", @@ -573,6 +580,7 @@ def __init__(self, *args, **kwargs): self.specialize_dims = kwargs.pop("specialize_dims", False) self.passes = kwargs.pop("passes", None) self.codegen_opt = kwargs.pop("codegen_opt", 3) + self.reset_mode = kwargs.get("reset_mode", None) self.output_ll = kwargs.pop("output_ll", None) diff --git a/python/mneme/device.py b/python/mneme/device.py index 24fbd7d3..6f18d86c 100644 --- a/python/mneme/device.py +++ b/python/mneme/device.py @@ -25,6 +25,7 @@ import weakref from ctypes import POINTER, Structure, c_char_p, c_float, c_int, c_uint, c_void_p +from typing import Optional from .llvm import ffi from .llvm.buffer import MemBufferRef @@ -62,6 +63,7 @@ MnemeRecordStateRef, c_int, c_int, + c_char_p, ] @@ -228,6 +230,7 @@ def profile( epilogue_state: MemBufferRef, shared_mem_size: int, iterations=5, + reset_mode: Optional[str] = None, ): """ Execute the kernel under Mneme record/replay profiling. @@ -253,6 +256,11 @@ def profile( Dynamic shared memory size (bytes) for the launch. iterations : int, optional Number of kernel executions to perform for profiling. + reset_mode : str, optional + Replay memory reset mode: ``"bytes"`` or ``"diff"``. + When omitted, the native runtime uses ``MNEME_REPLAY_RESET_MODE`` or + selects ``"diff"`` for diff epilogue snapshots and ``"bytes"`` + otherwise. Raises ------ @@ -273,6 +281,7 @@ def profile( epilogue_state, shared_mem_size, iterations, + _encode_string(reset_mode or ""), ) return diff --git a/python/mneme/replay_executor.py b/python/mneme/replay_executor.py index 026cf83c..b4136378 100644 --- a/python/mneme/replay_executor.py +++ b/python/mneme/replay_executor.py @@ -34,7 +34,7 @@ import os from datetime import datetime, timezone from multiprocessing import Event, Queue -from typing import Tuple +from typing import Optional, Tuple from mneme.device import ( DeviceFunction, @@ -54,6 +54,22 @@ from mneme.transforms import transform from mneme.utils import cond_gpu_time, cond_time +VALID_RESET_MODES = {"bytes", "diff"} + + +def normalize_reset_mode(reset_mode: Optional[str]) -> str: + mode = reset_mode or os.environ.get("MNEME_REPLAY_RESET_MODE", "") + if not mode: + return "" + mode = mode.lower() + if mode == "full": + mode = "bytes" + if mode not in VALID_RESET_MODES: + raise ValueError( + f"Unknown reset mode '{mode}'. Expected one of: bytes, diff" + ) + return mode + class BaseExecutor: """ @@ -109,6 +125,7 @@ def __init__( iterations: int = 3, device_id: int = 0, warmup: int = 2, + reset_mode: Optional[str] = None, ): self.record_db = record_db self.record_id = record_id @@ -124,6 +141,7 @@ def __init__( self._page_manager = None self._iterations = iterations self._warmup = warmup + self.reset_mode = normalize_reset_mode(reset_mode) self.num_devices = get_device_count() set_device(device_id) logger.debug( @@ -403,6 +421,7 @@ def _run_kernel( self._epilogue._state, config.shared_mem, iterations, + reset_mode=self.reset_mode, ) def _build( @@ -598,21 +617,34 @@ def _execute( if self._prologue._state is None or self._epilogue._state is None: raise RuntimeError("States should never be none when executing a kernel") - # NOTE: 1. First we need to verify. - ver_mod = ir_module.clone() - mem_buffer = self._build(result, config, ver_mod, False) - self._run(result, config, mem_buffer, self.prologue, self.epilogue, True, False, 1) - - # NOTE: 2. We apply a custom pass to delete all clang insered code. + # NOTE: 1. We apply a custom pass to delete all clang insered code. # It is hard to identify these cases, So we delete only things # that have been attributed by clang ir_module = transform.remove_auto_initialize(ir_module.clone()) - # Done with verification. Moving to next stage - # NOTE: 3. We build and run. We set tracking on and execute warmups plus iterations, - # to enalbe later computation of statistical metrics etc. + # NOTE: 2. Build the final replay object, verify that exact object once, + # then run the tracked repetitions. Diff reset is only used after the + # native runtime has performed a full-reset first replay in each profile + # call, so failed verification never depends on the optimized reset path. mem_buffer = self._build(result, config, ir_module, True) - self._run(result, config, mem_buffer, self.prologue, self.epilogue, False, True, self._iterations + self._warmup) + self._run( + result, config, mem_buffer, self.prologue, self.epilogue, True, False, 1 + ) + if not result.verified: + return ir_module + + # NOTE: 3. We set tracking on and execute warmups plus iterations, + # to enalbe later computation of statistical metrics etc. + self._run( + result, + config, + mem_buffer, + self.prologue, + self.epilogue, + False, + True, + self._iterations + self._warmup, + ) result.executed = True return ir_module @@ -719,6 +751,7 @@ def run( results_db_dir: str, state: Event, warmup: int = 2, + reset_mode: Optional[str] = None, ): """ Worker process entry point: initialize resources and serve requests from a queue. @@ -768,6 +801,8 @@ def run( state : multiprocessing.Event Event used to signal to the parent process that initialization is complete and the worker is ready to accept requests. + reset_mode : str, optional + Replay memory reset mode passed to the native runtime. Notes ----- @@ -792,6 +827,7 @@ def run( device_id=device_id, iterations=iterations, warmup=warmup, + reset_mode=reset_mode, ) # Open GPU memory, setup prologue epilogue and create a single # LLVM IR file to start working on optimizations diff --git a/python/tests/test_device.py b/python/tests/test_device.py index 5ad6c2c4..63f86a40 100644 --- a/python/tests/test_device.py +++ b/python/tests/test_device.py @@ -153,9 +153,11 @@ def test_device_function_profile_calls_ffi(): epilogue_state=buf, shared_mem_size=0, iterations=5, + reset_mode="diff", ) fake_lib.MnemePy_profile.assert_called_once() + assert fake_lib.MnemePy_profile.call_args[0][-1] == b"diff" def test_device_function_profile_fails_if_module_gc(): @@ -202,4 +204,3 @@ def test_set_device_calls_ffi(): set_device(2) fake_lib.MnemePy_setDevice.assert_called_once_with(2) - diff --git a/python/tests/test_replay_executor.py b/python/tests/test_replay_executor.py index cb3f3ab0..721fd459 100644 --- a/python/tests/test_replay_executor.py +++ b/python/tests/test_replay_executor.py @@ -106,9 +106,18 @@ def __init__(self): self.const_mem = 33 self.local_mem = 44 - def profile(self, grid, block, pro_state, epi_state, shared_mem, iterations): + def profile( + self, + grid, + block, + pro_state, + epi_state, + shared_mem, + iterations, + reset_mode="", + ): self.profile_calls.append( - (grid, block, pro_state, epi_state, shared_mem, iterations) + (grid, block, pro_state, epi_state, shared_mem, iterations, reset_mode) ) @@ -217,7 +226,11 @@ def fake_from_json(path): ) ex = mod.BaseExecutor( - record_db="db.json", record_id="rid", device_id=3, iterations=5 + record_db="db.json", + record_id="rid", + device_id=3, + iterations=5, + reset_mode="diff", ) assert ex.records is rec @@ -227,6 +240,28 @@ def fake_from_json(path): assert calls["set_device"] == [3] assert calls["from_json"] == ["db.json"] assert ex._iterations == 5 + assert ex.reset_mode == "diff" + + +def test_baseexecutor_reset_mode_env_and_validation(monkeypatch): + mod = _reload_with_identity_decorators(monkeypatch) + + kernel = FakeKernelDescr() + rec = FakeRecordedExecution(kernel) + + monkeypatch.setattr( + mod.RecordedExecution, "from_json", staticmethod(lambda _: rec), raising=True + ) + monkeypatch.setattr(mod, "set_device", lambda _: None, raising=True) + monkeypatch.setattr(mod, "get_device_arch", lambda: "sm", raising=True) + monkeypatch.setattr(mod, "get_device_count", lambda: 1, raising=True) + monkeypatch.setenv("MNEME_REPLAY_RESET_MODE", "full") + + ex = mod.BaseExecutor(record_db="x", record_id="rid") + assert ex.reset_mode == "bytes" + + with pytest.raises(ValueError): + mod.BaseExecutor(record_db="x", record_id="rid", reset_mode="bogus") def test_open_close_context_manager_opens_and_closes_resources(monkeypatch): @@ -496,9 +531,10 @@ def fake_run(result, cfg, mem_buf, prologue, epilogue, verify, track, iters): out_ir = ex._execute(res, cfg, ir) - # Verification stage: track=False, iterations=1 + # Final object verification stage: track=False, iterations=1 # Tracked stage: track=True, iterations=self._iterations + 2 = 5 assert run_calls == [(False, 1), (True, 5)] + assert build_calls == [("root_clone", True)] assert res.executed is True assert res.verified is True @@ -556,12 +592,15 @@ def test_tuneworker_run_process_and_terminate(monkeypatch, tmp_path): monkeypatch.setattr(mod.os, "dup2", lambda *a, **k: None, raising=True) class FakeWorker: - def __init__(self, record_db, record_id, device_id, iterations, warmup): + def __init__( + self, record_db, record_id, device_id, iterations, warmup, reset_mode=None + ): self.record_db = record_db self.record_id = record_id self.device_id = device_id self.iterations = iterations self.warmup = warmup + self.reset_mode = reset_mode def link_ir(self): return FakeModule("root_ir") diff --git a/src/MnemeDeviceCode.cpp b/src/MnemeDeviceCode.cpp index dc62531d..e40cbeec 100644 --- a/src/MnemeDeviceCode.cpp +++ b/src/MnemeDeviceCode.cpp @@ -1,6 +1,8 @@ #include "mneme/DeviceTraits.hpp" #include "mneme/MnemeAnnotation.hpp" #include "mneme/MnemeComparators.hpp" +#include "mneme/MnemeDeviceKernels.hpp" +#include #include #include #include @@ -13,11 +15,13 @@ using namespace mneme; #ifdef MNEME_ENABLE_HIP #include using DeviceVendorTraits = DeviceTraits; +constexpr DeviceVendors ActiveDeviceVendor = DeviceVendors::HIP; #define MNEME_DEV __device__ __forceinline__ #elif defined(MNEME_ENABLE_CUDA) #include #include // or hip/hip_fp16.h for HIP using DeviceVendorTraits = DeviceTraits; +constexpr DeviceVendors ActiveDeviceVendor = DeviceVendors::CUDA; #define MNEME_DEV __device__ __forceinline__ #endif @@ -209,12 +213,43 @@ compare_builtin_kernel(const T *__restrict__ a, const T *__restrict__ b, atomicAddDouble(&out->Agg, LOut.Agg); } } + +__global__ void diff_reset_scatter_kernel(const DiffResetScatterTask *Tasks, + std::size_t NumTasks) { + for (std::size_t TaskIndex = blockIdx.x; TaskIndex < NumTasks; + TaskIndex += gridDim.x) { + auto Task = Tasks[TaskIndex]; + for (std::size_t I = threadIdx.x; I < Task.Size; I += blockDim.x) + Task.Dst[I] = Task.Src[I]; + } +} } // namespace dev } // namespace mneme constexpr int NUM_THREADS_PER_BLOCK = 256; +constexpr int MAX_DIFF_SCATTER_BLOCKS = 131072; namespace mneme { +template <> +DeviceVendorTraits::DeviceError_t +launchDiffResetScatterKernel( + const DiffResetScatterTask *Tasks, size_t NumTasks, + DeviceVendorTraits::DeviceStream_t Stream) { + if (NumTasks == 0) + return DeviceVendorTraits::DeviceSuccess; + + auto NumBlocks = + std::min(NumTasks, static_cast(MAX_DIFF_SCATTER_BLOCKS)); + mneme::dev::diff_reset_scatter_kernel<<>>(Tasks, NumTasks); + +#ifdef MNEME_ENABLE_HIP + return hipGetLastError(); +#else + return cudaGetLastError(); +#endif +} + CompareResult compareDeviceBlobs(const char *Blob1, const char *Blob2, uint64_t NumBytes, Metadata Md) { CompareResult Init{}; diff --git a/src/python/device.cpp b/src/python/device.cpp index 279ce393..cfbdeac8 100644 --- a/src/python/device.cpp +++ b/src/python/device.cpp @@ -1,4 +1,7 @@ #include "llvm/core.h" +#include +#include +#include #include #include #include @@ -103,8 +106,8 @@ API_EXPORT(const char *) MnemePy_getDeviceArch() { API_EXPORT(void) MnemePy_profile(void *WrappedModule, void *Func, dim3 Grid, dim3 Block, MnemeDeviceMemStateRef Prologue, - MnemeDeviceMemStateRef Epilogue, int SharedMemSize, - int repeats) { + MnemeDeviceMemStateRef Epilogue, int SharedMemSize, int repeats, + const char *ResetModeName) { auto VendorModule = reinterpret_cast(WrappedModule); @@ -118,38 +121,109 @@ MnemePy_profile(void *WrappedModule, void *Func, dim3 Grid, dim3 Block, DeviceVendorTraits::DeviceStream_t ReplayStream; auto DevFunc = reinterpret_cast(Func); - auto PrologueState = unwrap(Prologue); - auto EpilogueState = unwrap(Epilogue); + auto *ProState = unwrap(Prologue)->asPrologue(); + auto *EpiState = unwrap(Epilogue)->asEpilogue(); + if (!ProState || !EpiState) + LOG_FATAL("MnemePy_profile expects a prologue and an epilogue state"); + std::string ResetModeString = + ResetModeName == nullptr ? "" : std::string(ResetModeName); + ReplayResetMode ResetMode; + if (ResetModeString.empty()) { + const char *EnvMode = std::getenv("MNEME_REPLAY_RESET_MODE"); + if (EnvMode != nullptr) { + ResetMode = parseReplayResetMode(EnvMode); + } else if (MnemeSnapshot::isDiffSnapshotFile( + EpiState->getSnapshotPath())) { + ResetMode = ReplayResetMode::Diff; + } else { + ResetMode = ReplayResetMode::Bytes; + } + } else { + ResetMode = parseReplayResetMode(ResetModeString); + } + bool TraceResetTiming = std::getenv("MNEME_REPLAY_RESET_TIMING") != nullptr; + EC = DeviceVendorTraits::DeviceErrorCheck( DeviceVendorTraits::DeviceStreamCreate(&ReplayStream)); if (EC) LOG_FATAL("Error when creating a stream for replay\n" + EC.value()); - PrologueState->initializeGlobals(VendorModule); + ProState->initializeGlobals(VendorModule); + ProState->prepareResetPlan(*EpiState, ResetMode); for (int i = 0; i < repeats; i++) { LOG_DEBUG("Run {}/{}", i + 1, repeats); - PrologueState->reset(); - auto Args = PrologueState->getArgs(); + auto IterStart = std::chrono::steady_clock::now(); + auto ResetStart = std::chrono::steady_clock::now(); + if (i == 0) + ProState->reset(); + else + ProState->reset(ResetMode, ReplayStream); + auto ResetStop = std::chrono::steady_clock::now(); + if (TraceResetTiming) { + auto ResetNs = std::chrono::duration_cast( + ResetStop - ResetStart) + .count(); + bool UsedDiff = i != 0 && ResetMode == ReplayResetMode::Diff; + const char *ResetKind = "full"; + if (UsedDiff) + ResetKind = "diff"; + std::fprintf(stderr, + "MNEME_RESET_TIMING repeat=%d/%d kind=%s ns=%lld\n", + i + 1, repeats, ResetKind, static_cast(ResetNs)); + } + auto Args = ProState->getArgs(); + auto PreSyncStart = std::chrono::steady_clock::now(); EC = DeviceVendorTraits::DeviceErrorCheck( DeviceVendorTraits::DeviceStreamSynchronize(ReplayStream)); + auto PreSyncStop = std::chrono::steady_clock::now(); if (EC) LOG_FATAL("Error when synchronizing device stream " + EC.value()); + auto LaunchStart = std::chrono::steady_clock::now(); EC = DeviceVendorTraits::DeviceErrorCheck( DeviceVendorTraits::launchKernelFunction(DevFunc, Grid, Block, Args, SharedMemSize, ReplayStream)); + auto LaunchStop = std::chrono::steady_clock::now(); if (EC) LOG_FATAL("Error When Launching Kernel: " + EC.value()); + auto PostSyncStart = std::chrono::steady_clock::now(); EC = DeviceVendorTraits::DeviceErrorCheck( DeviceVendorTraits::DeviceStreamSynchronize(ReplayStream)); + auto PostSyncStop = std::chrono::steady_clock::now(); if (EC) LOG_FATAL("Error When synchronizing with kernel stream: " + EC.value()); + + if (TraceResetTiming) { + auto IterNs = std::chrono::duration_cast( + PostSyncStop - IterStart) + .count(); + auto PreSyncNs = std::chrono::duration_cast( + PreSyncStop - PreSyncStart) + .count(); + auto LaunchNs = std::chrono::duration_cast( + LaunchStop - LaunchStart) + .count(); + auto PostSyncNs = std::chrono::duration_cast( + PostSyncStop - PostSyncStart) + .count(); + bool UsedDiff = i != 0 && ResetMode == ReplayResetMode::Diff; + const char *ReplayKind = "full"; + if (UsedDiff) + ReplayKind = "diff"; + std::fprintf(stderr, + "MNEME_REPLAY_TIMING repeat=%d/%d kind=%s iter_ns=%lld " + "pre_sync_ns=%lld launch_ns=%lld post_sync_ns=%lld\n", + i + 1, repeats, ReplayKind, static_cast(IterNs), + static_cast(PreSyncNs), + static_cast(LaunchNs), + static_cast(PostSyncNs)); + } } EC = DeviceVendorTraits::DeviceErrorCheck( diff --git a/tests/unit_tests/CMakeLists.txt b/tests/unit_tests/CMakeLists.txt index fe93e071..389faeaf 100644 --- a/tests/unit_tests/CMakeLists.txt +++ b/tests/unit_tests/CMakeLists.txt @@ -57,7 +57,7 @@ add_test(NAME Serialize::Globals COMMAND SerializeGlobals) BUILD_UNIT_TEST("SerializeBlob" SerializeBlob.cpp) add_test(NAME Serialize::Blob COMMAND SerializeBlob) -BUILD_UNIT_TEST("ReadWriteSnapshot" ReadWriteSnapshot.cpp) +BUILD_DEVICE_UNIT_TEST("ReadWriteSnapshot" ReadWriteSnapshot.cpp) add_test(NAME Serialize::RWSnapshot COMMAND ReadWriteSnapshot) BUILD_DEVICE_UNIT_TEST("CompareMemoryBlob" DeviceKernels.cpp) diff --git a/tests/unit_tests/ReadWriteSnapshot.cpp b/tests/unit_tests/ReadWriteSnapshot.cpp index 016bbc1c..a9e9a9b4 100644 --- a/tests/unit_tests/ReadWriteSnapshot.cpp +++ b/tests/unit_tests/ReadWriteSnapshot.cpp @@ -2,6 +2,7 @@ #include "mneme/MnemeAnnotation.hpp" #include "mneme/MnemeKernelInfo.hpp" #include "mneme/MnemeLogger.hpp" +#include "mneme/MnemeReplay.hpp" #include "mneme/MnemeSnapshot.hpp" #include #include @@ -97,9 +98,9 @@ int main(int argc, char **argv) { std::filesystem::path SnapshotFN("./test.mneme"); MnemeSnapshot::GlobalSnapshotData PrologueGlobals; - MnemeSnapshot::takeMnemeBytesSnapshot( - GVars, DeviceMemMap, SnapshotFN, TestKernel->KernelArgSizes, Args, 0, - &PrologueGlobals); + MnemeSnapshot::takeMnemeBytesSnapshot(GVars, DeviceMemMap, SnapshotFN, + TestKernel->KernelArgSizes, + Args, 0, &PrologueGlobals); auto ReadSnap = MnemeSnapshot::readBytesSnapshot(KernelName, SnapshotFN.string()); @@ -229,15 +230,15 @@ int main(int argc, char **argv) { GlobalData.second[3] ^= 0x9; GlobalData.second[4] ^= 0x13; - auto EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceCopy( - BlobData.first, BlobData.second, 128, - MnemeDeviceRT::MemcpyHostToDeviceKind())); + auto EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceCopy(BlobData.first, BlobData.second, 128, + MnemeDeviceRT::MemcpyHostToDeviceKind())); if (EC) LOG_FATAL("Could not update device blob data"); - EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceCopy( - GlobalData.first, GlobalData.second, 128, - MnemeDeviceRT::MemcpyHostToDeviceKind())); + EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceCopy(GlobalData.first, GlobalData.second, 128, + MnemeDeviceRT::MemcpyHostToDeviceKind())); if (EC) LOG_FATAL("Could not update device global data"); @@ -251,6 +252,30 @@ int main(int argc, char **argv) { auto &DiffDeviceMemMap = DiffSnap.DeviceMemory; auto &DiffKernel = DiffSnap.KInfo; + std::string DiffPlanError; + auto DiffPlan = MnemeSnapshot::readDiffPlan(DiffSnapshotFN.string(), + &DiffPlanError); + auto ValidateDiffPlan = [&]() { + if (!DiffPlan) { + std::cerr << "Could not parse diff plan: " << DiffPlanError << "\n"; + return 64; + } + if (DiffPlan->Globals.size() != 1 || DiffPlan->Blobs.size() != 1) { + std::cerr << "Diff plan counts differ\n"; + return 64; + } + if (DiffPlan->RawRangeCount == 0 || DiffPlan->ChangedBytes == 0) { + std::cerr << "Diff plan did not record changed ranges\n"; + return 64; + } + if (DiffPlan->Globals[0].Ranges.empty() || + DiffPlan->Blobs[0].Ranges.empty()) { + std::cerr << "Diff plan missing global or blob ranges\n"; + return 64; + } + return 0; + }(); + auto ValidateDiffGlobalMem = [&]() { auto it = DiffGVars.find("Test"); if (it == DiffGVars.end()) @@ -319,14 +344,79 @@ int main(int argc, char **argv) { return 0; }(); + auto ValidateReplayDiffReset = [&]() { + PrologueState PrologueReplay(KernelName, SnapshotFN.string()); + EpilogueState EpilogueReplay( + MnemeSnapshot::readDiffSnapshot( + KernelName, DiffSnapshotFN.string(), SnapshotFN.string()), + DiffSnapshotFN.string()); + + PrologueReplay.load(); + + auto CheckMode = [&](ReplayResetMode Mode) { + PrologueReplay.prepareResetPlan(EpilogueReplay, Mode); + + auto EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceCopy(BlobData.first, BlobData.second, 128, + MnemeDeviceRT::MemcpyHostToDeviceKind())); + if (EC) + LOG_FATAL("Could not poison reset device blob data"); + + EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceCopy(GlobalData.first, GlobalData.second, 128, + MnemeDeviceRT::MemcpyHostToDeviceKind())); + if (EC) + LOG_FATAL("Could not poison reset global data"); + + PrologueReplay.reset(Mode); + + std::unique_ptr DeviceCheck(new uint8_t[128]); + EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceCopy(DeviceCheck.get(), BlobData.first, 128, + MnemeDeviceRT::MemcpyDeviceToHostKind())); + if (EC) + LOG_FATAL("Could not copy reset device blob data"); + + auto BlobIt = ReadDeviceMemMap.find((void *)BlobData.first); + if (BlobIt == ReadDeviceMemMap.end() || + std::memcmp(DeviceCheck.get(), BlobIt->second.getHostData().get(), + 128) != 0) { + std::cerr << "Diff reset did not restore prologue blob data\n"; + return 128; + } + + EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceCopy(DeviceCheck.get(), GlobalData.first, 128, + MnemeDeviceRT::MemcpyDeviceToHostKind())); + if (EC) + LOG_FATAL("Could not copy reset global data"); + + auto GlobalIt = ReadGVars.find("Test"); + if (GlobalIt == ReadGVars.end() || + std::memcmp(DeviceCheck.get(), GlobalIt->second.HostAddr, 128) != 0) { + std::cerr << "Diff reset did not restore prologue global data\n"; + return 128; + } + + return 0; + }; + + auto DiffRet = CheckMode(ReplayResetMode::Diff); + + PrologueReplay.release(); + return DiffRet; + }(); + auto Ret = ValidateGlobalMem | ValidateDeviceMem | ValidateKernelArgs | ValidateDiffGlobalMem | ValidateDiffDeviceMem | - ValidateDiffKernelArgs; + ValidateDiffKernelArgs | ValidateDiffPlan | + ValidateReplayDiffReset; delete[] GlobalData.second; delete[] BlobData.second; - EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceFree(GlobalData.first)); + EC = MnemeDeviceRT::DeviceErrorCheck( + MnemeDeviceRT::DeviceFree(GlobalData.first)); if (EC) LOG_FATAL("Could not release device memory\n");