[MemProf] Change histogram storage from uint64_t to uint8_t

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;
 
@@ -219,7 +218,11 @@ void Merge(const MemInfoBlock &newMIB) {
     ShorterHistogramSize = newMIB.AccessHistogramSize;
   }
   for (size_t i = 0; i < ShorterHistogramSize; ++i) {
-    ((uint64_t *)AccessHistogram)[i] += ((uint64_t *)ShorterHistogram)[i];
+    // Cast to uint8_t* and cap the sum at 255 to prevent overflow
+    uint8_t *CurrentHistPtr = (uint8_t *)AccessHistogram;
+    uint8_t *ShorterHistPtr = (uint8_t *)ShorterHistogram;
+    uint32_t sum = CurrentHistPtr[i] + ShorterHistPtr[i];
+    CurrentHistPtr[i] = (sum > 255) ? 255 : (uint8_t)sum;
   }
 }
 

compiler-rt/lib/memprof/memprof_allocator.cpp

@@ -77,7 +77,7 @@ void Print(const MemInfoBlock &M, const u64 id, bool print_terse) {
                              ? MAX_HISTOGRAM_PRINT_SIZE
                              : M.AccessHistogramSize;
     for (size_t i = 0; i < PrintSize; ++i) {
-      Printf("%llu ", ((uint64_t *)M.AccessHistogram)[i]);
+      Printf("%u ", ((uint8_t *)M.AccessHistogram)[i]);
     }
     Printf("\n");
   }
