[memprof] Add a raw binary format to serialize memprof profiles.

[snehasish]

Summary

This change implements the raw binary format discussed in
https://lists.llvm.org/pipermail/llvm-dev/2021-September/153007.html

Summary of changes

• Add a new memprof option to choose binary or text (default) format.
• Add a rawprofile library which serializes the MIB map to profile.
• Add a unit test for rawprofile.
• Mark sanitizer procmaps methods as virtual to be able to mock them.
• Extend memprof_profile_dump regression test.

[davidxl]

This is raw format only change, right? Probably update the summary to reflect that.

[snehasish]

Update description, rebase on stacked changes.

[davidxl]

compiler-rt/lib/memprof/memprof_allocator.cpp

263

Document 'true' argument with comment. Also should it be controlled by a flag?

269

Add a comment for the function.

compiler-rt/lib/memprof/memprof_flags.inc

38

should it be false by default?

compiler-rt/lib/memprof/memprof_rawprofile.cpp

17

Does it need to be packed?

25

is the assertion supposed be on?

63

Add as function comments about the stack section layout.

compiler-rt/test/memprof/TestCases/memprof_profile_dump.cpp

28

This depends on endianess.

28

Should more info be checked?

[tejohnson]

compiler-rt/lib/memprof/memprof_allocator.cpp

17

Is this change related to this patch?

232

We're now printing this twice afaict (here and in FinishAndWrite) - is this intentional?

290

In D111676 it was suggested to move the unlock until after the ForEach PrintCallback to address a race, but this effectively removes that change. Intentional?

compiler-rt/lib/memprof/memprof_rawprofile.cpp

64

Comment about the first value being counted here.

94

Nit: suggest reversing the above 2 lines, logic flows better.

97

Why multiply by sizeof(Count) and not just directly sizeof(uptr) or sizeof(u64)? It seems a little odd to use the size of the Count variable, which happens to be u64 but isn't necessarily related to the size of the pc values.

105

Suggest being more explicit about what value you are talking about. I.e. you don't touch the provided MIB value. Does it need to be marked as unused to avoid warnings (looks like there is an UNUSED macro defined in sanitizer_internal_defs.h that you could use).

112

Comment about first value being counted here.

116

Comment about u64 value being written with each MIB.

186

