[MemProf] Change histogram storage from uint64_t to uint16_t (#147854)

Change memory access histogram storage from uint64_t to uint16_t to
reduce profile size on disk. This change updates the raw profile format
to v5. Also add a histogram test in compiler-rt since we didn't have one
before. With this change the histogram memprof raw for the basic test
reduces from 75KB -> 20KB.

Author: Snehasish Kumar

compiler-rt/include/profile/MemProfData.inc

@@ -33,11 +33,10 @@
    (uint64_t)'o' << 24 | (uint64_t)'f' << 16 | (uint64_t)'r' << 8 | (uint64_t)129)
 
 // The version number of the raw binary format.
-#define MEMPROF_RAW_VERSION 4ULL
+#define MEMPROF_RAW_VERSION 5ULL
 
 // Currently supported versions.
-#define MEMPROF_RAW_SUPPORTED_VERSIONS                                         \
-  { 3ULL, 4ULL }
+#define MEMPROF_RAW_SUPPORTED_VERSIONS {3ULL, 4ULL, 5ULL}
 
 #define MEMPROF_V3_MIB_SIZE 132ULL;
 
@@ -229,6 +228,41 @@ void Merge(const MemInfoBlock &newMIB) {
 } __attribute__((__packed__));
 #endif
 
+constexpr int MantissaBits = 12;
+constexpr int ExponentBits = 4;
+constexpr uint16_t MaxMantissa = (1U << MantissaBits) - 1;
+constexpr uint16_t MaxExponent = (1U << ExponentBits) - 1;
+constexpr uint64_t MaxRepresentableValue = static_cast<uint64_t>(MaxMantissa)
+                                           << MaxExponent;
+
+// Encodes a 64-bit unsigned integer into a 16-bit scaled integer format.
+inline uint16_t encodeHistogramCount(uint64_t Count) {
+  if (Count == 0)
+    return 0;
+
+  if (Count > MaxRepresentableValue)
+    Count = MaxRepresentableValue;
+
+  if (Count <= MaxMantissa)
+    return Count;
+
+  uint64_t M = Count;
+  uint16_t E = 0;
+  while (M > MaxMantissa) {
+    M = (M + 1) >> 1;
+    E++;
+  }
+  return (E << MantissaBits) | static_cast<uint16_t>(M);
+}
+
+// Decodes a 16-bit scaled integer and returns the
+// decoded 64-bit unsigned integer.
+inline uint64_t decodeHistogramCount(uint16_t EncodedValue) {
+  const uint16_t E = EncodedValue >> MantissaBits;
+  const uint16_t M = EncodedValue & MaxMantissa;
+  return static_cast<uint64_t>(M) << E;
+}
+
 } // namespace memprof
 } // namespace llvm
 

compiler-rt/lib/memprof/memprof_rawprofile.cpp

@@ -19,6 +19,7 @@ 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) {
@@ -169,13 +170,15 @@ 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) &&
@@ -249,7 +252,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);
 

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)
 

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

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

@@ -0,0 +1,37 @@
+// Test the histogram support in memprof using the text format output.
+// Shadow memory counters per object are limited to 8b. In memory counters
+// aggregating counts across multiple objects are 64b.
+
+// RUN: %clangxx_memprof -O0 -mllvm -memprof-histogram -mllvm -memprof-use-callbacks=true %s -o %t && %env_memprof_opts=print_text=1:histogram=1:log_path=stdout %run %t 2>&1 | FileCheck %s
+
+#include <stdio.h>
+#include <stdlib.h>
+
+int main() {
+  // Allocate memory that will create a histogram
+  char *buffer = (char *)malloc(1024);
+  if (!buffer)
+    return 1;
+
+  for (int i = 0; i < 10; ++i) {
+    // Access every 8th byte (since shadow granularity is 8b.
+    buffer[i * 8] = 'A';
+  }
+
+  for (int j = 0; j < 200; ++j) {
+    buffer[8] = 'B'; // Count = previous count + 200
+  }
+
+  for (int j = 0; j < 400; ++j) {
+    buffer[16] = 'B'; // Count is saturated at 255
+  }
+
+  // Free the memory to trigger MIB creation with histogram
+  free(buffer);
+
+  printf("Test completed successfully\n");
+  return 0;
+}
+
+// CHECK: AccessCountHistogram[128]: 1 201 255 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+// CHECK: Test completed successfully

llvm/include/llvm/ProfileData/MemProfData.inc

@@ -33,11 +33,10 @@
    (uint64_t)'o' << 24 | (uint64_t)'f' << 16 | (uint64_t)'r' << 8 | (uint64_t)129)
 
 // The version number of the raw binary format.
-#define MEMPROF_RAW_VERSION 4ULL
+#define MEMPROF_RAW_VERSION 5ULL
 
 // Currently supported versions.
-#define MEMPROF_RAW_SUPPORTED_VERSIONS                                         \
-  { 3ULL, 4ULL }
+#define MEMPROF_RAW_SUPPORTED_VERSIONS {3ULL, 4ULL, 5ULL}
 
 #define MEMPROF_V3_MIB_SIZE 132ULL;
 