@@ -327,12 +327,14 @@ struct Allocator {
     uint32_t HistogramSize =
         RoundUpTo(user_size, HISTOGRAM_GRANULARITY) / HISTOGRAM_GRANULARITY;
     uintptr_t Histogram =
-        (uintptr_t)InternalAlloc(HistogramSize * sizeof(uint64_t));
-    memset((void *)Histogram, 0, HistogramSize * sizeof(uint64_t));
+        (uintptr_t)InternalAlloc(HistogramSize * sizeof(uint8_t));
+    memset((void *)Histogram, 0, HistogramSize * sizeof(uint8_t));
     for (size_t i = 0; i < HistogramSize; ++i) {
       u8 Counter =
           *((u8 *)HISTOGRAM_MEM_TO_SHADOW(p + HISTOGRAM_GRANULARITY * i));
-      ((uint64_t *)Histogram)[i] = (uint64_t)Counter;
+      // Cap the counter at HISTOGRAM_MAX_COUNTER (255) to prevent overflow
+      ((uint8_t *)Histogram)[i] =
+          (Counter > HISTOGRAM_MAX_COUNTER) ? HISTOGRAM_MAX_COUNTER : Counter;
     }
     MemInfoBlock newMIB(user_size, c, m->timestamp_ms, curtime, m->cpu_id,
                         GetCpuId(), Histogram, HistogramSize);

compiler-rt/lib/memprof/memprof_rawprofile.cpp

@@ -171,7 +171,8 @@ void SerializeMIBInfoToBuffer(MIBMapTy &MIBMap, const Vector<u64> &StackIds,
     // deserialization.
     Ptr = WriteBytes((*h)->mib, Ptr);
     for (u64 j = 0; j < (*h)->mib.AccessHistogramSize; ++j) {
-      u64 HistogramEntry = ((u64 *)((*h)->mib.AccessHistogram))[j];
+      // Read as uint8_t and write as uint8_t
+      uint8_t HistogramEntry = ((uint8_t *)((*h)->mib.AccessHistogram))[j];
       Ptr = WriteBytes(HistogramEntry, Ptr);
     }
     if ((*h)->mib.AccessHistogramSize > 0) {
@@ -249,7 +250,7 @@ u64 SerializeToRawProfile(MIBMapTy &MIBMap, ArrayRef<LoadedModule> Modules,
       },
       reinterpret_cast<void *>(&TotalAccessHistogramEntries));
   const u64 NumHistogramBytes =
-      RoundUpTo(TotalAccessHistogramEntries * sizeof(uint64_t), 8);
+      RoundUpTo(TotalAccessHistogramEntries * sizeof(uint8_t), 8);
 
   const u64 NumStackBytes = RoundUpTo(StackSizeBytes(StackIds), 8);
 

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

@@ -0,0 +1,33 @@
+// 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;
 
@@ -219,7 +218,11 @@ void Merge(const MemInfoBlock &newMIB) {
     ShorterHistogramSize = newMIB.AccessHistogramSize;
   }
   for (size_t i = 0; i < ShorterHistogramSize; ++i) {
-    ((uint64_t *)AccessHistogram)[i] += ((uint64_t *)ShorterHistogram)[i];
+    // Cast to uint8_t* and cap the sum at 255 to prevent overflow
+    uint8_t *CurrentHistPtr = (uint8_t *)AccessHistogram;
+    uint8_t *ShorterHistPtr = (uint8_t *)ShorterHistogram;
+    uint32_t sum = CurrentHistPtr[i] + ShorterHistPtr[i];
+    CurrentHistPtr[i] = (sum > 255) ? 255 : (uint8_t)sum;
   }
 }
 

llvm/lib/ProfileData/MemProfReader.cpp

@@ -166,6 +166,68 @@ readMemInfoBlocksV4(const char *Ptr) {
   return Items;
 }
 
+llvm::SmallVector<std::pair<uint64_t, MemInfoBlock>>
+readMemInfoBlocksV5(const char *Ptr) {
+  using namespace support;
+
+  const uint64_t NumItemsToRead =
+      endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
+
+  llvm::SmallVector<std::pair<uint64_t, MemInfoBlock>> Items;
+  for (uint64_t I = 0; I < NumItemsToRead; I++) {
+    const uint64_t Id =
+        endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
+
+    MemInfoBlock MIB;
+#define READ_MIB_FIELD(FIELD)                                                  \
+  MIB.FIELD = endian::readNext<decltype(MIB.FIELD), llvm::endianness::little,  \
+                               unaligned>(Ptr)
+
+    READ_MIB_FIELD(AllocCount);
+    READ_MIB_FIELD(TotalAccessCount);
+    READ_MIB_FIELD(MinAccessCount);
+    READ_MIB_FIELD(MaxAccessCount);
+    READ_MIB_FIELD(TotalSize);
+    READ_MIB_FIELD(MinSize);
+    READ_MIB_FIELD(MaxSize);
+    READ_MIB_FIELD(AllocTimestamp);
+    READ_MIB_FIELD(DeallocTimestamp);
+    READ_MIB_FIELD(TotalLifetime);
+    READ_MIB_FIELD(MinLifetime);
+    READ_MIB_FIELD(MaxLifetime);
+    READ_MIB_FIELD(AllocCpuId);
+    READ_MIB_FIELD(DeallocCpuId);
+    READ_MIB_FIELD(NumMigratedCpu);
+    READ_MIB_FIELD(NumLifetimeOverlaps);
+    READ_MIB_FIELD(NumSameAllocCpu);
+    READ_MIB_FIELD(NumSameDeallocCpu);
+    READ_MIB_FIELD(DataTypeId);
+    READ_MIB_FIELD(TotalAccessDensity);
+    READ_MIB_FIELD(MinAccessDensity);
+    READ_MIB_FIELD(MaxAccessDensity);
+    READ_MIB_FIELD(TotalLifetimeAccessDensity);
+    READ_MIB_FIELD(MinLifetimeAccessDensity);
+    READ_MIB_FIELD(MaxLifetimeAccessDensity);
+    READ_MIB_FIELD(AccessHistogramSize);
+    READ_MIB_FIELD(AccessHistogram);
+#undef READ_MIB_FIELD
+
+    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++) {
+        // The on-disk format for V5 uses uint8_t.
+        const uint8_t Val =
+            endian::readNext<uint8_t, llvm::endianness::little, unaligned>(Ptr);
+        ((uint64_t *)MIB.AccessHistogram)[J] = Val;
+      }
+    }
+    Items.push_back({Id, MIB});
+  }
+  return Items;
+}
+
 CallStackMap readStackInfo(const char *Ptr) {
   using namespace support;
 
@@ -658,6 +720,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");
 }

llvm/test/tools/llvm-profdata/memprof-basic-histogram.test

@@ -7,7 +7,7 @@ We expect 5 MIBs, each with different AccessHistogramValues.
 
 CHECK: MemprofProfile:
 CHECK-NEXT:   Summary:
-CHECK-NEXT:     Version: 4
+CHECK-NEXT:     Version: 5
 CHECK-NEXT:     NumSegments: {{[0-9]+}}
 CHECK-NEXT:     NumMibInfo: 5
 CHECK-NEXT:     NumAllocFunctions: 3
@@ -241,4 +241,4 @@ CHECK-NEXT:         MinLifetimeAccessDensity: 56000
 CHECK-NEXT:         MaxLifetimeAccessDensity: 56000
 CHECK-NEXT:         AccessHistogramSize: 8
 CHECK-NEXT:         AccessHistogram: {{[0-9]+}}
