Draft changes to trim the histogram raw profile format.

Author: Snehasish Kumar

compiler-rt/include/profile/MemProfData.inc

@@ -33,11 +33,11 @@
    (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 }
+  { 3ULL, 4ULL, 5ULL }
 
 #define MEMPROF_V3_MIB_SIZE 132ULL;
 
@@ -219,7 +219,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,13 @@ 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,32 @@
+// 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,11 @@
    (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 }
+  { 3ULL, 4ULL, 5ULL }
 
 #define MEMPROF_V3_MIB_SIZE 132ULL;
 
@@ -219,7 +219,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;
   }
 }