@@ -229,6 +228,41 @@ void Merge(const MemInfoBlock &newMIB) {
 } __attribute__((__packed__));
 #endif
 
+constexpr int MantissaBits = 12;
+constexpr int ExponentBits = 4;
+constexpr uint16_t MaxMantissa = (1U << MantissaBits) - 1;
+constexpr uint16_t MaxExponent = (1U << ExponentBits) - 1;
+constexpr uint64_t MaxRepresentableValue = static_cast<uint64_t>(MaxMantissa)
+                                           << MaxExponent;
+
+// Encodes a 64-bit unsigned integer into a 16-bit scaled integer format.
+inline uint16_t encodeHistogramCount(uint64_t Count) {
+  if (Count == 0)
+    return 0;
+
+  if (Count > MaxRepresentableValue)
+    Count = MaxRepresentableValue;
+
+  if (Count <= MaxMantissa)
+    return Count;
+
+  uint64_t M = Count;
+  uint16_t E = 0;
+  while (M > MaxMantissa) {
+    M = (M + 1) >> 1;
+    E++;
+  }
+  return (E << MantissaBits) | static_cast<uint16_t>(M);
+}
+
+// Decodes a 16-bit scaled integer and returns the
+// decoded 64-bit unsigned integer.
+inline uint64_t decodeHistogramCount(uint16_t EncodedValue) {
+  const uint16_t E = EncodedValue >> MantissaBits;
+  const uint16_t M = EncodedValue & MaxMantissa;
+  return static_cast<uint64_t>(M) << E;
+}
+
 } // namespace memprof
 } // namespace llvm
 

llvm/lib/ProfileData/MemProfReader.cpp

@@ -135,7 +135,7 @@ readMemInfoBlocksV3(const char *Ptr) {
 }
 
 llvm::SmallVector<std::pair<uint64_t, MemInfoBlock>>
-readMemInfoBlocksV4(const char *Ptr) {
+readMemInfoBlocksCommon(const char *Ptr, bool IsHistogramEncoded = false) {
   using namespace support;
 
   const uint64_t NumItemsToRead =
@@ -145,27 +145,43 @@ readMemInfoBlocksV4(const char *Ptr) {
   for (uint64_t I = 0; I < NumItemsToRead; I++) {
     const uint64_t Id =
         endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
-    // We cheat a bit here and remove the const from cast to set the
-    // Histogram Pointer to newly allocated buffer.
-    MemInfoBlock MIB = *reinterpret_cast<const MemInfoBlock *>(Ptr);
 
-    // Only increment by size of MIB since readNext implicitly increments.
+    MemInfoBlock MIB = *reinterpret_cast<const MemInfoBlock *>(Ptr);
     Ptr += sizeof(MemInfoBlock);
 
     if (MIB.AccessHistogramSize > 0) {
+      // The in-memory representation uses uint64_t for histogram entries.
       MIB.AccessHistogram =
           (uintptr_t)malloc(MIB.AccessHistogramSize * sizeof(uint64_t));
-    }
-
-    for (uint64_t J = 0; J < MIB.AccessHistogramSize; J++) {
-      ((uint64_t *)MIB.AccessHistogram)[J] =
-          endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
+      for (uint64_t J = 0; J < MIB.AccessHistogramSize; J++) {
+        if (!IsHistogramEncoded) {
+          ((uint64_t *)MIB.AccessHistogram)[J] =
+              endian::readNext<uint64_t, llvm::endianness::little, unaligned>(
+                  Ptr);
+        } else {
+          // The encoded on-disk format (V5 onwards) uses uint16_t.
+          const uint16_t Val =
+              endian::readNext<uint16_t, llvm::endianness::little, unaligned>(
+                  Ptr);
+          ((uint64_t *)MIB.AccessHistogram)[J] = decodeHistogramCount(Val);
+        }
+      }
     }
     Items.push_back({Id, MIB});
   }
   return Items;
 }
 
+llvm::SmallVector<std::pair<uint64_t, MemInfoBlock>>
+readMemInfoBlocksV4(const char *Ptr) {
+  return readMemInfoBlocksCommon(Ptr);
+}
+
+llvm::SmallVector<std::pair<uint64_t, MemInfoBlock>>
+readMemInfoBlocksV5(const char *Ptr) {
+  return readMemInfoBlocksCommon(Ptr, /*IsHistogramEncoded=*/true);
+}
+
 CallStackMap readStackInfo(const char *Ptr) {
   using namespace support;
 
@@ -658,6 +674,8 @@ RawMemProfReader::readMemInfoBlocks(const char *Ptr) {
     return readMemInfoBlocksV3(Ptr);
   if (MemprofRawVersion == 4ULL)
     return readMemInfoBlocksV4(Ptr);
+  if (MemprofRawVersion == 5ULL)
+    return readMemInfoBlocksV5(Ptr);
   llvm_unreachable(
       "Panic: Unsupported version number when reading MemInfoBlocks");
 }