-CHECK-NEXT:         AccessHistogramValues: 168 147 126 105 84 63 42 21
+CHECK-NEXT:         AccessHistogramValues: 168 147 126 105 84 63 42 21

llvm/test/tools/llvm-profdata/memprof-basic.test

@@ -8,7 +8,7 @@ additional allocations which do not originate from the main binary are pruned.
 
 CHECK:  MemprofProfile:
 CHECK-NEXT:   Summary:
-CHECK-NEXT:     Version: 4
+CHECK-NEXT:     Version: 5
 CHECK-NEXT:     NumSegments: {{[0-9]+}}
 CHECK-NEXT:     NumMibInfo: 2
 CHECK-NEXT:     NumAllocFunctions: 1
@@ -96,4 +96,4 @@ CHECK-NEXT:         TotalLifetimeAccessDensity: 20000
 CHECK-NEXT:         MinLifetimeAccessDensity: 20000
 CHECK-NEXT:         MaxLifetimeAccessDensity: 20000
 CHECK-NEXT:         AccessHistogramSize: 0
-CHECK-NEXT:         AccessHistogram: 0
+CHECK-NEXT:         AccessHistogram: 0

llvm/test/tools/llvm-profdata/memprof-inline.test

@@ -5,7 +5,7 @@ RUN: llvm-profdata show --memory %p/Inputs/inline.memprofraw --profiled-binary %
 
 CHECK:  MemprofProfile:
 CHECK-NEXT:  Summary:
-CHECK-NEXT:    Version: 4
+CHECK-NEXT:    Version: 5
 CHECK-NEXT:    NumSegments: {{[0-9]+}}
 CHECK-NEXT:    NumMibInfo: 2
 CHECK-NEXT:    NumAllocFunctions: 2

llvm/test/tools/llvm-profdata/memprof-multi.test

@@ -7,7 +7,7 @@ We expect 2 MIB entries, 1 each for the malloc calls in the program.
 
 CHECK:  MemprofProfile:
 CHECK-NEXT:  Summary:
-CHECK-NEXT:    Version: 4
+CHECK-NEXT:    Version: 5
 CHECK-NEXT:    NumSegments: {{[0-9]+}}
 CHECK-NEXT:    NumMibInfo: 2
 CHECK-NEXT:    NumAllocFunctions: 1

llvm/test/tools/llvm-profdata/memprof-padding-histogram.test

@@ -7,7 +7,7 @@ We expect 2 different MIBs with histogram values. This test is to make sure we p
 
 CHECK: MemprofProfile:
 CHECK-NEXT:   Summary:
-CHECK-NEXT:     Version: 4
+CHECK-NEXT:     Version: 5
 CHECK-NEXT:     NumSegments: {{[0-9]+}}
 CHECK-NEXT:     NumMibInfo: 2
 CHECK-NEXT:     NumAllocFunctions: 1
@@ -96,4 +96,4 @@ CHEC-NEXT        MinLifetimeAccessDensity: 8000
 CHEC-NEXT        MaxLifetimeAccessDensity: 8000
 CHEC-NEXT        AccessHistogramSize: 6
 CHEC-NEXT        AccessHistogram: {{[0-9]+}}
-CHEC-NEXT        AccessHistogramValues: -2 -0 -0 -0 -1 -1
+CHEC-NEXT        AccessHistogramValues: -2 -0 -0 -0 -1 -1

llvm/test/tools/llvm-profdata/memprof-pic.test

@@ -11,7 +11,7 @@ RUN: llvm-profdata show --memory %p/Inputs/pic.memprofraw --profiled-binary %p/I
 
 CHECK:  MemprofProfile:
 CHECK-NEXT:   Summary:
-CHECK-NEXT:     Version: 4
+CHECK-NEXT:     Version: 5
 CHECK-NEXT:     NumSegments: {{[0-9]+}}
 CHECK-NEXT:     NumMibInfo: 2
 CHECK-NEXT:     NumAllocFunctions: 1
@@ -100,4 +100,4 @@ CHECK-NEXT:         TotalLifetimeAccessDensity: 20000
 CHECK-NEXT:         MinLifetimeAccessDensity: 20000
 CHECK-NEXT:         MaxLifetimeAccessDensity: 20000
 CHECK-NEXT:         AccessHistogramSize: 0
-CHECK-NEXT:         AccessHistogram: 0
+CHECK-NEXT:         AccessHistogram: 0