4 files changed, 66 insertions, 12 deletions
diff --git a/compiler-rt/lib/memprof/memprof_interface_internal.h b/compiler-rt/lib/memprof/memprof_interface_internal.h
index 7d3a937..1fd0748 100644
--- a/compiler-rt/lib/memprof/memprof_interface_internal.h
+++ b/compiler-rt/lib/memprof/memprof_interface_internal.h
@@ -36,8 +36,14 @@ SANITIZER_INTERFACE_ATTRIBUTE
 void __memprof_record_access(void const volatile *addr);
 
 SANITIZER_INTERFACE_ATTRIBUTE
+void __memprof_record_access_hist(void const volatile *addr);
+
+SANITIZER_INTERFACE_ATTRIBUTE
 void __memprof_record_access_range(void const volatile *addr, uptr size);
 
+SANITIZER_INTERFACE_ATTRIBUTE
+void __memprof_record_access_range_hist(void const volatile *addr, uptr size);
+
 SANITIZER_INTERFACE_ATTRIBUTE void __memprof_print_accumulated_stats();
 
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE extern char
@@ -51,6 +57,10 @@ extern uptr __memprof_shadow_memory_dynamic_address;
 
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE extern char
     __memprof_profile_filename[1];
+
+SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE extern bool
+    __memprof_histogram;
+
 SANITIZER_INTERFACE_ATTRIBUTE int __memprof_profile_dump();
 SANITIZER_INTERFACE_ATTRIBUTE void __memprof_profile_reset();
 
diff --git a/compiler-rt/lib/memprof/memprof_rawprofile.cpp b/compiler-rt/lib/memprof/memprof_rawprofile.cpp
index a897648..f579e12 100644
--- a/compiler-rt/lib/memprof/memprof_rawprofile.cpp
+++ b/compiler-rt/lib/memprof/memprof_rawprofile.cpp
@@ -7,10 +7,7 @@
 #include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_array_ref.h"
 #include "sanitizer_common/sanitizer_common.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"
 
@@ -19,10 +16,20 @@ using ::__sanitizer::Vector;
 using ::llvm::memprof::MemInfoBlock;
 using SegmentEntry = ::llvm::memprof::SegmentEntry;
 using Header = ::llvm::memprof::Header;
+using ::llvm::memprof::encodeHistogramCount;
 
 namespace {
 template <class T> char *WriteBytes(const T &Pod, char *Buffer) {
-  *(T *)Buffer = Pod;
+  static_assert(is_trivially_copyable<T>::value, "T must be POD");
+  const uint8_t *Src = reinterpret_cast<const uint8_t *>(&Pod);
+
+  for (size_t I = 0; I < sizeof(T); ++I)
+#if defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+    // Reverse byte order since reader is little-endian.
+    Buffer[I] = Src[sizeof(T) - 1 - I];
+#else
+    Buffer[I] = Src[I];
+#endif
   return Buffer + sizeof(T);
 }
 
@@ -32,7 +39,6 @@ void RecordStackId(const uptr Key, UNUSED LockedMemInfoBlock *const &MIB,
   auto *StackIds = reinterpret_cast<Vector<u64> *>(Arg);
   StackIds->PushBack(Key);
 }
-} // namespace
 
 u64 SegmentSizeBytes(ArrayRef<LoadedModule> Modules) {
   u64 NumSegmentsToRecord = 0;
@@ -169,18 +175,21 @@ void SerializeMIBInfoToBuffer(MIBMapTy &MIBMap, const Vector<u64> &StackIds,
     // FIXME: We unnecessarily serialize the AccessHistogram pointer. Adding a
     // serialization schema will fix this issue. See also FIXME in
     // deserialization.
-    Ptr = WriteBytes((*h)->mib, Ptr);
-    for (u64 j = 0; j < (*h)->mib.AccessHistogramSize; ++j) {
-      u64 HistogramEntry = ((u64 *)((*h)->mib.AccessHistogram))[j];
+    auto &MIB = (*h)->mib;
+    Ptr = WriteBytes(MIB, Ptr);
+    for (u64 j = 0; j < MIB.AccessHistogramSize; ++j) {
+      u16 HistogramEntry =
+          encodeHistogramCount(((u64 *)(MIB.AccessHistogram))[j]);
       Ptr = WriteBytes(HistogramEntry, Ptr);
     }
-    if ((*h)->mib.AccessHistogramSize > 0) {
-      InternalFree((void *)((*h)->mib.AccessHistogram));
+    if (MIB.AccessHistogramSize > 0) {
+      InternalFree((void *)MIB.AccessHistogram);
     }
   }
   CHECK(ExpectedNumBytes >= static_cast<u64>(Ptr - Buffer) &&
         "Expected num bytes != actual bytes written");
 }
+} // namespace
 
 // Format
 // ---------- Header
@@ -249,7 +258,7 @@ u64 SerializeToRawProfile(MIBMapTy &MIBMap, ArrayRef<LoadedModule> Modules,
       },
       reinterpret_cast<void *>(&TotalAccessHistogramEntries));
   const u64 NumHistogramBytes =