How much extra memory do we carry by essentially duplicating the necessary allocations (e.g. we have allocated memory for the segment/mibinfo/stack already, then allocate it a second time here)? Since the 3 functions that allocate these individual buffers have 2 parts to them (compute # bytes, then write bytes), should we consider splitting each of those into 2 separate functions, call the computation parts first, then do a single allocation here, then write the individual components into this Buffer directly?

compiler-rt/test/memprof/TestCases/interface_test.cpp

5

Seems like an unrelated fix - can you just commit it NFC separately and remove from here?

[snehasish]

Address review comments.

Also

• Update test to use od instead of xxd since its not available on debian buildbots.
• Some smaller changes to address warnings.

[snehasish]

Thanks for the prompt review! PTAL.

compiler-rt/lib/memprof/memprof_allocator.cpp

17

No, reverted.

232

No, its leftover from a merge. Reverted.

263

Done. I think setting it to true is more conservative - forces an update of the cache and then reads the proc maps again to check if anything was missed (i.e. due to concurrent updates). Setting it to false could be faster but the gains are probably negligible so I think hard-coding the value is fine.

290

Not intentional. Thinking about this some more - the deallocation path is guarded by the destructing var - so print/merge conflict should not happen. Thus we can remove the locks from this path altogether. I updated the code - converted destructing and constructed to atomics and removed the ForceLock/ForceUnlock calls since they are not used anywhere else. Wdyt?

compiler-rt/lib/memprof/memprof_flags.inc

38

Not planning on making it the default in this change since it would mean updating the tests which check the printed out output. I plan on adding this flag to the tests in compiler-rt and making binary the default in a future patch.

compiler-rt/lib/memprof/memprof_rawprofile.cpp

17

Yes, it is packed to make it easy to serialize/deserialize. I plan on reusing this typedef in the llvm/ProfileData reader code by moving it to the shared include header in the future.

25

Removed commented code. It would be nice to have this static check but we don't have the c++ standard library here. There are some sanitizer defined trait checks but is_pod isn't implemented.

63

Added a comment - the contents are the same as the format description in the SerializeMIBInfo function. Let me know if you think I should clarify further.

186

I started with an approach which precomputed the necessary space and performed a single allocation but then abandoned the approach in favour of the current implementation which IMO is easier to read (and test individually though I didn't take advantage of this fact).

For a clang workload (eg ~35K contexts), the overhead would be ~ equal to the size of the raw profile on disk, the largest I observed was ~5M. For a larger internal workload with ~1M observed contexts, I expect the overhead to be ~100M. I don't have a strong opinion and I can rewrite this patch to precompute and only allocate once if you would like.

compiler-rt/test/memprof/TestCases/interface_test.cpp

5

Pushed in rG9305e3b6d7e7

compiler-rt/test/memprof/TestCases/memprof_profile_dump.cpp

28

This depends on endianess.

Yes. The sanitizer tests are only enabled on x86_64 [1]. I don't think we intend to support memprof on big endian arch. So I think we can leave this as is?

Should more info be checked?

I have a more robust test for contents in the unit test (compiler-rt/lib/memprof/tests/rawprofile.cpp). Once we have a llvm-profdata reader interface to dump as text we can add something to cross-project-tests? If you feel strongly I could add more checks here for the binary octets we expect.

https://git.io/JXgFi

[davidxl]

compiler-rt/lib/memprof/tests/rawprofile.cpp

35

Perhaps make the make contain two stacks and two MIB entries?

66

Can offset be larger than start?

100

Can the actual size be checked?

130

Why using NE instead of comparing with actual id with EQ?

compiler-rt/test/memprof/TestCases/memprof_profile_dump.cpp

28

Ok, perhaps add a comment about endianess assumption

[tejohnson]

compiler-rt/lib/memprof/memprof_allocator.cpp

290

When we call this from __memprof_profile_dump we aren't destructing, so wouldn't we have an issue then?

compiler-rt/lib/memprof/memprof_rawprofile.cpp

105

Suggest changing comment to be more explicit, e.g. s/the value/the MIB value/

186

It seems large enough that I think it would be worth splitting if it isn't too much trouble.

[snehasish]

Address review comments.

• Enhance the unit test with multiple entries and strict checks.
• Re-introduce the allocator locks.
• Rewrite to memprof_rawprofile.cpp to allocate memory only once.

[snehasish]

Updated fake offset in test to match start.

[snehasish]

Rewrote the serialization logic to remove unnecessary allocations and addressed issues in the test. PTAL, thanks!

compiler-rt/lib/memprof/memprof_allocator.cpp

290

Yes it will be. I missed that case when I was reasoning about this. I reintroduced the check and updated the code here. While working on this I also realized that the malloc call memprof_rawprofile was being instrumented to and caused a deadlock when attempting to serialize the rawprofile. I converted that to use InternalAlloc. I left the constructed and destructing vars as atomics since its stricter.

compiler-rt/lib/memprof/tests/rawprofile.cpp

66

Not in practice. Updated to be the same.

100

Done here and all other checks.

[tejohnson]

lgtm with small fix noted below. Wait to see if @davidxl has more comments too

compiler-rt/lib/memprof/memprof_rawprofile.cpp

93

Include a meaningful message instead of "" (here and in other similar CHECKs below).

[snehasish]

Fill in missing messages for CHECKs.

[snehasish]

Fix typo.

[davidxl]

lgtm

[snehasish]

The fuzzer test failure seems to be unrelated so I'll go ahead and push this set of patches.

compiler-rt/lib/memprof/CMakeLists.txt

 # Build for the Memory Profiler runtime support library.
 
 set(MEMPROF_SOURCES
   memprof_allocator.cpp
   memprof_descriptions.cpp
   memprof_flags.cpp
   memprof_interceptors.cpp
   memprof_interceptors_memintrinsics.cpp
   memprof_linux.cpp
   memprof_malloc_linux.cpp
   memprof_mibmap.cpp
   memprof_posix.cpp
+  memprof_rawprofile.cpp
   memprof_rtl.cpp
   memprof_shadow_setup.cpp
   memprof_stack.cpp
   memprof_stats.cpp
   memprof_thread.cpp
   )
 
 set(MEMPROF_CXX_SOURCES
SET(MEMPROF_HEADERS
   memprof_init_version.h
   memprof_interceptors.h
   memprof_interceptors_memintrinsics.h
   memprof_interface_internal.h
   memprof_internal.h
   memprof_mapping.h
   memprof_meminfoblock.h
   memprof_mibmap.h
+  memprof_rawprofile.h
   memprof_stack.h
   memprof_stats.h
   memprof_thread.h
   )
 
 include_directories(..)
 
 set(MEMPROF_CFLAGS ${SANITIZER_COMMON_CFLAGS})
add_sanitizer_rt_symbols(clang_rt.memprof_cxx
       ARCHS ${arch})
     add_dependencies(memprof clang_rt.memprof_cxx-${arch}-symbols)
     add_sanitizer_rt_symbols(clang_rt.memprof
       ARCHS ${arch}
       EXTRA memprof.syms.extra)
     add_dependencies(memprof clang_rt.memprof-${arch}-symbols)
   endif()
 endforeach()
+
+
+if(COMPILER_RT_INCLUDE_TESTS)
+  add_subdirectory(tests)
+endif()

compiler-rt/lib/memprof/memprof_allocator.cpp

 //===-- memprof_allocator.cpp --------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of MemProfiler, a memory profiler.
 //
 // Implementation of MemProf's memory allocator, which uses the allocator
 // from sanitizer_common.
 //
 //===----------------------------------------------------------------------===//
 
 #include "memprof_allocator.h"
-#include "memprof_mapping.h"
 #include "memprof_meminfoblock.h"
 #include "memprof_mibmap.h"
+#include "memprof_rawprofile.h"
 #include "memprof_stack.h"
 #include "memprof_thread.h"
 #include "sanitizer_common/sanitizer_allocator_checks.h"
 #include "sanitizer_common/sanitizer_allocator_interface.h"
 #include "sanitizer_common/sanitizer_allocator_report.h"
 #include "sanitizer_common/sanitizer_errno.h"
 #include "sanitizer_common/sanitizer_file.h"
 #include "sanitizer_common/sanitizer_flags.h"
 #include "sanitizer_common/sanitizer_internal_defs.h"
 #include "sanitizer_common/sanitizer_list.h"
+#include "sanitizer_common/sanitizer_procmaps.h"
 #include "sanitizer_common/sanitizer_stackdepot.h"
+#include "sanitizer_common/sanitizer_vector.h"
 
 #include <sched.h>
 #include <time.h>
 
 namespace __memprof {
 
 static int GetCpuId(void) {
   // _memprof_preinit is called via the preinit_array, which subsequently calls
struct Allocator {
   uptr max_user_defined_malloc_size;
   atomic_uint8_t rss_limit_exceeded;
 
   // Holds the mapping of stack ids to MemInfoBlocks.
   MIBMapTy MIBMap;
 
   bool destructing;
   bool constructed = false;
+  bool PrintText;
 
   // ------------------- Initialization ------------------------
   explicit Allocator(LinkerInitialized)
-      : destructing(false), constructed(true) {}
-  ~Allocator() { FinishAndPrint(); }
+      : destructing(false), constructed(true), PrintText(flags()->print_text) {
+    if (common_flags()->print_module_map)
+      DumpProcessMap();
+    if (flags()->print_terse)
+      MemInfoBlock::printHeader();
+  }
+
+  ~Allocator() {
+    destructing = true;
+    FinishAndWrite();
+  }
 
   static void PrintCallback(const uptr Key, LockedMemInfoBlock *const &Value,
                             void *Arg) {
     SpinMutexLock(&Value->mutex);
     Value->mib.Print(Key, bool(Arg));
   }
 
-  void FinishAndPrint() {
+  void FinishAndWrite() {
     if (common_flags()->print_module_map)
       DumpProcessMap();
-    if (!flags()->print_terse)
+    InsertLiveBlocks();
+
+    if (PrintText) {
+      WriteTextProfile();
+      return;
+    }
+
+    // Serialize the contents to a raw profile. Format documented in
+    // memprof_rawprofile.h.
+    char *Buffer = nullptr;
+
+    MemoryMappingLayout Layout(true);
+    int BytesSerialized = SerializeToRawProfile(MIBMap, Layout, Buffer);
+    CHECK(Buffer && BytesSerialized && "could not serialize to buffer");
+    report_file.Write(Buffer, BytesSerialized);
+  }
+
+  void InsertLiveBlocks() {
+    if (PrintText && !flags()->print_terse)
       Printf("Live on exit:\n");
+
     allocator.ForceLock();
     allocator.ForEachChunk(
         [](uptr chunk, void *alloc) {
           u64 user_requested_size;
           Allocator *A = (Allocator *)alloc;
           MemprofChunk *m =
               A->GetMemprofChunk((void *)chunk, user_requested_size);
           if (!m)
             return;
           uptr user_beg = ((uptr)m) + kChunkHeaderSize;
           u64 c = GetShadowCount(user_beg, user_requested_size);
           long curtime = GetTimestamp();
           MemInfoBlock newMIB(user_requested_size, c, m->timestamp_ms, curtime,
                               m->cpu_id, GetCpuId());
           InsertOrMerge(m->alloc_context_id, newMIB, A->MIBMap);
         },
         this);
+    allocator.ForceUnlock();
+  }
 
-    destructing = true;
+  void WriteTextProfile() {
     MIBMap.ForEach(PrintCallback,
                    reinterpret_cast<void *>(flags()->print_terse));
     StackDepotPrintAll();
-    allocator.ForceUnlock();
   }
 
   void InitLinkerInitialized() {
     SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
     allocator.InitLinkerInitialized(
         common_flags()->allocator_release_to_os_interval_ms);
     max_user_defined_malloc_size = common_flags()->max_allocation_size_mb
                                        ? common_flags()->max_allocation_size_mb
int __sanitizer_get_ownership(const void *p) {
   return memprof_malloc_usable_size(p, 0, 0) != 0;
 }
 
 uptr __sanitizer_get_allocated_size(const void *p) {
   return memprof_malloc_usable_size(p, 0, 0);
 }
 
 int __memprof_profile_dump() {
-  instance.FinishAndPrint();
+  instance.FinishAndWrite();
   // In the future we may want to return non-zero if there are any errors
   // detected during the dumping process.
   return 0;
 }

compiler-rt/lib/memprof/memprof_flags.inc

 
 MEMPROF_FLAG(bool, halt_on_error, true,
              "Crash the program after printing the first error report "
              "(WARNING: USE AT YOUR OWN RISK!)")
 MEMPROF_FLAG(bool, allocator_frees_and_returns_null_on_realloc_zero, true,
              "realloc(p, 0) is equivalent to free(p) by default (Same as the "
              "POSIX standard). If set to false, realloc(p, 0) will return a "
              "pointer to an allocated space which can not be used.")
+MEMPROF_FLAG(bool, print_text, true,
+  "If set, prints the heap profile in text format. Else use the raw binary serialization format.")
 MEMPROF_FLAG(bool, print_terse, false,
-             "If set, prints memory profile in a terse format.")
+             "If set, prints memory profile in a terse format. Only applicable if print_text = true.")

compiler-rt/lib/memprof/memprof_rawprofile.h

+#ifndef MEMPROF_RAWPROFILE_H_
+#define MEMPROF_RAWPROFILE_H_
+
+#include "memprof_mibmap.h"
+#include "sanitizer_common/sanitizer_procmaps.h"
+
+namespace __memprof {
+
+// TODO: pull these in from MemProfData.inc
+#define MEMPROF_RAW_MAGIC_64                                                   \
+  (u64)255 << 56 | (u64)'m' << 48 | (u64)'p' << 40 | (u64)'r' << 32 |          \
+      (u64)'o' << 24 | (u64)'f' << 16 | (u64)'r' << 8 | (u64)129
+
+#define MEMPROF_RAW_VERSION 1ULL
+
+int SerializeToRawProfile(MIBMapTy &BlockCache, MemoryMappingLayoutBase &Layout,
+                          char *&Buffer);
+
+} // namespace __memprof
+
+#endif // MEMPROF_RAWPROFILE_H_

compiler-rt/lib/memprof/memprof_rawprofile.cpp

+#include "memprof_rawprofile.h"
+#include "memprof_meminfoblock.h"
+#include "sanitizer_common/sanitizer_linux.h"
+#include "sanitizer_common/sanitizer_procmaps.h"
+#include "sanitizer_common/sanitizer_stackdepot.h"
+#include "sanitizer_common/sanitizer_stackdepotbase.h"
+#include "sanitizer_common/sanitizer_stacktrace.h"
+#include "sanitizer_common/sanitizer_vector.h"
+
+#include <stdlib.h>
+#include <string.h>
+
+namespace __memprof {
+using ::__sanitizer::Vector;
+
+namespace {
+typedef struct __attribute__((__packed__)) {
+  u64 start;
+  u64 end;
+  u64 offset;
+  u8 buildId[32];
+} SegmentEntry;
+
+template <class T> char *WriteBytes(T Pod, char *&Buffer) {
+  // static_assert(std::is_pod<T>::value, "Must be a POD type.");
+  *(T *)Buffer = Pod;
+  return Buffer + sizeof(T);
+}
+
+} // namespace
+
+int SerializeSegments(MemoryMappingLayoutBase &Layout, char *&Buffer) {
+  u32 NumSegmentsToRecord = 0;
+  MemoryMappedSegment segment;
+
+  for (Layout.Reset(); Layout.Next(&segment);)
+    if (segment.IsReadable() && segment.IsExecutable())
+      NumSegmentsToRecord++;
+
+  CHECK(Buffer == nullptr);
+  // Allocate a buffer
+  int NumBytesToWrite =
+      sizeof(u64) // A header which stores the number of records.
+      + sizeof(SegmentEntry) * NumSegmentsToRecord;
+  Buffer = (char *)malloc(NumBytesToWrite);
+
+  char *Ptr = Buffer;
+  *((u64 *)Ptr) = NumSegmentsToRecord;
+  Ptr += sizeof(u64);
+
+  for (Layout.Reset(); Layout.Next(&segment);) {
+    if (segment.IsReadable() && segment.IsExecutable()) {
+      SegmentEntry entry{segment.start, segment.end, segment.offset};
+      memcpy(entry.buildId, segment.uuid, sizeof(segment.uuid));
+      memcpy(Ptr, &entry, sizeof(SegmentEntry));
+      Ptr += sizeof(SegmentEntry);
+    }
+  }
+
+  return NumBytesToWrite;
+}
+
+int SerializeStack(const Vector<u64> &StackIds, char *&Buffer) {
+  int NumBytesToWrite = sizeof(u64);
+
+  const u64 NumIds = StackIds.Size();
+  for (unsigned k = 0; k < NumIds; ++k) {
+    const u64 Id = StackIds[k];
+    // One entry for the id and then one more for the number of stack pcs.
+    NumBytesToWrite += 2 * sizeof(u64);
+    const StackTrace St = StackDepotGet(Id);
+
+    CHECK(St.trace != nullptr && St.size > 0 && "Empty stack trace");
+    for (uptr i = 0; i < St.size && St.trace[i] != 0; i++) {
+      NumBytesToWrite += sizeof(u64);
+    }
+  }
+
+  Buffer = (char *)malloc(NumBytesToWrite);
+  char *Ptr = Buffer;
+  Ptr = WriteBytes(static_cast<u64>(NumIds), Ptr);
+
+  for (unsigned k = 0; k < NumIds; ++k) {
+    const u64 Id = StackIds[k];
+    Ptr = WriteBytes(Id, Ptr);
+    Ptr += sizeof(u64); // Bump it by u64, we will fill this in later.
+    u64 Count = 0;
+    const StackTrace St = StackDepotGet(Id);
+    for (uptr i = 0; i < St.size && St.trace[i] != 0; i++) {
+      // PCs in stack traces are actually the return addresses, that is,
+      // addresses of the next instructions after the call.
+      uptr pc = StackTrace::GetPreviousInstructionPc(St.trace[i]);
+      ++Count;
+      Ptr = WriteBytes(static_cast<u64>(pc), Ptr);
+    }
+    // Store the count in the space we reserved earlier.
+    *(u64 *)(Ptr - (Count + 1) * sizeof(Count)) = Count;
+  }
+
+  return NumBytesToWrite;
+}
+
+static void RecordStackId(const uptr Key, LockedMemInfoBlock *const &MIB,
+                          void *Arg) {
+  // No need to touch the value here since we are only recording the key.
+  auto *StackIds = reinterpret_cast<Vector<u64> *>(Arg);
+  StackIds->PushBack(Key);
+}
+
+int SerializeMIBInfo(MIBMapTy &MIBMap, Vector<u64> &StackIds, char *&Buffer) {
+  CHECK(StackIds.Size() == 0);
+  int NumBytesToWrite = sizeof(u64);
+
+  MIBMap.ForEach(RecordStackId, reinterpret_cast<void *>(&StackIds));
+  const u64 NumEntries = StackIds.Size();
+  NumBytesToWrite += NumEntries * (sizeof(u64) + sizeof(MemInfoBlock));
+
+  Buffer = (char *)malloc(NumBytesToWrite);
+  char *Ptr = Buffer;
+  Ptr = WriteBytes(NumEntries, Ptr);
+
+  for (uptr i = 0; i < StackIds.Size(); i++) {
+    const u64 Key = StackIds[i];
+    MIBMapTy::Handle h(&MIBMap, Key, /*remove=*/true, /*create=*/false);
+    CHECK(h.exists());
+    Ptr = WriteBytes(Key, Ptr);
+    Ptr = WriteBytes((*h)->mib, Ptr);
+  }
+
+  return NumBytesToWrite;
+}
+
+// Format
+// ---------- Header
+// Magic
+// Version
+// Total Size
+// Segment Offset
+// MIB Info Offset
+// Stack Offset
+// ---------- Segment Info
+// Num Entries
+// ---------- Segment Entry
+// Start
+// End
+// Offset
+// BuildID 32B
+// ----------
+// ...
+// ---------- MIB Info
+// Num Entries
+// ---------- MIB Entry
+// Alloc Count
+// ...
+// ---------- Stack Info
+// Num Entries
+// ---------- Stack Entry
+// Num Stacks
+// PC1
+// PC2
+// ...
+// ----------
+// ...
+int SerializeToRawProfile(MIBMapTy &MIBMap, MemoryMappingLayoutBase &Layout,
+                          char *&Buffer) {
+  char *SegmentBuffer = nullptr;
+  int NumSegmentBytes = SerializeSegments(Layout, SegmentBuffer);
+
+  char *MIBInfoBuffer = nullptr;
+  Vector<u64> StackIds;
+  int NumMIBInfoBytes = SerializeMIBInfo(MIBMap, StackIds, MIBInfoBuffer);
+
+  char *StackBuffer = nullptr;
+  int NumStackBytes = SerializeStack(StackIds, StackBuffer);
+
+  // Compute how many bytes we need for the header.
+  int NumBytesHeader = sizeof(MEMPROF_RAW_MAGIC_64) +
+                       sizeof(MEMPROF_RAW_VERSION) +
+                       4 * sizeof(u64); // Offset into Segment and Stack
+                                        // sections, MIB and total size.
+
+  int TotalSizeBytes =
+      NumBytesHeader + NumSegmentBytes + NumStackBytes + NumMIBInfoBytes;
+
+  // Allocate the memory for the entire buffer incl. info blocks.
+  Buffer = (char *)malloc(TotalSizeBytes);
+  char *Ptr = Buffer;
+
+  // Write Header
+  Ptr = WriteBytes(MEMPROF_RAW_MAGIC_64, Ptr);
+  Ptr = WriteBytes(MEMPROF_RAW_VERSION, Ptr);
+  Ptr = WriteBytes(static_cast<u64>(TotalSizeBytes), Ptr);
+  Ptr = WriteBytes(static_cast<u64>(NumBytesHeader), Ptr);
+  Ptr = WriteBytes(static_cast<u64>(NumBytesHeader + NumSegmentBytes), Ptr);
+  Ptr = WriteBytes(
+      static_cast<u64>(NumBytesHeader + NumSegmentBytes + NumMIBInfoBytes),
+      Ptr);
+
+  memcpy(Ptr, SegmentBuffer, NumSegmentBytes);
+  Ptr += NumSegmentBytes;
+  memcpy(Ptr, MIBInfoBuffer, NumMIBInfoBytes);
+  Ptr += NumMIBInfoBytes;
+  memcpy(Ptr, StackBuffer, NumStackBytes);
+  // We don't need Ptr anymore.. just return the total size.
+
+  return TotalSizeBytes;
+}
+
+} // namespace __memprof

compiler-rt/lib/memprof/tests/CMakeLists.txt

+include(CheckCXXCompilerFlag)
+include(CompilerRTCompile)
+include(CompilerRTLink)
+
+set(MEMPROF_UNITTEST_CFLAGS
+  ${COMPILER_RT_UNITTEST_CFLAGS}
+  ${COMPILER_RT_GTEST_CFLAGS}
+  ${COMPILER_RT_GMOCK_CFLAGS}
+  -I${COMPILER_RT_SOURCE_DIR}/lib/
+  -O2
+  -g
+  -fno-rtti
+  -Wno-gnu-zero-variadic-macro-arguments
+  -fno-omit-frame-pointer)
+
+file(GLOB MEMPROF_HEADERS ../*.h)
+
+set(MEMPROF_SOURCES
+  ../memprof_mibmap.cpp
+  ../memprof_rawprofile.cpp)
+
+set(MEMPROF_UNITTESTS
+  rawprofile.cpp
+  driver.cpp)
+
+set(MEMPROF_UNIT_TEST_HEADERS
+  ${MEMPROF_HEADERS})
+
+if(NOT WIN32)
+  list(APPEND MEMPROF_UNITTEST_LINK_FLAGS -pthread)
+endif()
+
+if(COMPILER_RT_DEFAULT_TARGET_ARCH IN_LIST MEMPROF_SUPPORTED_ARCH)
+  # MemProf unit tests are only run on the host machine.
+  set(arch ${COMPILER_RT_DEFAULT_TARGET_ARCH})
+
+  add_executable(MemProfUnitTests 
+    ${MEMPROF_UNITTESTS}
+    ${COMPILER_RT_GTEST_SOURCE}
+    ${COMPILER_RT_GMOCK_SOURCE}
+    ${MEMPROF_SOURCES}
+    $<TARGET_OBJECTS:RTSanitizerCommon.${arch}>
+    $<TARGET_OBJECTS:RTSanitizerCommonCoverage.${arch}>
+    $<TARGET_OBJECTS:RTSanitizerCommonLibc.${arch}>
+    $<TARGET_OBJECTS:RTSanitizerCommonSymbolizer.${arch}>)
+  set_target_compile_flags(MemProfUnitTests ${MEMPROF_UNITTEST_CFLAGS})
+  set_target_link_flags(MemProfUnitTests ${MEMPROF_UNITTEST_LINK_FLAGS})
+  target_link_libraries(MemProfUnitTests dl) 
+
+  set_target_properties(MemProfUnitTests PROPERTIES
+    RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+endif()

compiler-rt/lib/memprof/tests/driver.cpp

+//===-- driver.cpp ----------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "gtest/gtest.h"
+
+int main(int argc, char **argv) {
+  testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

compiler-rt/lib/memprof/tests/rawprofile.cpp

+#include "memprof/memprof_rawprofile.h"
+
+#include "memprof/memprof_meminfoblock.h"
+#include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_procmaps.h"
+#include "sanitizer_common/sanitizer_stackdepot.h"
+#include "sanitizer_common/sanitizer_stacktrace.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include <memory>
+
+namespace {
+
+using ::__memprof::MemInfoBlock;
+using ::__memprof::MIBMapTy;
+using ::__memprof::SerializeToRawProfile;
+using ::__sanitizer::MemoryMappedSegment;
+using ::__sanitizer::MemoryMappingLayoutBase;
+using ::__sanitizer::StackDepotPut;
+using ::__sanitizer::StackTrace;
+using ::testing::_;
+using ::testing::Action;
+using ::testing::DoAll;
+using ::testing::Return;
+using ::testing::SetArgPointee;
+
+class MockMemoryMappingLayout final : public MemoryMappingLayoutBase {
+public:
+  MOCK_METHOD(bool, Next, (MemoryMappedSegment *), (override));
+  MOCK_METHOD(void, Reset, (), (override));
+};
+
+std::unique_ptr<MIBMapTy> MakeFakeMap() {
+  MemInfoBlock FakeMIB;
+  memset(&FakeMIB, 0, sizeof(MemInfoBlock));
+  FakeMIB.alloc_count = 0x1;
+  FakeMIB.total_access_count = 0x2;
+
+  // Create a fake stack trace, start from 2 since we deduct 1 to get prior pc
+  // address.
+  uptr array[] = {2, 3, 4, 5, 6};
+  StackTrace St(array, ARRAY_SIZE(array));
+  u32 Id = StackDepotPut(St);
+
+  auto FakeMap = std::make_unique<MIBMapTy>();
+  InsertOrMerge(Id, FakeMIB, *FakeMap);
+
+  return FakeMap;
+}
+
+template <class T = u64> T Read(char *&Buffer) {
+  static_assert(std::is_pod<T>::value, "Must be a POD type.");
+  T t = *reinterpret_cast<T *>(Buffer);
+  Buffer += sizeof(T);
+  return t;
+}
+
+TEST(MemProf, Basic) {
+  MockMemoryMappingLayout Layout;
+  MemoryMappedSegment FakeSegment;
+  memset(&FakeSegment, 0, sizeof(FakeSegment));
+  FakeSegment.start = 0x10;
+  FakeSegment.end = 0x20;
+  FakeSegment.offset = 0x20;
+  uint8_t uuid[__sanitizer::kModuleUUIDSize] = {0xC, 0x0, 0xF, 0xF, 0xE, 0xE};
+  memcpy(FakeSegment.uuid, uuid, __sanitizer::kModuleUUIDSize);
+  FakeSegment.protection =
+      __sanitizer::kProtectionExecute | __sanitizer::kProtectionRead;
+
+  const Action<bool(MemoryMappedSegment *)> SetSegment =
+      DoAll(SetArgPointee<0>(FakeSegment), Return(true));
+  EXPECT_CALL(Layout, Next(_))
+      .WillOnce(SetSegment)
+      .WillOnce(Return(false))
+      .WillOnce(SetSegment)
+      .WillRepeatedly(Return(false));
+
+  EXPECT_CALL(Layout, Reset).Times(2);
+
+  auto FakeMap = MakeFakeMap();
+  char *Ptr = nullptr;
+  int NumBytes = SerializeToRawProfile(*FakeMap, Layout, Ptr);
+  const char *Buffer = Ptr;
+
+  ASSERT_GT(NumBytes, 0);
+  ASSERT_TRUE(Ptr);
+
+  // Check the header.
+  EXPECT_THAT(Read(Ptr), MEMPROF_RAW_MAGIC_64);
+  EXPECT_THAT(Read(Ptr), MEMPROF_RAW_VERSION);
+  const u64 TotalSize = Read(Ptr);
+  const u64 SegmentOffset = Read(Ptr);
+  const u64 MIBOffset = Read(Ptr);
+  const u64 StackOffset = Read(Ptr);
+
+  // ============= Check sizes.
+
+  EXPECT_GT(TotalSize, 0ULL);
+
+  EXPECT_GT(SegmentOffset, 0ULL);
+  // We expect only 1 segment entry, 8b for the count and 56b for SegmentEntry
+  // in memprof_rawprofile.cpp.
+  EXPECT_EQ(MIBOffset - SegmentOffset, 64ULL);
+
+  EXPECT_GT(MIBOffset, 0ULL);
+  // We expect 1 mib entry, 8b for the count and sizeof(u64) +
+  // sizeof(MemInfoBlock) contains stack id + MeminfoBlock.
+  EXPECT_EQ(StackOffset - MIBOffset, 8 + 8 + sizeof(MemInfoBlock));
+
+  EXPECT_GT(StackOffset, 0ULL);
+  // We expect 1 stack entries, with 5 frames - 8b for total count, 8b for
+  // id, 8b for frame count and 5*8b for fake frame.
+  EXPECT_EQ(TotalSize - StackOffset, 8ULL + (8 + 8 + 5 * 8));
+
+  // ============= Check contents.
+  unsigned char ExpectedSegmentBytes[64] = {
+      0x01, 0,   0,   0,   0,   0,   0, 0, // Number of entries
+      0x10, 0,   0,   0,   0,   0,   0, 0, // Start
+      0x20, 0,   0,   0,   0,   0,   0, 0, // End
+      0x20, 0,   0,   0,   0,   0,   0, 0, // Offset
+      0x0C, 0x0, 0xF, 0xF, 0xE, 0xE,       // Uuid
+  };
+  EXPECT_EQ(memcmp(Buffer + SegmentOffset, ExpectedSegmentBytes, 64), 0);
+
+  // Check that the number of entries is 1.
+  EXPECT_EQ(*reinterpret_cast<const u64 *>(Buffer + MIBOffset), 1ULL);
+  // Check that stack id is set.
+  EXPECT_NE(*reinterpret_cast<const u64 *>(Buffer + MIBOffset + 8), 0ULL);
+
+  // Only check a few fields of the first MemInfoBlock.
+  unsigned char ExpectedMIBBytes[sizeof(MemInfoBlock)] = {
+      0x01, 0, 0, 0, // Alloc count
+      0x02, 0, 0, 0, // Total access count
+  };
+  // Compare contents after skipping count and stack id.
+  EXPECT_EQ(
+      memcmp(Buffer + MIBOffset + 16, ExpectedMIBBytes, sizeof(MemInfoBlock)),
+      0);
+
+  // Check that the number of entries is 1.
+  EXPECT_EQ(*reinterpret_cast<const u64 *>(Buffer + StackOffset), 1ULL);
+  // Check that the stack id is not null.
+  EXPECT_NE(*reinterpret_cast<const u64 *>(Buffer + StackOffset + 8), 0ULL);
+  // Contents are num pcs, value of each pc - 1.
+  unsigned char ExpectedStackBytes[6 * 8] = {
+      0x5, 0, 0, 0, 0, 0, 0, 0, // Number of PCs
+      0x1, 0, 0, 0, 0, 0, 0, 0, // PC ...
+      0x2, 0, 0, 0, 0, 0, 0, 0, 0x3, 0, 0, 0, 0, 0, 0, 0,
+      0x4, 0, 0, 0, 0, 0, 0, 0, 0x5, 0, 0, 0, 0, 0, 0, 0,
+  };
+  EXPECT_EQ(memcmp(Buffer + StackOffset + 16, ExpectedStackBytes,
+                   sizeof(ExpectedStackBytes)),
+            0);
+}
+
+} // namespace

compiler-rt/lib/sanitizer_common/sanitizer_procmaps.h

 
  private:
   friend class MemoryMappingLayout;
 
   // This field is assigned and owned by MemoryMappingLayout if needed
   MemoryMappedSegmentData *data_;
 };
 
-class MemoryMappingLayout {
+class MemoryMappingLayoutBase {
+ public:
+  virtual bool Next(MemoryMappedSegment *segment) { UNIMPLEMENTED(); }
+  virtual bool Error() const { UNIMPLEMENTED(); };
+  virtual void Reset() { UNIMPLEMENTED(); }
+
+ protected:
+  ~MemoryMappingLayoutBase() {}
+};
+
+class MemoryMappingLayout final : public MemoryMappingLayoutBase {
  public:
   explicit MemoryMappingLayout(bool cache_enabled);
   ~MemoryMappingLayout();
-  bool Next(MemoryMappedSegment *segment);
-  bool Error() const;
-  void Reset();
+  virtual bool Next(MemoryMappedSegment *segment) override;
+  virtual bool Error() const override;
+  virtual void Reset() override;
   // In some cases, e.g. when running under a sandbox on Linux, ASan is unable
   // to obtain the memory mappings. It should fall back to pre-cached data
   // instead of aborting.
   static void CacheMemoryMappings();
 
   // Adds all mapped objects into a vector.
   void DumpListOfModules(InternalMmapVectorNoCtor<LoadedModule> *modules);
 

compiler-rt/test/memprof/TestCases/interface_test.cpp

 // Check that user may include MemProf interface header.
 // Also check that interfaces declared in the sanitizer's allocator_interface
 // are defined for MemProf.
-// RUN: %clang_memprof %s -o %t -DMEMPROF && %run %t
-// RUN: %clang_memprof -x c %s -o %t -DMEMPROF && %run %t
+// RUN: %clangxx_memprof %s -o %t -DMEMPROF && %run %t
+// RUN: %clangxx_memprof -x c %s -o %t -DMEMPROF && %run %t
 // RUN: %clang %s -pie -o %t && %run %t
 // RUN: %clang -x c %s -pie -o %t && %run %t
 #include <sanitizer/allocator_interface.h>
 #include <sanitizer/memprof_interface.h>
 #include <stdlib.h>
 
 int main() {
   int *p = (int *)malloc(10 * sizeof(int));

compiler-rt/test/memprof/TestCases/memprof_profile_dump.cpp

 // RUN: %clangxx_memprof  %s -o %t
 
-// RUN: %env_memprof_opts=log_path=stdout %run %t | FileCheck %s
+// RUN: %env_memprof_opts=log_path=stdout %run %t | FileCheck --check-prefix=CHECK-TEXT %s
+// RUN: %env_memprof_opts=log_path=stdout,print_text=false %run %t > %t.memprofraw
+// RUN: xxd -l 8 %t.memprofraw | FileCheck --check-prefix=CHECK-RAW %s
 
 #include <sanitizer/memprof_interface.h>
 #include <stdlib.h>
 #include <string.h>
 int main(int argc, char **argv) {
   char *x = (char *)malloc(10);
   memset(x, 0, 10);
   free(x);
   __memprof_profile_dump();
   x = (char *)malloc(10);
   memset(x, 0, 10);
   free(x);
   return 0;
 }
 // We should get 2 rounds of profile info, one from the explicit dump request,
 // and one at exit.
-// CHECK: Memory allocation stack id
-// CHECK: Stack for id
-// CHECK: Memory allocation stack id
-// CHECK: Stack for id
+// CHECK-TEXT: Memory allocation stack id
+// CHECK-TEXT: Stack for id
+// CHECK-TEXT: Memory allocation stack id
+// CHECK-TEXT: Stack for id
+//
+// For the raw profile just check the header.
+// CHECK-RAW: 00000000: 8172 666f 7270 6dff                      .rforpm.