From 69bb8457a85e16e7c57ab9271110681e2034d96e Mon Sep 17 00:00:00 2001 From: Caetano Melone Date: Tue, 18 Aug 2026 11:23:15 -0700 Subject: [PATCH 1/5] add subregion annotations Signed-off-by: Caetano Melone --- include/mneme/MnemeAnnotation.hpp | 33 +++- include/mneme/MnemeMemory.hpp | 199 ++++++++++++++++++++++-- include/mneme/MnemeNullRecorder.hpp | 3 + include/mneme/MnemeRecord.hpp | 4 + include/mneme/MnemeRecorderBackend.hpp | 2 + include/mneme/MnemeRecordingBackend.hpp | 76 +++++++++ include/mneme/MnemeSnapshot.hpp | 123 +++++++++++++++ src/MnemeAnnotation.cpp | 20 ++- src/MnemeAnnotationRuntime.hpp | 8 +- src/MnemeRecordPreloadCUDA.cpp | 6 + src/MnemeRecordPreloadHIP.cpp | 6 + 11 files changed, 461 insertions(+), 19 deletions(-) diff --git a/include/mneme/MnemeAnnotation.hpp b/include/mneme/MnemeAnnotation.hpp index de946bf4..78132f6f 100644 --- a/include/mneme/MnemeAnnotation.hpp +++ b/include/mneme/MnemeAnnotation.hpp @@ -2,7 +2,8 @@ //===----------------------------------------------------------------------===// // mneme_annotate.h - User-facing annotation API for Mneme // -// This header declares mneme::annotate(ptr, Metadata{...}) and related types. +// This header declares mneme::annotate(ptr, Metadata{...}) and sub-region +// annotation overloads such as mneme::annotate(ptr, nbytes, Metadata{...}). // It is intentionally *interface-only*: no lambdas / comparators yet. // // Example: @@ -18,9 +19,14 @@ // You can also use the typed helper: // mneme::annotate(p, mneme::Metadata{ .threshold = 0.1 }); // +// Or annotate a sub-region explicitly: +// mneme::annotate(p, 128 * sizeof(double), +// mneme::Metadata{ .threshold = 0.1 }); +// //===----------------------------------------------------------------------===// #include +#include #include #include #include @@ -72,6 +78,16 @@ struct Metadata { std::optional tag = std::nullopt; }; +inline bool operator==(const Metadata &LHS, const Metadata &RHS) { + return LHS.builtin == RHS.builtin && LHS.threshold == RHS.threshold && + LHS.threshold_kind == RHS.threshold_kind && LHS.norm == RHS.norm && + LHS.tag == RHS.tag; +} + +inline bool operator!=(const Metadata &LHS, const Metadata &RHS) { + return !(LHS == RHS); +} + // --------- Builtin dtype mapping helpers (optional sugar) ------------------ template struct builtin_dtype_of { @@ -115,17 +131,26 @@ template <> struct builtin_dtype_of { // Primary user API: annotate a pointer with metadata. void annotate(const void *ptr, Metadata md); -// Convenience overload for non-const pointers. -void annotate(void *ptr, Metadata md); +// Annotate a sub-region beginning at ptr and spanning bytes bytes. +void annotate(const void *ptr, std::size_t bytes, Metadata md); // Typed helper: sets builtin dtype automatically when T maps to a known // BuiltinDType. If T is unknown, this will leave builtin=Unknown (still useful // if you set dtype=Custom later). -template inline void annotate(T *ptr, Metadata md = {}) { +template >, int> = 0> +inline void annotate(T *ptr, Metadata md = {}) { // If the user didn't specify dtype explicitly, keep default Builtin. // If they *did* specify Custom, we don't override anything here. md.builtin = builtin_dtype_of>::value; annotate(static_cast(ptr), std::move(md)); } +template >, int> = 0> +inline void annotate(T *ptr, std::size_t bytes, Metadata md) { + md.builtin = builtin_dtype_of>::value; + annotate(static_cast(ptr), bytes, std::move(md)); +} + } // namespace mneme diff --git a/include/mneme/MnemeMemory.hpp b/include/mneme/MnemeMemory.hpp index dde65b10..6171472c 100644 --- a/include/mneme/MnemeMemory.hpp +++ b/include/mneme/MnemeMemory.hpp @@ -1,10 +1,12 @@ #pragma once +#include #include #include #include #include #include +#include #include #include @@ -16,6 +18,39 @@ #include "mneme/MnemeUtils.hpp" namespace mneme { +// Metadata attached to a byte subrange inside one owning blob. +struct MemoryRegionMetadata { + uint64_t Offset = 0; + uint64_t Extent = 0; + Metadata MD; + + uint64_t endOffset() const { return Offset + Extent; } + + bool contains(uint64_t ByteOffset) const { + return Offset <= ByteOffset && ByteOffset < endOffset(); + } + + bool hasSameRange(uint64_t OtherOffset, uint64_t OtherExtent) const { + return Offset == OtherOffset && Extent == OtherExtent; + } + + bool overlaps(uint64_t OtherOffset, uint64_t OtherExtent) const { + return Offset < OtherOffset + OtherExtent && + OtherOffset < Offset + Extent; + } +}; + +inline bool operator==(const MemoryRegionMetadata &LHS, + const MemoryRegionMetadata &RHS) { + return LHS.Offset == RHS.Offset && LHS.Extent == RHS.Extent && + LHS.MD == RHS.MD; +} + +inline bool operator!=(const MemoryRegionMetadata &LHS, + const MemoryRegionMetadata &RHS) { + return !(LHS == RHS); +} + template class MnemeMemoryBlob { public: using MnemeDeviceRT = DeviceTraits; @@ -26,13 +61,41 @@ template class MnemeMemoryBlob { typename MnemeDeviceRT::MemoryAllocationHandle_t; protected: + static bool compareRange(const char *Expected, const char *Actual, + uint64_t NumBytes, const Metadata &Md) { + if (NumBytes == 0) + return true; + + auto Compare = compareDeviceBlobs(Expected, Actual, NumBytes, Md); + + // Comparator semantics: + // - Norm::None => per-element thresholding reports AnyFail/FirstBadIdx + // - Norm::L1/L2/Linf => aggregated error is reported in Agg + if (Md.norm == Norm::None) + return Compare.AnyFail == 0; + + return Compare.Agg <= Md.threshold; + } + Metadata PtrMD; + std::vector RegionMD; uint64_t ActualSize; void *BlobAddr; uint64_t Size; std::unique_ptr HostData; bool IsMapped; + bool isValidRegionRange(uint64_t Offset, uint64_t Extent) const { + return Extent != 0 && Offset <= Size && Extent <= Size - Offset; + } + + static bool regionSortLess(const MemoryRegionMetadata &LHS, + const MemoryRegionMetadata &RHS) { + if (LHS.Offset != RHS.Offset) + return LHS.Offset < RHS.Offset; + return LHS.Extent < RHS.Extent; + } + public: MnemeMemoryBlob(uint64_t ActualSize = 0, void *BlobAddr = nullptr, uint64_t Size = 0) @@ -111,6 +174,7 @@ template class MnemeMemoryBlob { HostData = std::move(other.HostData); IsMapped = other.IsMapped; PtrMD = other.PtrMD; + RegionMD = std::move(other.RegionMD); other.BlobAddr = 0; other.HostData = nullptr; } @@ -120,7 +184,8 @@ template class MnemeMemoryBlob { MnemeMemoryBlob(MnemeMemoryBlob &&other) noexcept : BlobAddr(other.BlobAddr), Size(other.Size), ActualSize(other.ActualSize), HostData(std::move(other.HostData)), - IsMapped(other.IsMapped), PtrMD(other.PtrMD) { + IsMapped(other.IsMapped), PtrMD(other.PtrMD), + RegionMD(std::move(other.RegionMD)) { other.BlobAddr = 0; other.HostData = nullptr; } @@ -138,24 +203,136 @@ template class MnemeMemoryBlob { Metadata getMetadata() const { return PtrMD; } + // Distinct overlapping regions are rejected. Re-registering the exact same + // [offset, extent) updates the existing metadata in place. + bool setRegionMetadata(uint64_t Offset, uint64_t Extent, Metadata Md) { + if (!isValidRegionRange(Offset, Extent)) + return false; + + MemoryRegionMetadata Key{Offset, Extent, {}}; + auto It = std::lower_bound(RegionMD.begin(), RegionMD.end(), Key, + regionSortLess); + + if (It != RegionMD.end() && It->hasSameRange(Offset, Extent)) { + It->MD = std::move(Md); + return true; + } + + if (It != RegionMD.begin()) { + auto Prev = std::prev(It); + if (Prev->overlaps(Offset, Extent)) + return false; + } + if (It != RegionMD.end() && It->overlaps(Offset, Extent)) + return false; + + RegionMD.insert(It, {Offset, Extent, std::move(Md)}); + return true; + } + + // Replace the entire region table after validating a bulk-loaded set, such + // as snapshot deserialization. + bool replaceRegionMetadata(std::vector Regions) { + for (const auto &Region : Regions) { + if (!isValidRegionRange(Region.Offset, Region.Extent)) + return false; + } + + std::sort(Regions.begin(), Regions.end(), regionSortLess); + for (size_t I = 1; I < Regions.size(); ++I) { + if (Regions[I - 1].overlaps(Regions[I].Offset, Regions[I].Extent)) + return false; + } + + RegionMD = std::move(Regions); + return true; + } + + void clearRegionMetadata() { RegionMD.clear(); } + + bool hasRegionMetadata() const { return !RegionMD.empty(); } + + const std::vector &getRegionMetadata() const { + return RegionMD; + } + + // RegionMD is kept sorted by offset, so the owning region for a byte can be + // resolved by finding the last region whose start is <= ByteOffset. + const MemoryRegionMetadata *findRegionMetadata(uint64_t ByteOffset) const { + auto It = std::upper_bound( + RegionMD.begin(), RegionMD.end(), ByteOffset, + [](uint64_t Offset, const MemoryRegionMetadata &Region) { + return Offset < Region.Offset; + }); + if (It == RegionMD.begin()) + return nullptr; + + --It; + return It->contains(ByteOffset) ? &*It : nullptr; + } + bool operator==(const MnemeMemoryBlob &other) const { if (getSize() != other.getSize()) { LOG_WARN("Sizes Differ {} vs {}", getSize(), other.getSize()); return false; } - auto Md = getMetadata(); - auto Compare = compareDeviceBlobs((const char *)other.getBlobAddr(), - (const char *)getBlobAddr(), getSize(), - Md); + if (getMetadata() != other.getMetadata()) { + LOG_WARN("Whole-blob metadata differs for blob {}", getBlobAddr()); + return false; + } + + // Compare the region tables first so the subsequent byte comparison can + // assume both blobs partition the address range the same way. + const auto &Regions = getRegionMetadata(); + const auto &OtherRegions = other.getRegionMetadata(); - // Comparator semantics: - // - Norm::None => per-element thresholding reports AnyFail/FirstBadIdx - // - Norm::L1/L2/Linf => aggregated error is reported in Agg - if (Md.norm == Norm::None) - return Compare.AnyFail == 0; + if (Regions.size() != OtherRegions.size()) { + LOG_WARN("Region metadata count differs for blob {}", getBlobAddr()); + return false; + } - return Compare.Agg <= Md.threshold; + for (size_t I = 0; I < Regions.size(); ++I) { + const auto &ExpectedRegion = Regions[I]; + const auto &ActualRegion = OtherRegions[I]; + if (ExpectedRegion != ActualRegion) { + LOG_WARN("Region metadata differs for blob {} at region index {}", + getBlobAddr(), I); + return false; + } + } + + if (Regions.empty()) + return compareRange(static_cast(other.getBlobAddr()), + static_cast(getBlobAddr()), getSize(), + getMetadata()); + + uint64_t CurrentOffset = 0; + for (const auto &Region : Regions) { + if (!compareRange(static_cast(other.getBlobAddr()) + + CurrentOffset, + static_cast(getBlobAddr()) + + CurrentOffset, + Region.Offset - CurrentOffset, getMetadata())) { + return false; + } + + if (!compareRange(static_cast(other.getBlobAddr()) + + Region.Offset, + static_cast(getBlobAddr()) + + Region.Offset, + Region.Extent, Region.MD)) { + return false; + } + + CurrentOffset = Region.endOffset(); + } + + return compareRange(static_cast(other.getBlobAddr()) + + CurrentOffset, + static_cast(getBlobAddr()) + + CurrentOffset, + getSize() - CurrentOffset, getMetadata()); } bool operator!=(const MnemeMemoryBlob &other) const { return !(*this == other); diff --git a/include/mneme/MnemeNullRecorder.hpp b/include/mneme/MnemeNullRecorder.hpp index c33bf001..68187744 100644 --- a/include/mneme/MnemeNullRecorder.hpp +++ b/include/mneme/MnemeNullRecorder.hpp @@ -16,6 +16,9 @@ class NullRecorder final : public RecorderBackend { public: bool setMetadataForPointer(const void *, Metadata) override { return false; } + bool setMetadataForRegion(const void *, size_t, Metadata) override { + return false; + } bool getMetadataForPointer(const void *, Metadata &) const override { return false; } diff --git a/include/mneme/MnemeRecord.hpp b/include/mneme/MnemeRecord.hpp index 4c52a8a8..561035c8 100644 --- a/include/mneme/MnemeRecord.hpp +++ b/include/mneme/MnemeRecord.hpp @@ -51,6 +51,10 @@ template class MnemeRecorder { return Backend->setMetadataForPointer(ptr, std::move(md)); } + bool setMetadataForRegion(const void *ptr, size_t bytes, Metadata md) { + return Backend->setMetadataForRegion(ptr, bytes, std::move(md)); + } + bool getMetadataForPointer(const void *ptr, Metadata &md) const { return Backend->getMetadataForPointer(ptr, md); } diff --git a/include/mneme/MnemeRecorderBackend.hpp b/include/mneme/MnemeRecorderBackend.hpp index 4c0be652..19276a94 100644 --- a/include/mneme/MnemeRecorderBackend.hpp +++ b/include/mneme/MnemeRecorderBackend.hpp @@ -80,6 +80,8 @@ template class RecorderBackend { virtual ~RecorderBackend() = default; virtual bool setMetadataForPointer(const void *ptr, Metadata md) = 0; + virtual bool setMetadataForRegion(const void *ptr, size_t bytes, + Metadata md) = 0; virtual bool getMetadataForPointer(const void *ptr, Metadata &md) const = 0; virtual bool eraseMetadataForPointer(const void *ptr) = 0; diff --git a/include/mneme/MnemeRecordingBackend.hpp b/include/mneme/MnemeRecordingBackend.hpp index 77ce94d8..758c88ec 100644 --- a/include/mneme/MnemeRecordingBackend.hpp +++ b/include/mneme/MnemeRecordingBackend.hpp @@ -8,6 +8,8 @@ #include "mneme/MnemeSnapshot.hpp" #include +#include +#include #include #include @@ -27,6 +29,7 @@ class RecordingBackend final : public RecorderBackend { RecorderRuntimeFunctions Runtime; RecordDatabase DB; llvm::DenseMap> AllocatedBlobs; + std::map AllocationIndex; std::unique_ptr> PM; // NOTE: We only keep track of the first time we set the device id. Once we @@ -48,6 +51,49 @@ class RecordingBackend final : public RecorderBackend { DeviceID, (void *)MnemeDeviceRT::getSuggestedAddr()); } + static std::size_t builtinByteWidth(BuiltinDType builtin) { + switch (builtin) { + case BuiltinDType::U8: + case BuiltinDType::I8: + return 1; + case BuiltinDType::U16: + case BuiltinDType::I16: + case BuiltinDType::F16: + return 2; + case BuiltinDType::U32: + case BuiltinDType::I32: + case BuiltinDType::F32: + return 4; + case BuiltinDType::U64: + case BuiltinDType::I64: + case BuiltinDType::F64: + return 8; + } + return 1; + } + + typename llvm::DenseMap>::iterator + findOwningBlob(const void *ptr) { + if (!ptr) + return AllocatedBlobs.end(); + + auto Addr = reinterpret_cast(ptr); + auto It = AllocationIndex.upper_bound(Addr); + if (It == AllocationIndex.begin()) + return AllocatedBlobs.end(); + + --It; + auto BlobIt = AllocatedBlobs.find(It->second); + if (BlobIt == AllocatedBlobs.end()) + return AllocatedBlobs.end(); + + auto Base = reinterpret_cast(BlobIt->first); + auto Size = BlobIt->second.getSize(); + if (Addr < Base || (Addr - Base) >= Size) + return AllocatedBlobs.end(); + return BlobIt; + } + public: bool setMetadataForPointer(const void *ptr, Metadata md) override { if (!ptr) @@ -61,6 +107,34 @@ class RecordingBackend final : public RecorderBackend { return true; } + bool setMetadataForRegion(const void *ptr, size_t bytes, Metadata md) override { + if (!ptr) + return false; + + if (bytes == 0) + return false; + + auto It = findOwningBlob(ptr); + if (It == AllocatedBlobs.end()) + return false; + + auto Base = reinterpret_cast(It->first); + auto Addr = reinterpret_cast(ptr); + auto Offset = static_cast(Addr - Base); + auto BlobSize = It->second.getSize(); + if (bytes > BlobSize || Offset > BlobSize - bytes) + return false; + + auto ElemSize = builtinByteWidth(md.builtin); + if ((bytes % ElemSize) != 0) + return false; + + if (!It->second.setRegionMetadata(Offset, bytes, md)) + return false; + + return true; + } + bool getMetadataForPointer(const void *ptr, Metadata &md) const override { if (!ptr) return false; @@ -94,6 +168,7 @@ class RecordingBackend final : public RecorderBackend { auto ret = MemBlob.map(reinterpret_cast(Addr), ReservedSize, size); *ptr = MemBlob.ptr(); AllocatedBlobs.insert({*ptr, std::move(MemBlob)}); + AllocationIndex[reinterpret_cast(*ptr)] = *ptr; LOG_DEBUG("Intercepted Device Malloc PTR:{} SIZE:{} ACTUALSIZE:{}", *ptr, size, ReservedSize); return ret; @@ -126,6 +201,7 @@ class RecordingBackend final : public RecorderBackend { } PM->releaseAddr(AllocatedBlobs[ptr].getActualSize(), ptr); auto ret = AllocatedBlobs[ptr].release(); + AllocationIndex.erase(reinterpret_cast(ptr)); LOG_DEBUG("Intercepted device Free PTR:{} SIZE:{} ACTUALSIZE:{}", ptr, AllocatedBlobs[ptr].getSize(), AllocatedBlobs[ptr].getActualSize()); diff --git a/include/mneme/MnemeSnapshot.hpp b/include/mneme/MnemeSnapshot.hpp index 5aa8cc27..bc631aaa 100644 --- a/include/mneme/MnemeSnapshot.hpp +++ b/include/mneme/MnemeSnapshot.hpp @@ -131,6 +131,9 @@ template class MnemeSnapshot { using KernelFunction_t = typename MnemeDeviceRT::KernelFunction_t; static constexpr const char DiffMagic[] = "MNEME_DIFF_V1"; static constexpr size_t DiffMagicSize = sizeof(DiffMagic) - 1; + static constexpr const char RegionTrailerMagic[] = "MNEME_REGION_V1"; + static constexpr size_t RegionTrailerMagicSize = + sizeof(RegionTrailerMagic) - 1; static constexpr size_t DiffChunkSize = 1 << 20; class CountingRawOStream : public llvm::raw_ostream { @@ -281,12 +284,98 @@ template class MnemeSnapshot { } } + static void writeRegionMetadataTrailer( + llvm::raw_ostream &OS, + const llvm::DenseMap> &DeviceMemory) { + size_t TotalBlobsWithRegions = 0; + for (const auto &[Ptr, Blob] : DeviceMemory) { + if (Blob.hasRegionMetadata()) + ++TotalBlobsWithRegions; + } + // The "trailer" is an optional block of bytes appended after the normal + // snapshot payload (globals, blobs, kernel args). Keep older snapshots + // valid by omitting it entirely when no blobs carry subregion metadata. + if (TotalBlobsWithRegions == 0) + return; + + // Key trailer entries by blob device address so they can be attached after + // the main blob records have already been deserialized. + util::writeBytes(OS, + llvm::StringRef(RegionTrailerMagic, RegionTrailerMagicSize)); + util::writeScalar(OS, TotalBlobsWithRegions); + for (const auto &[Ptr, Blob] : DeviceMemory) { + if (!Blob.hasRegionMetadata()) + continue; + auto *BlobAddr = Blob.getBlobAddr(); + util::writeScalar(OS, BlobAddr); + + const auto &Regions = Blob.getRegionMetadata(); + size_t NumRegions = Regions.size(); + util::writeScalar(OS, NumRegions); + for (const auto &Region : Regions) { + util::writeScalar(OS, Region.Offset); + util::writeScalar(OS, Region.Extent); + metadata::serialize(OS, Region.MD); + } + } + } + + static void readRegionMetadataTrailer( + const std::string &Filename, const char *&CurrentPtr, + const char *BufferEnd, + llvm::DenseMap> &DeviceMemory) { + // No trailing bytes means this snapshot predates region metadata support. + if (CurrentPtr == BufferEnd) + return; + + // If trailing bytes exist, they must be the region trailer. Anything else + // indicates a malformed or incompatible snapshot payload. + if (static_cast(BufferEnd - CurrentPtr) < RegionTrailerMagicSize || + llvm::StringRef(CurrentPtr, RegionTrailerMagicSize) != + llvm::StringRef(RegionTrailerMagic, RegionTrailerMagicSize)) { + LOG_FATAL("Unexpected trailing bytes in Mneme snapshot " + Filename); + } + CurrentPtr += RegionTrailerMagicSize; + + size_t TotalBlobs = util::extractScalar(CurrentPtr); + for (size_t I = 0; I < TotalBlobs; ++I) { + void *DeviceAddr = util::extractScalar(CurrentPtr); + size_t NumRegions = util::extractScalar(CurrentPtr); + + auto It = DeviceMemory.find(DeviceAddr); + if (It == DeviceMemory.end()) + LOG_FATAL("Mneme region trailer references device allocation missing " + "from snapshot"); + + // Rebuild the full region table for this blob, then install it as one + // validated replacement so ordering/overlap checks happen once. + std::vector Regions; + Regions.reserve(NumRegions); + for (size_t R = 0; R < NumRegions; ++R) { + MemoryRegionMetadata Region; + Region.Offset = util::extractScalar(CurrentPtr); + Region.Extent = util::extractScalar(CurrentPtr); + Region.MD = metadata::fromBuffer(CurrentPtr); + Regions.push_back(std::move(Region)); + } + + if (!It->second.replaceRegionMetadata(std::move(Regions))) + LOG_FATAL("Mneme region trailer contains invalid or overlapping " + "regions for device allocation"); + } + + if (CurrentPtr != BufferEnd) + LOG_FATAL("Unexpected extra bytes after Mneme region trailer in " + + Filename); + } + static void readFullMnemeSnapShot( llvm::MemoryBuffer *Buffer, std::unordered_map &GlobalVars, llvm::DenseMap> &DeviceMemory, std::shared_ptr KInfo) { auto *Start = Buffer->getBufferStart(); + auto *BufferEnd = Buffer->getBufferEnd(); auto *CurrentPtr = Start; size_t TotalGlobals = util::extractScalar(CurrentPtr); LOG_DEBUG("Snapshot contains {} Globals at location {}", TotalGlobals, @@ -315,6 +404,9 @@ template class MnemeSnapshot { KInfo->KernelArgSizes[A] = util::extractScalar(CurrentPtr); KInfo->setArgData(CurrentPtr, A); } + + readRegionMetadataTrailer(Buffer->getBufferIdentifier().str(), CurrentPtr, + BufferEnd, DeviceMemory); } // Applies the diff ranges from DiffBuffer onto the already-loaded base @@ -377,6 +469,9 @@ template class MnemeSnapshot { Blob.getSize()), NumRanges); } + + readRegionMetadataTrailer(Filename, CurrentPtr, DiffBuffer->getBufferEnd(), + DeviceMemory); } static size_t getSerializedMetadataSize(const Metadata &MD) { @@ -385,6 +480,29 @@ template class MnemeSnapshot { (MD.tag ? MD.tag->size() : 0); } + static size_t + getSerializedRegionMetadataSize(const MemoryRegionMetadata &Region) { + return sizeof(Region.Offset) + sizeof(Region.Extent) + + getSerializedMetadataSize(Region.MD); + } + + static size_t computeRegionMetadataTrailerSize( + const llvm::DenseMap> &DeviceMemory) { + size_t TotalBlobsWithRegions = 0; + size_t Size = 0; + for (const auto &[Ptr, Blob] : DeviceMemory) { + if (!Blob.hasRegionMetadata()) + continue; + if (TotalBlobsWithRegions++ == 0) + Size = RegionTrailerMagicSize + sizeof(size_t); + Size += sizeof(void *); + Size += sizeof(size_t); + for (const auto &Region : Blob.getRegionMetadata()) + Size += getSerializedRegionMetadataSize(Region); + } + return Size; + } + static size_t computeMnemeBytesSnapshotSize( const proteus::runtime::GlobalMetadataMap &GlobalVars, const llvm::DenseMap> &DeviceMemory, @@ -412,6 +530,7 @@ template class MnemeSnapshot { Size += sizeof(size_t); Size += ArgSize; } + Size += computeRegionMetadataTrailerSize(DeviceMemory); return Size; } @@ -508,6 +627,8 @@ template class MnemeSnapshot { KernelArgSizes[I]); } + writeRegionMetadataTrailer(OutBC, DeviceMemory); + return Filename.filename(); } @@ -557,6 +678,8 @@ template class MnemeSnapshot { writeCountAndWriteChangedRanges(OutBC, Blob, UpdateBaseData); } + + writeRegionMetadataTrailer(OutBC, DeviceMemory); } static size_t measureMnemeDiffSnapshotSize( diff --git a/src/MnemeAnnotation.cpp b/src/MnemeAnnotation.cpp index 68081c32..1bdfeb79 100644 --- a/src/MnemeAnnotation.cpp +++ b/src/MnemeAnnotation.cpp @@ -6,8 +6,8 @@ void annotate(const void *ptr, Metadata md) { detail::annotate_impl(ptr, std::move(md)); } -void annotate(void *ptr, Metadata md) { - detail::annotate_impl(ptr, std::move(md)); +void annotate(const void *ptr, std::size_t bytes, Metadata md) { + detail::annotate_region_impl(ptr, bytes, std::move(md)); } namespace detail { @@ -16,8 +16,22 @@ void annotate_impl(const void *ptr, Metadata md) { if (!ptr) return; - if (mneme_set_metadata_for_ptr) + if (mneme_set_metadata_for_ptr) { mneme_set_metadata_for_ptr(ptr, std::move(md)); + return; + } +} + +void annotate_region_impl(const void *ptr, std::size_t bytes, Metadata md) { + if (!ptr) + return; + + if (bytes == 0) + return; + + if (mneme_set_metadata_for_region) { + mneme_set_metadata_for_region(ptr, bytes, std::move(md)); + } } bool get_annotation(const void *ptr, Metadata &md) { diff --git a/src/MnemeAnnotationRuntime.hpp b/src/MnemeAnnotationRuntime.hpp index 17abd26e..39a2eb51 100644 --- a/src/MnemeAnnotationRuntime.hpp +++ b/src/MnemeAnnotationRuntime.hpp @@ -2,9 +2,14 @@ #include "mneme/MnemeAnnotation.hpp" +#include + extern "C" { bool mneme_set_metadata_for_ptr(const void *ptr, mneme::Metadata md) __attribute__((weak)); +bool mneme_set_metadata_for_region(const void *ptr, std::size_t bytes, + mneme::Metadata md) + __attribute__((weak)); bool mneme_get_metadata_for_ptr(const void *ptr, mneme::Metadata *md) __attribute__((weak)); bool mneme_erase_metadata_for_ptr(const void *ptr) __attribute__((weak)); @@ -14,8 +19,9 @@ namespace mneme { namespace detail { void annotate_impl(const void *ptr, Metadata md); +void annotate_region_impl(const void *ptr, std::size_t bytes, Metadata md); bool get_annotation(const void *ptr, Metadata &md); void erase_annotation(const void *ptr); } // namespace detail -} // namespace mneme \ No newline at end of file +} // namespace mneme diff --git a/src/MnemeRecordPreloadCUDA.cpp b/src/MnemeRecordPreloadCUDA.cpp index ec0e92ec..3e262ca9 100644 --- a/src/MnemeRecordPreloadCUDA.cpp +++ b/src/MnemeRecordPreloadCUDA.cpp @@ -94,6 +94,12 @@ bool mneme_set_metadata_for_ptr(const void *ptr, mneme::Metadata md) { return mneme.setMetadataForPointer(ptr, std::move(md)); } +bool mneme_set_metadata_for_region(const void *ptr, size_t bytes, + mneme::Metadata md) { + auto &mneme = MnemeRecorderCUDAPreload::instance(); + return mneme.setMetadataForRegion(ptr, bytes, std::move(md)); +} + bool mneme_get_metadata_for_ptr(const void *ptr, mneme::Metadata *md) { if (!md) return false; diff --git a/src/MnemeRecordPreloadHIP.cpp b/src/MnemeRecordPreloadHIP.cpp index 993c3afe..868be71f 100644 --- a/src/MnemeRecordPreloadHIP.cpp +++ b/src/MnemeRecordPreloadHIP.cpp @@ -86,6 +86,12 @@ bool mneme_set_metadata_for_ptr(const void *ptr, mneme::Metadata md) { return mneme.setMetadataForPointer(ptr, std::move(md)); } +bool mneme_set_metadata_for_region(const void *ptr, size_t bytes, + mneme::Metadata md) { + auto &mneme = MnemeRecorderHIPPreload::instance(); + return mneme.setMetadataForRegion(ptr, bytes, std::move(md)); +} + bool mneme_get_metadata_for_ptr(const void *ptr, mneme::Metadata *md) { if (!md) return false; From 4a1ceb114e1109566e629f179a7a637abb5a7132 Mon Sep 17 00:00:00 2001 From: Caetano Melone Date: Tue, 18 Aug 2026 11:24:50 -0700 Subject: [PATCH 2/5] add annotation and replay logging Signed-off-by: Caetano Melone --- include/mneme/MnemeMemory.hpp | 35 ++++++++++- include/mneme/MnemeRecordingBackend.hpp | 83 +++++++++++++++++++++---- python/mneme/commands.py | 5 ++ src/MnemeAnnotation.cpp | 34 +++++++++- 4 files changed, 141 insertions(+), 16 deletions(-) diff --git a/include/mneme/MnemeMemory.hpp b/include/mneme/MnemeMemory.hpp index 6171472c..7b302dd2 100644 --- a/include/mneme/MnemeMemory.hpp +++ b/include/mneme/MnemeMemory.hpp @@ -302,13 +302,30 @@ template class MnemeMemoryBlob { } } - if (Regions.empty()) + if (Regions.empty()) { + LOG_DEBUG( + "Comparing blob {} full range [0, {}) with whole-blob metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", + getBlobAddr(), getSize(), + static_cast(getMetadata().builtin), getMetadata().threshold, + static_cast(getMetadata().threshold_kind), + static_cast(getMetadata().norm), + getMetadata().tag.value_or("no_tag")); return compareRange(static_cast(other.getBlobAddr()), static_cast(getBlobAddr()), getSize(), getMetadata()); + } uint64_t CurrentOffset = 0; for (const auto &Region : Regions) { + if (Region.Offset > CurrentOffset) { + LOG_DEBUG( + "Comparing blob {} gap [{}, {}) with whole-blob metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", + getBlobAddr(), CurrentOffset, Region.Offset, + static_cast(getMetadata().builtin), getMetadata().threshold, + static_cast(getMetadata().threshold_kind), + static_cast(getMetadata().norm), + getMetadata().tag.value_or("no_tag")); + } if (!compareRange(static_cast(other.getBlobAddr()) + CurrentOffset, static_cast(getBlobAddr()) + @@ -317,6 +334,13 @@ template class MnemeMemoryBlob { return false; } + LOG_DEBUG( + "Comparing blob {} region [{}, {}) with region metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", + getBlobAddr(), Region.Offset, Region.endOffset(), + static_cast(Region.MD.builtin), Region.MD.threshold, + static_cast(Region.MD.threshold_kind), + static_cast(Region.MD.norm), + Region.MD.tag.value_or("no_tag")); if (!compareRange(static_cast(other.getBlobAddr()) + Region.Offset, static_cast(getBlobAddr()) + @@ -328,6 +352,15 @@ template class MnemeMemoryBlob { CurrentOffset = Region.endOffset(); } + if (CurrentOffset < getSize()) { + LOG_DEBUG( + "Comparing blob {} tail [{}, {}) with whole-blob metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", + getBlobAddr(), CurrentOffset, getSize(), + static_cast(getMetadata().builtin), getMetadata().threshold, + static_cast(getMetadata().threshold_kind), + static_cast(getMetadata().norm), + getMetadata().tag.value_or("no_tag")); + } return compareRange(static_cast(other.getBlobAddr()) + CurrentOffset, static_cast(getBlobAddr()) + diff --git a/include/mneme/MnemeRecordingBackend.hpp b/include/mneme/MnemeRecordingBackend.hpp index 758c88ec..fb69f2bf 100644 --- a/include/mneme/MnemeRecordingBackend.hpp +++ b/include/mneme/MnemeRecordingBackend.hpp @@ -96,65 +96,124 @@ class RecordingBackend final : public RecorderBackend { public: bool setMetadataForPointer(const void *ptr, Metadata md) override { - if (!ptr) + if (!ptr) { + LOG_WARN("setMetadataForPointer called with null pointer"); return false; + } auto It = AllocatedBlobs.find(const_cast(ptr)); - if (It == AllocatedBlobs.end()) + if (It == AllocatedBlobs.end()) { + LOG_DEBUG("exact-match miss ptr={} tag={} allocated_blobs={}", ptr, + md.tag.value_or("no_tag"), AllocatedBlobs.size()); + LOG_WARN( + "annotation dropped after exact-match miss ptr={} tag={} allocated_blobs={}", + ptr, md.tag.value_or("no_tag"), AllocatedBlobs.size()); return false; + } + LOG_INFO( + "exact-match hit ptr={} size={} actual_size={} builtin={} threshold={} threshold_kind={} norm={} tag={}", + ptr, It->second.getSize(), It->second.getActualSize(), + static_cast(md.builtin), md.threshold, + static_cast(md.threshold_kind), static_cast(md.norm), + md.tag.value_or("no_tag")); It->second.setMetadata(std::move(md)); return true; } bool setMetadataForRegion(const void *ptr, size_t bytes, Metadata md) override { - if (!ptr) + if (!ptr) { + LOG_WARN("setMetadataForRegion called with null pointer"); return false; + } - if (bytes == 0) + if (bytes == 0) { + LOG_WARN("setMetadataForRegion called with zero bytes ptr={} tag={}", + ptr, md.tag.value_or("no_tag")); return false; + } auto It = findOwningBlob(ptr); - if (It == AllocatedBlobs.end()) + if (It == AllocatedBlobs.end()) { + LOG_WARN( + "region annotation could not resolve owning allocation ptr={} bytes={} tag={} allocated_blobs={}", + ptr, bytes, md.tag.value_or("no_tag"), AllocatedBlobs.size()); return false; + } auto Base = reinterpret_cast(It->first); auto Addr = reinterpret_cast(ptr); auto Offset = static_cast(Addr - Base); auto BlobSize = It->second.getSize(); - if (bytes > BlobSize || Offset > BlobSize - bytes) + if (bytes > BlobSize || Offset > BlobSize - bytes) { + LOG_WARN( + "region annotation exceeds owning allocation ptr={} base_ptr={} offset={} bytes={} blob_size={} tag={}", + ptr, It->first, Offset, bytes, BlobSize, md.tag.value_or("no_tag")); return false; + } auto ElemSize = builtinByteWidth(md.builtin); - if ((bytes % ElemSize) != 0) + if ((bytes % ElemSize) != 0) { + LOG_WARN( + "region annotation byte extent is not aligned to builtin type ptr={} bytes={} builtin={} elem_size={} tag={}", + ptr, bytes, static_cast(md.builtin), ElemSize, + md.tag.value_or("no_tag")); return false; + } - if (!It->second.setRegionMetadata(Offset, bytes, md)) + if (!It->second.setRegionMetadata(Offset, bytes, md)) { + LOG_WARN( + "region annotation rejected due to overlap ptr={} base_ptr={} offset={} bytes={} tag={}", + ptr, It->first, Offset, bytes, md.tag.value_or("no_tag")); return false; + } + LOG_INFO( + "region annotation registered ptr={} base_ptr={} offset={} bytes={} builtin={} threshold={} threshold_kind={} norm={} tag={}", + ptr, It->first, Offset, bytes, static_cast(md.builtin), + md.threshold, static_cast(md.threshold_kind), + static_cast(md.norm), md.tag.value_or("no_tag")); return true; } bool getMetadataForPointer(const void *ptr, Metadata &md) const override { - if (!ptr) + if (!ptr) { + LOG_WARN("getMetadataForPointer called with null pointer"); return false; + } auto It = AllocatedBlobs.find(const_cast(ptr)); - if (It == AllocatedBlobs.end()) + if (It == AllocatedBlobs.end()) { + LOG_DEBUG("getMetadataForPointer exact-match miss ptr={}", ptr); + LOG_DEBUG( + "getMetadataForPointer returning false after exact-match miss ptr={}", + ptr); return false; + } md = It->second.getMetadata(); + LOG_DEBUG("getMetadataForPointer hit ptr={} tag={}", ptr, + md.tag.value_or("no_tag")); return true; } bool eraseMetadataForPointer(const void *ptr) override { - if (!ptr) + if (!ptr) { + LOG_WARN("eraseMetadataForPointer called with null pointer"); return false; + } auto It = AllocatedBlobs.find(const_cast(ptr)); - if (It == AllocatedBlobs.end()) + if (It == AllocatedBlobs.end()) { + LOG_DEBUG("eraseMetadataForPointer exact-match miss ptr={}", ptr); + LOG_DEBUG( + "eraseMetadataForPointer returning false after exact-match miss ptr={}", + ptr); return false; + } + LOG_INFO("eraseMetadataForPointer hit ptr={} tag={}", ptr, + It->second.getMetadata().tag.value_or("no_tag")); It->second.setMetadata(Metadata{}); return true; } diff --git a/python/mneme/commands.py b/python/mneme/commands.py index 76c35e3a..6ede6e1b 100644 --- a/python/mneme/commands.py +++ b/python/mneme/commands.py @@ -682,6 +682,11 @@ def run(args, verbosity): kwargs = vars(args) kwargs.pop("command") kwargs.pop("func") + + if verbosity is not None: + logger.debug(f"MNEME_LOG_LEVEL={verbosity}") + os.environ["MNEME_LOG_LEVEL"] = verbosity + executor = Replay(**kwargs) # We currently link all LLVM IR modules together diff --git a/src/MnemeAnnotation.cpp b/src/MnemeAnnotation.cpp index 1bdfeb79..db8df557 100644 --- a/src/MnemeAnnotation.cpp +++ b/src/MnemeAnnotation.cpp @@ -1,4 +1,5 @@ #include "MnemeAnnotationRuntime.hpp" +#include "mneme/MnemeLogger.hpp" namespace mneme { @@ -13,25 +14,52 @@ void annotate(const void *ptr, std::size_t bytes, Metadata md) { namespace detail { void annotate_impl(const void *ptr, Metadata md) { - if (!ptr) + if (!ptr) { + LOG_WARN("annotate called with null pointer"); return; + } + + LOG_DEBUG( + "request ptr={} builtin={} threshold={} threshold_kind={} norm={} tag={}", + ptr, static_cast(md.builtin), md.threshold, + static_cast(md.threshold_kind), static_cast(md.norm), + md.tag.value_or("no_tag")); if (mneme_set_metadata_for_ptr) { mneme_set_metadata_for_ptr(ptr, std::move(md)); return; } + + LOG_WARN("no recorder hook installed for ptr={} tag={}; annotation ignored", + ptr, md.tag.value_or("no_tag")); } void annotate_region_impl(const void *ptr, std::size_t bytes, Metadata md) { - if (!ptr) + if (!ptr) { + LOG_WARN("region annotate called with null pointer"); return; + } - if (bytes == 0) + if (bytes == 0) { + LOG_WARN("region annotate called with zero bytes ptr={} tag={}", ptr, + md.tag.value_or("no_tag")); return; + } + + LOG_DEBUG( + "region request ptr={} bytes={} builtin={} threshold={} threshold_kind={} norm={} tag={}", + ptr, bytes, static_cast(md.builtin), md.threshold, + static_cast(md.threshold_kind), static_cast(md.norm), + md.tag.value_or("no_tag")); if (mneme_set_metadata_for_region) { mneme_set_metadata_for_region(ptr, bytes, std::move(md)); + return; } + + LOG_WARN( + "no recorder region hook installed for ptr={} bytes={} tag={}; annotation ignored", + ptr, bytes, md.tag.value_or("no_tag")); } bool get_annotation(const void *ptr, Metadata &md) { From 4f9a7ecbf3e0f702b0102eb57a889fefe011f410 Mon Sep 17 00:00:00 2001 From: Caetano Melone Date: Wed, 2 Sep 2026 11:03:47 -0700 Subject: [PATCH 3/5] Refactor annotation storage around unified byte range type This is a breaking change to existing recordings. It replaces the split entire-blob/region metadata types with one internal annotation list from ByteSpan, Metadata pairs. It treats the whole-blob annotation as the owning blob and keeps the region annotations as overrides. I also updated the snapshot serialization to r/w the annotation list directly, and switched the snapshots to MNEME_BYTES_V2/MNEME_DIFF_V2. I also removed the compatibility layers the original commits had. Signed-off-by: Caetano Melone --- include/mneme/MnemeMemory.hpp | 304 +++++++++++++++--------- include/mneme/MnemeRecordingBackend.hpp | 20 +- include/mneme/MnemeSnapshot.hpp | 167 +++---------- 3 files changed, 243 insertions(+), 248 deletions(-) diff --git a/include/mneme/MnemeMemory.hpp b/include/mneme/MnemeMemory.hpp index 7b302dd2..2f434845 100644 --- a/include/mneme/MnemeMemory.hpp +++ b/include/mneme/MnemeMemory.hpp @@ -14,15 +14,14 @@ #include "mneme/MnemeAnnotation.hpp" #include "mneme/MnemeAnnotationInternal.hpp" #include "mneme/MnemeComparators.hpp" +#include "mneme/MnemeLLVMUtils.hpp" #include "mneme/MnemeLogger.hpp" #include "mneme/MnemeUtils.hpp" namespace mneme { -// Metadata attached to a byte subrange inside one owning blob. -struct MemoryRegionMetadata { +struct ByteSpan { uint64_t Offset = 0; uint64_t Extent = 0; - Metadata MD; uint64_t endOffset() const { return Offset + Extent; } @@ -34,23 +33,65 @@ struct MemoryRegionMetadata { return Offset == OtherOffset && Extent == OtherExtent; } + bool hasSameRange(const ByteSpan &Other) const { + return hasSameRange(Other.Offset, Other.Extent); + } + bool overlaps(uint64_t OtherOffset, uint64_t OtherExtent) const { - return Offset < OtherOffset + OtherExtent && - OtherOffset < Offset + Extent; + return Offset < OtherOffset + OtherExtent && OtherOffset < endOffset(); + } + + bool overlaps(const ByteSpan &Other) const { + return overlaps(Other.Offset, Other.Extent); } }; -inline bool operator==(const MemoryRegionMetadata &LHS, - const MemoryRegionMetadata &RHS) { - return LHS.Offset == RHS.Offset && LHS.Extent == RHS.Extent && - LHS.MD == RHS.MD; +inline bool operator==(const ByteSpan &LHS, const ByteSpan &RHS) { + return LHS.Offset == RHS.Offset && LHS.Extent == RHS.Extent; } -inline bool operator!=(const MemoryRegionMetadata &LHS, - const MemoryRegionMetadata &RHS) { +inline bool operator!=(const ByteSpan &LHS, const ByteSpan &RHS) { return !(LHS == RHS); } +// Metadata attached to a byte range inside one owning blob. +struct MemoryAnnotation { + ByteSpan Range; + Metadata MD; +}; + +inline bool operator==(const MemoryAnnotation &LHS, const MemoryAnnotation &RHS) { + return LHS.Range == RHS.Range && LHS.MD == RHS.MD; +} + +inline bool operator!=(const MemoryAnnotation &LHS, const MemoryAnnotation &RHS) { + return !(LHS == RHS); +} + +inline std::vector +readAnnotationsFromBuffer(const char *&Buffer, size_t NumAnnotations) { + std::vector LoadedAnnotations; + LoadedAnnotations.reserve(NumAnnotations); + for (size_t I = 0; I < NumAnnotations; ++I) { + MemoryAnnotation Annotation; + Annotation.Range.Offset = util::extractScalar(Buffer); + Annotation.Range.Extent = util::extractScalar(Buffer); + Annotation.MD = metadata::fromBuffer(Buffer); + LoadedAnnotations.push_back(std::move(Annotation)); + } + return LoadedAnnotations; +} + +inline void writeAnnotationsToStream( + llvm::raw_ostream &OS, const std::vector &Annotations) { + util::writeScalar(OS, Annotations.size()); + for (const auto &Annotation : Annotations) { + util::writeScalar(OS, Annotation.Range.Offset); + util::writeScalar(OS, Annotation.Range.Extent); + metadata::serialize(OS, Annotation.MD); + } +} + template class MnemeMemoryBlob { public: using MnemeDeviceRT = DeviceTraits; @@ -77,30 +118,88 @@ template class MnemeMemoryBlob { return Compare.Agg <= Md.threshold; } - Metadata PtrMD; - std::vector RegionMD; + std::vector Annotations; uint64_t ActualSize; void *BlobAddr; uint64_t Size; std::unique_ptr HostData; bool IsMapped; - bool isValidRegionRange(uint64_t Offset, uint64_t Extent) const { - return Extent != 0 && Offset <= Size && Extent <= Size - Offset; + bool isValidAnnotationRange(const ByteSpan &Range) const { + if (Range.Offset > Size) + return false; + if (Range.Extent == 0) + return Size == 0 && Range.Offset == 0; + return Range.Extent <= Size - Range.Offset; + } + + bool isWholeBlobRange(const ByteSpan &Range) const { + return Range.Offset == 0 && Range.Extent == Size; } - static bool regionSortLess(const MemoryRegionMetadata &LHS, - const MemoryRegionMetadata &RHS) { - if (LHS.Offset != RHS.Offset) - return LHS.Offset < RHS.Offset; - return LHS.Extent < RHS.Extent; + static bool annotationSortLess(const MemoryAnnotation &LHS, + const MemoryAnnotation &RHS) { + if (LHS.Range.Offset != RHS.Range.Offset) + return LHS.Range.Offset < RHS.Range.Offset; + return LHS.Range.Extent < RHS.Range.Extent; + } + + void resetDefaultAnnotation() { Annotations = {{{0, Size}, {}}}; } + + const MemoryAnnotation *findWholeBlobAnnotation() const { + for (const auto &Annotation : Annotations) { + if (isWholeBlobRange(Annotation.Range)) + return &Annotation; + } + return nullptr; + } + + static std::vector + regionAnnotations(const std::vector &Annotations, + uint64_t BlobSize) { + std::vector Regions; + Regions.reserve(Annotations.size()); + for (const auto &Annotation : Annotations) { + if (Annotation.Range.Offset == 0 && Annotation.Range.Extent == BlobSize) + continue; + Regions.push_back(Annotation); + } + return Regions; } public: + bool registerAnnotation(ByteSpan Range, Metadata Md) { + if (!isValidAnnotationRange(Range)) + return false; + + for (auto &Annotation : Annotations) { + if (Annotation.Range.hasSameRange(Range)) { + Annotation.MD = std::move(Md); + return true; + } + + if (!Annotation.Range.overlaps(Range)) + continue; + + // The full-range annotation is the blob's default policy, so narrower + // region annotations may overlap it. Overlaps between two region + // annotations remain invalid. + if (isWholeBlobRange(Annotation.Range) || isWholeBlobRange(Range)) + continue; + + return false; + } + + Annotations.push_back({Range, std::move(Md)}); + std::sort(Annotations.begin(), Annotations.end(), annotationSortLess); + return true; + } MnemeMemoryBlob(uint64_t ActualSize = 0, void *BlobAddr = nullptr, uint64_t Size = 0) : ActualSize(ActualSize), BlobAddr(BlobAddr), Size(Size), - HostData(new uint8_t[Size]), IsMapped(false) {} + HostData(new uint8_t[Size]), IsMapped(false) { + resetDefaultAnnotation(); + } DeviceError_t map(void *VA, uint64_t ActualSize, uint64_t Size) { this->Size = Size; @@ -109,6 +208,8 @@ template class MnemeMemoryBlob { this->BlobAddr = VA; this->IsMapped = true; this->ActualSize = ActualSize; + // Preserve any annotations already loaded from a snapshot. Fresh blobs get + // their default whole-blob annotation from the constructor. return MnemeDeviceRT::DeviceSuccess; }; @@ -120,6 +221,8 @@ template class MnemeMemoryBlob { this->ActualSize = Size; this->Size = Size; this->IsMapped = false; + // Preserve any annotations already loaded from a snapshot. Fresh blobs get + // their default whole-blob annotation from the constructor. return ret; } @@ -157,7 +260,10 @@ template class MnemeMemoryBlob { Buffer += Size; LOG_DEBUG("Read memory blob at address {} SIZE: {} ActualSize:{}", DeviceAddr, Size, ActualSize); - Blob.PtrMD = metadata::fromBuffer(Buffer); + size_t NumAnnotations = util::extractScalar(Buffer); + auto LoadedAnnotations = readAnnotationsFromBuffer(Buffer, NumAnnotations); + if (!Blob.replaceAnnotations(std::move(LoadedAnnotations))) + LOG_FATAL("Invalid annotation set serialized for blob"); return std::make_pair(DeviceAddr, std::move(Blob)); } @@ -173,8 +279,7 @@ template class MnemeMemoryBlob { ActualSize = other.ActualSize; HostData = std::move(other.HostData); IsMapped = other.IsMapped; - PtrMD = other.PtrMD; - RegionMD = std::move(other.RegionMD); + Annotations = std::move(other.Annotations); other.BlobAddr = 0; other.HostData = nullptr; } @@ -184,8 +289,7 @@ template class MnemeMemoryBlob { MnemeMemoryBlob(MnemeMemoryBlob &&other) noexcept : BlobAddr(other.BlobAddr), Size(other.Size), ActualSize(other.ActualSize), HostData(std::move(other.HostData)), - IsMapped(other.IsMapped), PtrMD(other.PtrMD), - RegionMD(std::move(other.RegionMD)) { + IsMapped(other.IsMapped), Annotations(std::move(other.Annotations)) { other.BlobAddr = 0; other.HostData = nullptr; } @@ -199,76 +303,52 @@ template class MnemeMemoryBlob { uint64_t getSize() const { return Size; } const std::unique_ptr &getHostData() const { return HostData; } - void setMetadata(Metadata Md) { PtrMD = Md; } - - Metadata getMetadata() const { return PtrMD; } - - // Distinct overlapping regions are rejected. Re-registering the exact same - // [offset, extent) updates the existing metadata in place. - bool setRegionMetadata(uint64_t Offset, uint64_t Extent, Metadata Md) { - if (!isValidRegionRange(Offset, Extent)) + bool replaceAnnotations(std::vector CandidateAnnotations) { + if (CandidateAnnotations.empty()) return false; - MemoryRegionMetadata Key{Offset, Extent, {}}; - auto It = std::lower_bound(RegionMD.begin(), RegionMD.end(), Key, - regionSortLess); - - if (It != RegionMD.end() && It->hasSameRange(Offset, Extent)) { - It->MD = std::move(Md); - return true; - } - - if (It != RegionMD.begin()) { - auto Prev = std::prev(It); - if (Prev->overlaps(Offset, Extent)) + // Install a complete blob annotation set, typically after deserializing a + // snapshot. The set must contain exactly one full-range default annotation. + size_t WholeBlobCount = 0; + for (const auto &Annotation : CandidateAnnotations) { + if (!isValidAnnotationRange(Annotation.Range)) return false; + if (isWholeBlobRange(Annotation.Range)) + ++WholeBlobCount; } - if (It != RegionMD.end() && It->overlaps(Offset, Extent)) + if (WholeBlobCount != 1) return false; - RegionMD.insert(It, {Offset, Extent, std::move(Md)}); - return true; - } - - // Replace the entire region table after validating a bulk-loaded set, such - // as snapshot deserialization. - bool replaceRegionMetadata(std::vector Regions) { - for (const auto &Region : Regions) { - if (!isValidRegionRange(Region.Offset, Region.Extent)) - return false; - } - - std::sort(Regions.begin(), Regions.end(), regionSortLess); - for (size_t I = 1; I < Regions.size(); ++I) { - if (Regions[I - 1].overlaps(Regions[I].Offset, Regions[I].Extent)) + std::sort(CandidateAnnotations.begin(), CandidateAnnotations.end(), + annotationSortLess); + for (size_t I = 1; I < CandidateAnnotations.size(); ++I) { + if (CandidateAnnotations[I - 1].Range.hasSameRange( + CandidateAnnotations[I].Range)) + continue; + if (CandidateAnnotations[I - 1].Range.overlaps( + CandidateAnnotations[I].Range) && + !isWholeBlobRange(CandidateAnnotations[I - 1].Range) && + !isWholeBlobRange(CandidateAnnotations[I].Range)) { return false; + } } - RegionMD = std::move(Regions); + Annotations = std::move(CandidateAnnotations); return true; } - void clearRegionMetadata() { RegionMD.clear(); } - - bool hasRegionMetadata() const { return !RegionMD.empty(); } + const std::vector &getAnnotations() const { + return Annotations; + } - const std::vector &getRegionMetadata() const { - return RegionMD; + const MemoryAnnotation *getWholeBlobAnnotation() const { + return findWholeBlobAnnotation(); } - // RegionMD is kept sorted by offset, so the owning region for a byte can be - // resolved by finding the last region whose start is <= ByteOffset. - const MemoryRegionMetadata *findRegionMetadata(uint64_t ByteOffset) const { - auto It = std::upper_bound( - RegionMD.begin(), RegionMD.end(), ByteOffset, - [](uint64_t Offset, const MemoryRegionMetadata &Region) { - return Offset < Region.Offset; - }); - if (It == RegionMD.begin()) - return nullptr; - - --It; - return It->contains(ByteOffset) ? &*It : nullptr; + std::vector getRegionAnnotations() const { + // Region annotations are the narrower ranges, excluding the + // full-range default annotation for the blob. + return regionAnnotations(Annotations, Size); } bool operator==(const MnemeMemoryBlob &other) const { @@ -277,15 +357,17 @@ template class MnemeMemoryBlob { return false; } - if (getMetadata() != other.getMetadata()) { + auto *WholeBlob = getWholeBlobAnnotation(); + auto *OtherWholeBlob = other.getWholeBlobAnnotation(); + if (!WholeBlob || !OtherWholeBlob || WholeBlob->MD != OtherWholeBlob->MD) { LOG_WARN("Whole-blob metadata differs for blob {}", getBlobAddr()); return false; } - // Compare the region tables first so the subsequent byte comparison can - // assume both blobs partition the address range the same way. - const auto &Regions = getRegionMetadata(); - const auto &OtherRegions = other.getRegionMetadata(); + // First compare the region annotation structure, then compare bytes using + // the whole-blob annotation as fallback for any uncovered ranges. + auto Regions = getRegionAnnotations(); + auto OtherRegions = other.getRegionAnnotations(); if (Regions.size() != OtherRegions.size()) { LOG_WARN("Region metadata count differs for blob {}", getBlobAddr()); @@ -306,66 +388,66 @@ template class MnemeMemoryBlob { LOG_DEBUG( "Comparing blob {} full range [0, {}) with whole-blob metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", getBlobAddr(), getSize(), - static_cast(getMetadata().builtin), getMetadata().threshold, - static_cast(getMetadata().threshold_kind), - static_cast(getMetadata().norm), - getMetadata().tag.value_or("no_tag")); + static_cast(WholeBlob->MD.builtin), WholeBlob->MD.threshold, + static_cast(WholeBlob->MD.threshold_kind), + static_cast(WholeBlob->MD.norm), + WholeBlob->MD.tag.value_or("no_tag")); return compareRange(static_cast(other.getBlobAddr()), static_cast(getBlobAddr()), getSize(), - getMetadata()); + WholeBlob->MD); } uint64_t CurrentOffset = 0; for (const auto &Region : Regions) { - if (Region.Offset > CurrentOffset) { + if (Region.Range.Offset > CurrentOffset) { LOG_DEBUG( "Comparing blob {} gap [{}, {}) with whole-blob metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", - getBlobAddr(), CurrentOffset, Region.Offset, - static_cast(getMetadata().builtin), getMetadata().threshold, - static_cast(getMetadata().threshold_kind), - static_cast(getMetadata().norm), - getMetadata().tag.value_or("no_tag")); + getBlobAddr(), CurrentOffset, Region.Range.Offset, + static_cast(WholeBlob->MD.builtin), WholeBlob->MD.threshold, + static_cast(WholeBlob->MD.threshold_kind), + static_cast(WholeBlob->MD.norm), + WholeBlob->MD.tag.value_or("no_tag")); } if (!compareRange(static_cast(other.getBlobAddr()) + CurrentOffset, static_cast(getBlobAddr()) + CurrentOffset, - Region.Offset - CurrentOffset, getMetadata())) { + Region.Range.Offset - CurrentOffset, WholeBlob->MD)) { return false; } LOG_DEBUG( "Comparing blob {} region [{}, {}) with region metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", - getBlobAddr(), Region.Offset, Region.endOffset(), + getBlobAddr(), Region.Range.Offset, Region.Range.endOffset(), static_cast(Region.MD.builtin), Region.MD.threshold, static_cast(Region.MD.threshold_kind), static_cast(Region.MD.norm), Region.MD.tag.value_or("no_tag")); if (!compareRange(static_cast(other.getBlobAddr()) + - Region.Offset, + Region.Range.Offset, static_cast(getBlobAddr()) + - Region.Offset, - Region.Extent, Region.MD)) { + Region.Range.Offset, + Region.Range.Extent, Region.MD)) { return false; } - CurrentOffset = Region.endOffset(); + CurrentOffset = Region.Range.endOffset(); } if (CurrentOffset < getSize()) { LOG_DEBUG( "Comparing blob {} tail [{}, {}) with whole-blob metadata builtin={} threshold={} threshold_kind={} norm={} tag={}", getBlobAddr(), CurrentOffset, getSize(), - static_cast(getMetadata().builtin), getMetadata().threshold, - static_cast(getMetadata().threshold_kind), - static_cast(getMetadata().norm), - getMetadata().tag.value_or("no_tag")); + static_cast(WholeBlob->MD.builtin), WholeBlob->MD.threshold, + static_cast(WholeBlob->MD.threshold_kind), + static_cast(WholeBlob->MD.norm), + WholeBlob->MD.tag.value_or("no_tag")); } return compareRange(static_cast(other.getBlobAddr()) + CurrentOffset, static_cast(getBlobAddr()) + CurrentOffset, - getSize() - CurrentOffset, getMetadata()); + getSize() - CurrentOffset, WholeBlob->MD); } bool operator!=(const MnemeMemoryBlob &other) const { return !(*this == other); @@ -376,7 +458,8 @@ template llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const MnemeMemoryBlob &Blob) { // The format in the binary is the following: - // | Var Actual Size | Var-Size | Device Address | Var Data | Metadata + // | Var Actual Size | Var-Size | Device Address | Var Data | + // | NumAnnotations | (Offset | Extent | Metadata) * NumAnnotations OS << llvm::StringRef(reinterpret_cast(&Blob.ActualSize), sizeof(Blob.ActualSize)); OS << llvm::StringRef(reinterpret_cast(&Blob.Size), @@ -401,8 +484,7 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, EC.value() + "\n"); OS << llvm::StringRef( reinterpret_cast(Blob.getHostData().get()), Blob.Size); - auto MD = Blob.getMetadata(); - mneme::metadata::serialize(OS, MD); + writeAnnotationsToStream(OS, Blob.getAnnotations()); return OS; } diff --git a/include/mneme/MnemeRecordingBackend.hpp b/include/mneme/MnemeRecordingBackend.hpp index fb69f2bf..0abfd0f5 100644 --- a/include/mneme/MnemeRecordingBackend.hpp +++ b/include/mneme/MnemeRecordingBackend.hpp @@ -117,8 +117,7 @@ class RecordingBackend final : public RecorderBackend { static_cast(md.builtin), md.threshold, static_cast(md.threshold_kind), static_cast(md.norm), md.tag.value_or("no_tag")); - It->second.setMetadata(std::move(md)); - return true; + return It->second.registerAnnotation({0, It->second.getSize()}, std::move(md)); } bool setMetadataForRegion(const void *ptr, size_t bytes, Metadata md) override { @@ -161,7 +160,7 @@ class RecordingBackend final : public RecorderBackend { return false; } - if (!It->second.setRegionMetadata(Offset, bytes, md)) { + if (!It->second.registerAnnotation({Offset, bytes}, md)) { LOG_WARN( "region annotation rejected due to overlap ptr={} base_ptr={} offset={} bytes={} tag={}", ptr, It->first, Offset, bytes, md.tag.value_or("no_tag")); @@ -191,7 +190,11 @@ class RecordingBackend final : public RecorderBackend { return false; } - md = It->second.getMetadata(); + auto *WholeBlob = It->second.getWholeBlobAnnotation(); + if (!WholeBlob) + return false; + + md = WholeBlob->MD; LOG_DEBUG("getMetadataForPointer hit ptr={} tag={}", ptr, md.tag.value_or("no_tag")); return true; @@ -212,10 +215,13 @@ class RecordingBackend final : public RecorderBackend { return false; } + Metadata ExistingMd; + if (!getMetadataForPointer(ptr, ExistingMd)) + return false; + LOG_INFO("eraseMetadataForPointer hit ptr={} tag={}", ptr, - It->second.getMetadata().tag.value_or("no_tag")); - It->second.setMetadata(Metadata{}); - return true; + ExistingMd.tag.value_or("no_tag")); + return It->second.registerAnnotation({0, It->second.getSize()}, Metadata{}); } DeviceError_t rtMalloc(void **ptr, size_t size) override { diff --git a/include/mneme/MnemeSnapshot.hpp b/include/mneme/MnemeSnapshot.hpp index bc631aaa..937880f0 100644 --- a/include/mneme/MnemeSnapshot.hpp +++ b/include/mneme/MnemeSnapshot.hpp @@ -129,11 +129,10 @@ template class MnemeSnapshot { using DeviceError_t = typename MnemeDeviceRT::DeviceError_t; using DeviceStream_t = typename MnemeDeviceRT::DeviceStream_t; using KernelFunction_t = typename MnemeDeviceRT::KernelFunction_t; - static constexpr const char DiffMagic[] = "MNEME_DIFF_V1"; + static constexpr const char BytesMagic[] = "MNEME_BYTES_V2"; + static constexpr size_t BytesMagicSize = sizeof(BytesMagic) - 1; + static constexpr const char DiffMagic[] = "MNEME_DIFF_V2"; static constexpr size_t DiffMagicSize = sizeof(DiffMagic) - 1; - static constexpr const char RegionTrailerMagic[] = "MNEME_REGION_V1"; - static constexpr size_t RegionTrailerMagicSize = - sizeof(RegionTrailerMagic) - 1; static constexpr size_t DiffChunkSize = 1 << 20; class CountingRawOStream : public llvm::raw_ostream { @@ -147,6 +146,11 @@ template class MnemeSnapshot { uint64_t bytesWritten() const { return tell(); } }; + static bool isBytesBuffer(llvm::StringRef Buffer) { + return Buffer.size() >= BytesMagicSize && + Buffer.take_front(BytesMagicSize) == llvm::StringRef(BytesMagic); + } + static bool isDiffBuffer(llvm::StringRef Buffer) { return Buffer.size() >= DiffMagicSize && Buffer.take_front(DiffMagicSize) == llvm::StringRef(DiffMagic); @@ -284,91 +288,6 @@ template class MnemeSnapshot { } } - static void writeRegionMetadataTrailer( - llvm::raw_ostream &OS, - const llvm::DenseMap> &DeviceMemory) { - size_t TotalBlobsWithRegions = 0; - for (const auto &[Ptr, Blob] : DeviceMemory) { - if (Blob.hasRegionMetadata()) - ++TotalBlobsWithRegions; - } - // The "trailer" is an optional block of bytes appended after the normal - // snapshot payload (globals, blobs, kernel args). Keep older snapshots - // valid by omitting it entirely when no blobs carry subregion metadata. - if (TotalBlobsWithRegions == 0) - return; - - // Key trailer entries by blob device address so they can be attached after - // the main blob records have already been deserialized. - util::writeBytes(OS, - llvm::StringRef(RegionTrailerMagic, RegionTrailerMagicSize)); - util::writeScalar(OS, TotalBlobsWithRegions); - for (const auto &[Ptr, Blob] : DeviceMemory) { - if (!Blob.hasRegionMetadata()) - continue; - auto *BlobAddr = Blob.getBlobAddr(); - util::writeScalar(OS, BlobAddr); - - const auto &Regions = Blob.getRegionMetadata(); - size_t NumRegions = Regions.size(); - util::writeScalar(OS, NumRegions); - for (const auto &Region : Regions) { - util::writeScalar(OS, Region.Offset); - util::writeScalar(OS, Region.Extent); - metadata::serialize(OS, Region.MD); - } - } - } - - static void readRegionMetadataTrailer( - const std::string &Filename, const char *&CurrentPtr, - const char *BufferEnd, - llvm::DenseMap> &DeviceMemory) { - // No trailing bytes means this snapshot predates region metadata support. - if (CurrentPtr == BufferEnd) - return; - - // If trailing bytes exist, they must be the region trailer. Anything else - // indicates a malformed or incompatible snapshot payload. - if (static_cast(BufferEnd - CurrentPtr) < RegionTrailerMagicSize || - llvm::StringRef(CurrentPtr, RegionTrailerMagicSize) != - llvm::StringRef(RegionTrailerMagic, RegionTrailerMagicSize)) { - LOG_FATAL("Unexpected trailing bytes in Mneme snapshot " + Filename); - } - CurrentPtr += RegionTrailerMagicSize; - - size_t TotalBlobs = util::extractScalar(CurrentPtr); - for (size_t I = 0; I < TotalBlobs; ++I) { - void *DeviceAddr = util::extractScalar(CurrentPtr); - size_t NumRegions = util::extractScalar(CurrentPtr); - - auto It = DeviceMemory.find(DeviceAddr); - if (It == DeviceMemory.end()) - LOG_FATAL("Mneme region trailer references device allocation missing " - "from snapshot"); - - // Rebuild the full region table for this blob, then install it as one - // validated replacement so ordering/overlap checks happen once. - std::vector Regions; - Regions.reserve(NumRegions); - for (size_t R = 0; R < NumRegions; ++R) { - MemoryRegionMetadata Region; - Region.Offset = util::extractScalar(CurrentPtr); - Region.Extent = util::extractScalar(CurrentPtr); - Region.MD = metadata::fromBuffer(CurrentPtr); - Regions.push_back(std::move(Region)); - } - - if (!It->second.replaceRegionMetadata(std::move(Regions))) - LOG_FATAL("Mneme region trailer contains invalid or overlapping " - "regions for device allocation"); - } - - if (CurrentPtr != BufferEnd) - LOG_FATAL("Unexpected extra bytes after Mneme region trailer in " + - Filename); - } - static void readFullMnemeSnapShot( llvm::MemoryBuffer *Buffer, std::unordered_map &GlobalVars, @@ -376,7 +295,13 @@ template class MnemeSnapshot { std::shared_ptr KInfo) { auto *Start = Buffer->getBufferStart(); auto *BufferEnd = Buffer->getBufferEnd(); - auto *CurrentPtr = Start; + if (static_cast(BufferEnd - Start) < BytesMagicSize || + llvm::StringRef(Start, BytesMagicSize) != + llvm::StringRef(BytesMagic, BytesMagicSize)) { + LOG_FATAL("Unsupported Mneme bytes snapshot format in " + + Buffer->getBufferIdentifier().str()); + } + auto *CurrentPtr = Start + BytesMagicSize; size_t TotalGlobals = util::extractScalar(CurrentPtr); LOG_DEBUG("Snapshot contains {} Globals at location {}", TotalGlobals, (uintptr_t)CurrentPtr - (uintptr_t)Start); @@ -404,9 +329,9 @@ template class MnemeSnapshot { KInfo->KernelArgSizes[A] = util::extractScalar(CurrentPtr); KInfo->setArgData(CurrentPtr, A); } - - readRegionMetadataTrailer(Buffer->getBufferIdentifier().str(), CurrentPtr, - BufferEnd, DeviceMemory); + if (CurrentPtr != BufferEnd) + LOG_FATAL("Unexpected trailing bytes in Mneme snapshot " + + Buffer->getBufferIdentifier().str()); } // Applies the diff ranges from DiffBuffer onto the already-loaded base @@ -452,8 +377,7 @@ template class MnemeSnapshot { 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); + size_t NumAnnotations = util::extractScalar(CurrentPtr); auto It = DeviceMemory.find(DeviceAddr); if (It == DeviceMemory.end()) @@ -462,16 +386,18 @@ template class MnemeSnapshot { auto &Blob = It->second; if (Blob.getActualSize() != ActualSize || Blob.getSize() != Size) LOG_FATAL("Mneme diff memory blob size mismatch"); - Blob.setMetadata(MD); + auto Annotations = readAnnotationsFromBuffer(CurrentPtr, NumAnnotations); + if (!Blob.replaceAnnotations(std::move(Annotations))) + LOG_FATAL("Mneme diff blob annotation set is invalid"); + size_t NumRanges = util::extractScalar(CurrentPtr); applyDiffRanges( CurrentPtr, llvm::MutableArrayRef(Blob.getHostData().get(), Blob.getSize()), NumRanges); } - - readRegionMetadataTrailer(Filename, CurrentPtr, DiffBuffer->getBufferEnd(), - DeviceMemory); + if (CurrentPtr != DiffBuffer->getBufferEnd()) + LOG_FATAL("Unexpected trailing bytes in Mneme diff " + Filename); } static size_t getSerializedMetadataSize(const Metadata &MD) { @@ -480,34 +406,16 @@ template class MnemeSnapshot { (MD.tag ? MD.tag->size() : 0); } - static size_t - getSerializedRegionMetadataSize(const MemoryRegionMetadata &Region) { - return sizeof(Region.Offset) + sizeof(Region.Extent) + + static size_t getSerializedAnnotationSize(const MemoryAnnotation &Region) { + return sizeof(Region.Range.Offset) + sizeof(Region.Range.Extent) + getSerializedMetadataSize(Region.MD); } - static size_t computeRegionMetadataTrailerSize( - const llvm::DenseMap> &DeviceMemory) { - size_t TotalBlobsWithRegions = 0; - size_t Size = 0; - for (const auto &[Ptr, Blob] : DeviceMemory) { - if (!Blob.hasRegionMetadata()) - continue; - if (TotalBlobsWithRegions++ == 0) - Size = RegionTrailerMagicSize + sizeof(size_t); - Size += sizeof(void *); - Size += sizeof(size_t); - for (const auto &Region : Blob.getRegionMetadata()) - Size += getSerializedRegionMetadataSize(Region); - } - return Size; - } - static size_t computeMnemeBytesSnapshotSize( const proteus::runtime::GlobalMetadataMap &GlobalVars, const llvm::DenseMap> &DeviceMemory, llvm::ArrayRef KernelArgSizes) { - size_t Size = sizeof(size_t); + size_t Size = BytesMagicSize + sizeof(size_t); for (const auto &[VarName, GV] : GlobalVars) { Size += sizeof(size_t); Size += VarName.size(); @@ -522,7 +430,9 @@ template class MnemeSnapshot { Size += sizeof(size_t); Size += sizeof(void *); Size += Blob.getSize(); - Size += getSerializedMetadataSize(Blob.getMetadata()); + Size += sizeof(size_t); + for (const auto &Annotation : Blob.getAnnotations()) + Size += getSerializedAnnotationSize(Annotation); } Size += sizeof(size_t); @@ -530,7 +440,6 @@ template class MnemeSnapshot { Size += sizeof(size_t); Size += ArgSize; } - Size += computeRegionMetadataTrailerSize(DeviceMemory); return Size; } @@ -566,6 +475,7 @@ template class MnemeSnapshot { if (DEC) LOG_FATAL("Synnchronizing stream failed"); llvm::raw_fd_ostream OutBC(Filename.string(), EC); + util::writeBytes(OutBC, llvm::StringRef(BytesMagic, BytesMagicSize)); // First write Global Variables. size_t TotalGlobals = GlobalVars.size(); OutBC << llvm::StringRef(reinterpret_cast(&TotalGlobals), @@ -627,8 +537,6 @@ template class MnemeSnapshot { KernelArgSizes[I]); } - writeRegionMetadataTrailer(OutBC, DeviceMemory); - return Filename.filename(); } @@ -673,13 +581,10 @@ template class MnemeSnapshot { util::writeScalar(OutBC, Blob.getSize()); auto *BlobAddr = Blob.getBlobAddr(); util::writeScalar(OutBC, BlobAddr); - auto MD = Blob.getMetadata(); - mneme::metadata::serialize(OutBC, MD); + writeAnnotationsToStream(OutBC, Blob.getAnnotations()); writeCountAndWriteChangedRanges(OutBC, Blob, UpdateBaseData); } - - writeRegionMetadataTrailer(OutBC, DeviceMemory); } static size_t measureMnemeDiffSnapshotSize( @@ -746,8 +651,10 @@ template class MnemeSnapshot { if (isDiffBuffer(Buffer->getBuffer())) return std::make_unique>(Filename, std::move(Buffer)); - return std::make_unique>(Filename, - std::move(Buffer)); + if (isBytesBuffer(Buffer->getBuffer())) + return std::make_unique>(Filename, + std::move(Buffer)); + LOG_FATAL("Unsupported Mneme snapshot format in " + Filename); } static Snapshot readBytesSnapshot(std::string KernelName, From 57f6016e9d5cd0d9c5ae45e5f2a513ee04a784eb Mon Sep 17 00:00:00 2001 From: Caetano Melone Date: Wed, 2 Sep 2026 11:33:35 -0700 Subject: [PATCH 4/5] docs and tests Signed-off-by: Caetano Melone --- docs/usage/record.md | 23 +++++ docs/usage/verification.md | 35 +++++++ tests/record/read-record.py | 59 +++++++---- tests/record/test_annotation.cpp | 14 ++- tests/unit_tests/ReadWriteSnapshot.cpp | 130 ++++++++++++++++++++++--- 5 files changed, 230 insertions(+), 31 deletions(-) diff --git a/docs/usage/record.md b/docs/usage/record.md index 8522f7cb..d530ad60 100644 --- a/docs/usage/record.md +++ b/docs/usage/record.md @@ -33,10 +33,33 @@ mneme::annotate(d_output, mneme::Metadata{ }); ``` +For cases where the pointer you want to annotate is an interior pointer +into a larger allocation, use the sub-region annotation feature and +pass the byte extent of the logical field you want Mneme to verify: + + +```cpp +double* d_alias = base + offset; +std::size_t alias_bytes = count * sizeof(double); + +mneme::annotate(d_alias, alias_bytes, mneme::Metadata{ + .builtin = mneme::BuiltinDType::F64, + .threshold = 1e-6, + .threshold_kind = mneme::ThresholdKind::Relative, + .norm = mneme::Norm::Linf, + .tag = std::string("alias_region"), +}); +``` + Annotations must be applied **before** the kernel launch they should affect. You can update the annotation on the same pointer between launches to record different tolerance policies for different dynamic instances of the same kernel. +Use `mneme::annotate(ptr, md)` when `ptr` is the base of the whole +allocation. Use `mneme::annotate(ptr, nbytes, md)` when `ptr` is an +interior pointer or when only a sub-region of the allocation should +carry that metadata. + For the full API reference, supported data types, threshold semantics, and a complete example, see **[Usage → Verification](verification.md)**. diff --git a/docs/usage/verification.md b/docs/usage/verification.md index 92312877..5efa320f 100644 --- a/docs/usage/verification.md +++ b/docs/usage/verification.md @@ -265,11 +265,46 @@ replayed epilogue satisfies every annotated buffer's tolerance. --- +## Region annotations + +Use the region overload when the pointer to annotate is an alias or +interior pointer into a larger allocation, or when one allocation +contains multiple logical fields with different verification rules. + +```cpp +double* d_alias = base + offset; +std::size_t alias_bytes = count * sizeof(double); + +mneme::annotate(d_alias, alias_bytes, mneme::Metadata{ + .builtin = mneme::BuiltinDType::F64, + .threshold = 1e-12, + .threshold_kind = mneme::ThresholdKind::Relative, + .norm = mneme::Norm::L2, + .tag = std::string("alias_region"), +}); +``` + +- `mneme::annotate(ptr, nbytes, md)` applies metadata only to the byte range + `[ptr, ptr + nbytes)`. +- The full-range annotation acts as the default policy for uncovered bytes. +- Narrower region annotations override that default on covered bytes. + +Constraints: +- Overlapping region annotations on the same blob are rejected. +- `nbytes` must stay within the owning allocation. +- `nbytes` must be compatible with the selected `BuiltinDType` element size. + +This is the intended API for alias pointers produced by frameworks such as +MFEM on top of allocators such as Umpire. + +--- + ## Summary | Concept | Description | | ------- | ----------- | | `mneme::annotate(ptr, md)` | Attach verification metadata to a device pointer | +| `mneme::annotate(ptr, nbytes, md)` | Attach verification metadata to a byte sub-region of an owning allocation | | `BuiltinDType` | Scalar type used to interpret buffer contents | | `ThresholdKind` | Absolute or relative error formula | | `Norm` | Per-element (`None`) or aggregate (`L1`, `L2`, `Linf`) comparison | diff --git a/tests/record/read-record.py b/tests/record/read-record.py index d8e8b13c..b2812355 100644 --- a/tests/record/read-record.py +++ b/tests/record/read-record.py @@ -10,20 +10,24 @@ # Binary .mneme prologue parser # # On-disk format (all integers little-endian, pointer-sized = 8 bytes): +# char[14] "MNEME_BYTES_V2" # uint64 TotalGlobals # for each global: # uint64 StrLen; char[StrLen] name; uint64 VarSize; void* DevAddr; char[VarSize] data # uint64 TotalBlobs # for each blob: # uint64 ActualSize; uint64 Size; void* DevAddr; char[Size] data -# Metadata: uint8 builtin; double threshold; uint8 threshold_kind; -# uint8 norm; uint64 tag_len; char[tag_len] tag +# uint64 NumAnnotations +# for each annotation: +# uint64 Offset; uint64 Extent; +# Metadata: uint8 builtin; double threshold; uint8 threshold_kind; +# uint8 norm; uint64 tag_len; char[tag_len] tag # uint64 NumArgs # for each arg: uint64 ArgSize; char[ArgSize] data # --------------------------------------------------------------------------- def _parse_prologue_metadata(filename): - """Return a list of Metadata dicts for every blob in a prologue file.""" + """Return a list of annotation dicts for every blob in a prologue file.""" with open(filename, "rb") as f: data = f.read() off = 0 @@ -50,6 +54,11 @@ def skip(n): nonlocal off off += n + bytes_magic = b"MNEME_BYTES_V2" + if not data.startswith(bytes_magic): + raise ValueError("unsupported prologue format") + off += len(bytes_magic) + # globals for _ in range(u64()): str_len = u64() @@ -65,19 +74,23 @@ def skip(n): size = u64() skip(8) # dev_addr skip(size) # blob data - builtin = u8() - threshold = dbl() - threshold_kind = u8() - norm = u8() - tag_len = u64() - tag = None - if tag_len: - tag = data[off:off + tag_len].decode("utf-8", errors="replace") - skip(tag_len) - results.append( - dict(builtin=builtin, threshold=threshold, - threshold_kind=threshold_kind, norm=norm, tag=tag) - ) + for _ in range(u64()): + offset = u64() + extent = u64() + builtin = u8() + threshold = dbl() + threshold_kind = u8() + norm = u8() + tag_len = u64() + tag = None + if tag_len: + tag = data[off:off + tag_len].decode("utf-8", errors="replace") + skip(tag_len) + results.append( + dict(offset=offset, extent=extent, blob_size=size, + builtin=builtin, threshold=threshold, + threshold_kind=threshold_kind, norm=norm, tag=tag) + ) return results data_dir = sys.argv[1] if len(sys.argv) > 1 else "." @@ -129,7 +142,19 @@ def skip(n): ] if md["tag"] is not None: parts.append(f"tag={md['tag']}") - print("BlobAnnotation:", " ".join(parts)) + if md["offset"] == 0 and md["extent"] == md["blob_size"]: + print("BlobAnnotation:", " ".join(parts)) + else: + print( + "RegionAnnotation:", + " ".join( + [ + f"offset={md['offset']}", + f"extent={md['extent']}", + *parts, + ] + ), + ) except Exception as e: print(f"Warning: could not parse prologue metadata: {e}", file=sys.stderr) diff --git a/tests/record/test_annotation.cpp b/tests/record/test_annotation.cpp index e0cff203..7044df52 100644 --- a/tests/record/test_annotation.cpp +++ b/tests/record/test_annotation.cpp @@ -28,10 +28,13 @@ __global__ void annotation_kernel(float *data, int n) { int main() { const int N = 64; + const int AliasOffset = 16; + const int AliasLen = 24; float *d_ptr = nullptr; MnemeDeviceRT::DeviceMalloc(reinterpret_cast(&d_ptr), N * sizeof(float)); + float *d_alias = d_ptr + AliasOffset; // Annotate the device pointer before the kernel launch so the preload can // attach metadata to the recorded snapshot. @@ -43,7 +46,15 @@ int main() { .tag = std::string("test_ptr"), }); - annotation_kernel<<<1, N>>>(d_ptr, N); + mneme::annotate(d_alias, AliasLen * sizeof(float), mneme::Metadata{ + .builtin = mneme::BuiltinDType::F32, + .threshold = 0.002, + .threshold_kind = mneme::ThresholdKind::Relative, + .norm = mneme::Norm::L2, + .tag = std::string("test_region"), + }); + + annotation_kernel<<<1, AliasLen>>>(d_alias, AliasLen); auto EC = MnemeDeviceRT::DeviceErrorCheck(MnemeDeviceRT::DeviceSynchronize()); if (EC) { @@ -63,4 +74,5 @@ int main() { // CHECK-RR: NumModules: 1 // CHECK-RR: NumInstances: 1 // CHECK-RR: BlobAnnotation: threshold=0.001 threshold_kind=1 builtin=9 norm=3 tag=test_ptr +// CHECK-RR: RegionAnnotation: offset=64 extent=96 threshold=0.002 threshold_kind=1 builtin=9 norm=2 tag=test_region // clang-format on diff --git a/tests/unit_tests/ReadWriteSnapshot.cpp b/tests/unit_tests/ReadWriteSnapshot.cpp index 86e176fd..e82aada4 100644 --- a/tests/unit_tests/ReadWriteSnapshot.cpp +++ b/tests/unit_tests/ReadWriteSnapshot.cpp @@ -29,7 +29,7 @@ bool isDiffSnapshotFile(const std::filesystem::path &Path) { std::ifstream In(Path, std::ios::binary); std::string Magic(13, '\0'); In.read(Magic.data(), Magic.size()); - return Magic == "MNEME_DIFF_V1"; + return Magic == "MNEME_DIFF_V2"; } template void initializeRandomBuffer(T *Buffer, size_t Size) { @@ -44,6 +44,20 @@ template void initializeRandomBuffer(T *Buffer, size_t Size) { } } +const Metadata &getWholeBlobMetadata(const MnemeMemoryBlobDevice &Blob) { + auto *Annotation = Blob.getWholeBlobAnnotation(); + if (!Annotation) + LOG_FATAL("Blob is missing whole-blob annotation"); + return Annotation->MD; +} + +MemoryAnnotation getOnlyRegionAnnotation(const MnemeMemoryBlobDevice &Blob) { + auto Regions = Blob.getRegionAnnotations(); + if (Regions.size() != 1) + LOG_FATAL("Expected exactly one region annotation, got {}", Regions.size()); + return Regions.front(); +} + int main(int argc, char **argv) { // We allocate some "fake" globals auto initializeDeviceData = [&] { @@ -77,7 +91,17 @@ int main(int argc, char **argv) { Md.threshold = 0.5; Md.threshold_kind = ThresholdKind::Relative; Md.tag = std::string("Test"); - Blob.setMetadata(Md); + if (!Blob.registerAnnotation({0, Blob.getSize()}, Md)) + LOG_FATAL("Could not register whole-blob annotation"); + + mneme::Metadata RegionMd; + RegionMd.builtin = BuiltinDType::U8; + RegionMd.norm = Norm::None; + RegionMd.threshold = 0.0; + RegionMd.threshold_kind = ThresholdKind::Absolute; + RegionMd.tag = std::string("Region"); + if (!Blob.registerAnnotation({16, 32}, RegionMd)) + LOG_FATAL("Could not register region annotation"); Blob.setHostData(std::unique_ptr(new uint8_t[128])); @@ -158,31 +182,44 @@ int main(int argc, char **argv) { return 1; } - if (RBlob.getMetadata().builtin != BuiltinDType::F64) { + const auto &RBlobMd = getWholeBlobMetadata(RBlob); + + if (RBlobMd.builtin != BuiltinDType::F64) { std::cerr << "Metadata builtin differs\n"; return 1; } - if (RBlob.getMetadata().norm != Norm::L2) { + if (RBlobMd.norm != Norm::L2) { std::cerr << "Metadata norm differs\n"; return 1; } - if (RBlob.getMetadata().threshold != 0.5) { + if (RBlobMd.threshold != 0.5) { std::cerr << "Metadata threshold differs\n"; return 1; } - if (RBlob.getMetadata().threshold_kind != ThresholdKind::Relative) { + if (RBlobMd.threshold_kind != ThresholdKind::Relative) { std::cerr << "Metadata threshold_kind differs\n"; return 1; } - if (RBlob.getMetadata().tag.value() != "Test") { + if (RBlobMd.tag.value() != "Test") { std::cerr << "Metadata tag differs\n"; return 1; } + const auto &RRegion = getOnlyRegionAnnotation(RBlob); + if (RRegion.Range.Offset != 16 || RRegion.Range.Extent != 32 || + RRegion.MD.builtin != BuiltinDType::U8 || + RRegion.MD.norm != Norm::None || + RRegion.MD.threshold != 0.0 || + RRegion.MD.threshold_kind != ThresholdKind::Absolute || + RRegion.MD.tag.value() != "Region") { + std::cerr << "Region metadata differs\n"; + return 1; + } + uint8_t *WData = WBlob.getHostData().get(); uint8_t *RData = RBlob.getHostData().get(); if (std::memcmp(reinterpret_cast(WData), @@ -291,15 +328,26 @@ int main(int argc, char **argv) { return 16; } - if (RBlob.getMetadata().builtin != BuiltinDType::F64 || - RBlob.getMetadata().norm != Norm::L2 || - RBlob.getMetadata().threshold != 0.5 || - RBlob.getMetadata().threshold_kind != ThresholdKind::Relative || - RBlob.getMetadata().tag.value() != "Test") { + const auto &RBlobMd = getWholeBlobMetadata(RBlob); + if (RBlobMd.builtin != BuiltinDType::F64 || + RBlobMd.norm != Norm::L2 || RBlobMd.threshold != 0.5 || + RBlobMd.threshold_kind != ThresholdKind::Relative || + RBlobMd.tag.value() != "Test") { std::cerr << "Diff blob metadata differs\n"; return 16; } + const auto &RRegion = getOnlyRegionAnnotation(RBlob); + if (RRegion.Range.Offset != 16 || RRegion.Range.Extent != 32 || + RRegion.MD.builtin != BuiltinDType::U8 || + RRegion.MD.norm != Norm::None || + RRegion.MD.threshold != 0.0 || + RRegion.MD.threshold_kind != ThresholdKind::Absolute || + RRegion.MD.tag.value() != "Region") { + std::cerr << "Diff region metadata differs\n"; + return 16; + } + if (std::memcmp(BlobData.second, RBlob.getHostData().get(), 128) != 0) { std::cerr << "Diff device memory did not reconstruct epilogue data\n"; return 16; @@ -359,6 +407,61 @@ int main(int argc, char **argv) { return 0; }(); + // Replay loads annotations from snapshots first, then calls map()/allocate() + // while materializing prologue/epilogue memory. This regression check makes + // sure those setup paths do not clobber deserialized whole-blob or region + // annotations back to the default full-range state. + auto ValidateAnnotationPreservationAcrossReplaySetup = [&]() { + auto PrologueIt = ReadDeviceMemMap.find((void *)BlobData.first); + if (PrologueIt == ReadDeviceMemMap.end()) { + std::cerr << "Prologue snapshot missing blob for preservation test\n"; + return 256; + } + + auto &MappedBlob = PrologueIt->second; + auto WholeBeforeMap = getWholeBlobMetadata(MappedBlob); + auto RegionBeforeMap = getOnlyRegionAnnotation(MappedBlob); + EC = MnemeDeviceRT::DeviceErrorCheck( + MappedBlob.map((void *)BlobData.first, MappedBlob.getActualSize(), + MappedBlob.getSize())); + if (EC) { + std::cerr << "Could not map blob during preservation test\n"; + return 256; + } + const auto &WholeAfterMap = getWholeBlobMetadata(MappedBlob); + const auto &RegionAfterMap = getOnlyRegionAnnotation(MappedBlob); + if (WholeAfterMap != WholeBeforeMap || RegionAfterMap != RegionBeforeMap) { + std::cerr << "Annotations changed across map()\n"; + return 256; + } + MappedBlob.release(); + + auto DiffIt = DiffDeviceMemMap.find((void *)BlobData.first); + if (DiffIt == DiffDeviceMemMap.end()) { + std::cerr << "Diff snapshot missing blob for preservation test\n"; + return 256; + } + + auto &AllocatedBlob = DiffIt->second; + auto WholeBeforeAlloc = getWholeBlobMetadata(AllocatedBlob); + auto RegionBeforeAlloc = getOnlyRegionAnnotation(AllocatedBlob); + EC = MnemeDeviceRT::DeviceErrorCheck( + AllocatedBlob.allocate(AllocatedBlob.getSize())); + if (EC) { + std::cerr << "Could not allocate blob during preservation test\n"; + return 256; + } + const auto &WholeAfterAlloc = getWholeBlobMetadata(AllocatedBlob); + const auto &RegionAfterAlloc = getOnlyRegionAnnotation(AllocatedBlob); + if (WholeAfterAlloc != WholeBeforeAlloc || + RegionAfterAlloc != RegionBeforeAlloc) { + std::cerr << "Annotations changed across allocate()\n"; + return 256; + } + AllocatedBlob.release(); + return 0; + }(); + auto PrologueBlobIt = ReadDeviceMemMap.find((void *)BlobData.first); auto PrologueGlobalIt = ReadGVars.find("Test"); auto *PrologueBlob = PrologueBlobIt->second.getHostData().get(); @@ -412,7 +515,8 @@ int main(int argc, char **argv) { auto Ret = ValidateGlobalMem | ValidateDeviceMem | ValidateKernelArgs | ValidateDiffGlobalMem | ValidateDiffDeviceMem | ValidateDiffKernelArgs | ValidateBestSparse | - ValidateBestFragmented; + ValidateBestFragmented | + ValidateAnnotationPreservationAcrossReplaySetup; delete[] GlobalData.second; delete[] BlobData.second; From 4eadef5f9b6ac6b867d611c0d6440c75cf41bd0f Mon Sep 17 00:00:00 2001 From: Caetano Melone Date: Fri, 4 Sep 2026 12:31:41 -0700 Subject: [PATCH 5/5] use bytes_v2 in test Signed-off-by: Caetano Melone --- python/tests/test_record_annotations.py | 48 +++++++++++++++---------- 1 file changed, 29 insertions(+), 19 deletions(-) diff --git a/python/tests/test_record_annotations.py b/python/tests/test_record_annotations.py index 146cdf0a..6172da58 100644 --- a/python/tests/test_record_annotations.py +++ b/python/tests/test_record_annotations.py @@ -16,6 +16,7 @@ def _parse_prologue_blob_metadata(prologue_path: Path): data = prologue_path.read_bytes() off = 0 + bytes_magic = b"MNEME_BYTES_V2" def read_u64(): nonlocal off @@ -39,6 +40,9 @@ def skip(n): nonlocal off off += n + assert data.startswith(bytes_magic), "unsupported prologue snapshot format" + off += len(bytes_magic) + # globals for _ in range(read_u64()): name_len = read_u64() @@ -55,25 +59,31 @@ def skip(n): skip(8) # dev addr skip(blob_size) - builtin = read_u8() - threshold = read_f64() - threshold_kind = read_u8() - norm = read_u8() - tag_len = read_u64() - tag = None - if tag_len: - tag = data[off : off + tag_len].decode("utf-8", errors="replace") - skip(tag_len) - - out.append( - { - "builtin": builtin, - "threshold": threshold, - "threshold_kind": threshold_kind, - "norm": norm, - "tag": tag, - } - ) + for _ in range(read_u64()): + offset = read_u64() + extent = read_u64() + builtin = read_u8() + threshold = read_f64() + threshold_kind = read_u8() + norm = read_u8() + tag_len = read_u64() + tag = None + if tag_len: + tag = data[off : off + tag_len].decode("utf-8", errors="replace") + skip(tag_len) + + out.append( + { + "offset": offset, + "extent": extent, + "blob_size": blob_size, + "builtin": builtin, + "threshold": threshold, + "threshold_kind": threshold_kind, + "norm": norm, + "tag": tag, + } + ) return out