-      RoundUpTo(TotalAccessHistogramEntries * sizeof(uint64_t), 8);
+      RoundUpTo(TotalAccessHistogramEntries * sizeof(uint16_t), 8);
 
   const u64 NumStackBytes = RoundUpTo(StackSizeBytes(StackIds), 8);
 
@@ -285,5 +294,4 @@ u64 SerializeToRawProfile(MIBMapTy &MIBMap, ArrayRef<LoadedModule> Modules,
 
   return TotalSizeBytes;
 }
-
 } // namespace __memprof
diff --git a/compiler-rt/lib/memprof/tests/CMakeLists.txt b/compiler-rt/lib/memprof/tests/CMakeLists.txt
index 0b5c302..1603d47 100644
--- a/compiler-rt/lib/memprof/tests/CMakeLists.txt
+++ b/compiler-rt/lib/memprof/tests/CMakeLists.txt
@@ -26,6 +26,7 @@ set(MEMPROF_SOURCES
   ../memprof_rawprofile.cpp)
 
 set(MEMPROF_UNITTESTS
+  histogram_encoding.cpp
   rawprofile.cpp
   driver.cpp)
 
diff --git a/compiler-rt/lib/memprof/tests/histogram_encoding.cpp b/compiler-rt/lib/memprof/tests/histogram_encoding.cpp
new file mode 100644
index 0000000..be20595
--- /dev/null
+++ b/compiler-rt/lib/memprof/tests/histogram_encoding.cpp
@@ -0,0 +1,35 @@
+#include <cstdint>
+#include <vector>
+
+#include "profile/MemProfData.inc"
+#include "gtest/gtest.h"
+
+namespace llvm {
+namespace memprof {
+namespace {
+TEST(MemProf, F16EncodeDecode) {
+  const std::vector<uint64_t> TestCases = {
+      0, 100, 4095, 4096, 5000, 8191, 65535, 1000000, 134213640, 200000000,
+  };
+
+  for (const uint64_t TestCase : TestCases) {
+    const uint16_t Encoded = encodeHistogramCount(TestCase);
+    const uint64_t Decoded = decodeHistogramCount(Encoded);
+
+    const uint64_t MaxRepresentable = static_cast<uint64_t>(MaxMantissa)
+                                      << MaxExponent;
+
+    if (TestCase >= MaxRepresentable) {
+      EXPECT_EQ(Decoded, MaxRepresentable);
+    } else if (TestCase <= MaxMantissa) {
+      EXPECT_EQ(Decoded, TestCase);
+    } else {
+      // The decoded value should be close to the original value.
+      // The error should be less than 1/1024 for larger numbers.
+      EXPECT_NEAR(Decoded, TestCase, static_cast<double>(TestCase) / 1024.0);
+    }
+  }
+}
+} // namespace
+} // namespace memprof
+} // namespace llvm