[memprof] Track per-allocation MIBs by address

compiler-rt/lib/memprof/memprof_allocator.cpp

@@ -39,16 +39,21 @@
 extern bool __memprof_histogram;
 
 namespace __memprof {
+
+// Monotonically increasing counter; each allocation gets a unique ID.
+static atomic_uint64_t global_alloc_id;
+
 namespace {
 using ::llvm::memprof::MemInfoBlock;
 
-void Print(const MemInfoBlock &M, const u64 id, bool print_terse) {
+void Print(const MemInfoBlock &M, const u64 id, bool print_terse,
+           uptr addr = 0) {
   u64 p;
 
   if (print_terse) {
     p = M.TotalSize * 100 / M.AllocCount;
-    Printf("MIB:%llu/%u/%llu.%02llu/%u/%u/", id, M.AllocCount, p / 100, p % 100,
-           M.MinSize, M.MaxSize);
+    Printf("MIB:%llu/%llx/%u/%llu.%02llu/%u/%u/", id, (u64)addr, M.AllocCount,
+           p / 100, p % 100, M.MinSize, M.MaxSize);
     p = M.TotalAccessCount * 100 / M.AllocCount;
     Printf("%llu.%02llu/%llu/%llu/", p / 100, p % 100, M.MinAccessCount,
            M.MaxAccessCount);
@@ -60,6 +65,7 @@ void Print(const MemInfoBlock &M, const u64 id, bool print_terse) {
   } else {
     p = M.TotalSize * 100 / M.AllocCount;
     Printf("Memory allocation stack id = %llu\n", id);
+    Printf("\taddr = 0x%llx\n", (u64)addr);
     Printf("\talloc_count %u, size (ave/min/max) %llu.%02llu / %u / %u\n",
            M.AllocCount, p / 100, p % 100, M.MinSize, M.MaxSize);
     p = M.TotalAccessCount * 100 / M.AllocCount;
@@ -307,7 +313,8 @@ struct Allocator {
   static void PrintCallback(const uptr Key, LockedMemInfoBlock *const &Value,
                             void *Arg) {
     SpinMutexLock l(&Value->mutex);
-    Print(Value->mib, Key, bool(Arg));
+    // Key is the unique alloc ID; addr and stack_id are stored in Value.
+    Print(Value->mib, Value->stack_id, bool(Arg), Value->addr);
   }
 
   // See memprof_mapping.h for an overview on histogram counters.
@@ -389,7 +396,7 @@ struct Allocator {
             return;
           uptr user_beg = ((uptr)m) + kChunkHeaderSize;
           MemInfoBlock newMIB = CreateNewMIB(user_beg, m, user_requested_size);
-          InsertOrMerge(m->alloc_context_id, newMIB, A->MIBMap);
+          InsertOrMerge(m->data_type_id, newMIB, A->MIBMap, m->alloc_context_id, user_beg);
         },
         this);
   }
@@ -477,6 +484,8 @@ struct Allocator {
     m->cpu_id = GetCpuId();
     m->timestamp_ms = GetTimestamp();
     m->alloc_context_id = StackDepotPut(*stack);
+    // +1 so alloc IDs start at 1; AddrHashMap uses 0 as the empty-slot sentinel.
+    m->data_type_id = atomic_fetch_add(&global_alloc_id, 1, memory_order_relaxed) + 1;
 
     uptr size_rounded_down_to_granularity =
         RoundDownTo(size, SHADOW_GRANULARITY);
@@ -518,7 +527,7 @@ struct Allocator {
     if (memprof_inited && atomic_load_relaxed(&constructed) &&
         !atomic_load_relaxed(&destructing)) {
       MemInfoBlock newMIB = this->CreateNewMIB(p, m, user_requested_size);
-      InsertOrMerge(m->alloc_context_id, newMIB, MIBMap);
+      InsertOrMerge(m->data_type_id, newMIB, MIBMap, m->alloc_context_id, p);
     }
 
     MemprofStats &thread_stats = GetCurrentThreadStats();

compiler-rt/lib/memprof/memprof_mibmap.cpp

@@ -18,16 +18,19 @@
 namespace __memprof {
 using ::llvm::memprof::MemInfoBlock;
 
-void InsertOrMerge(const uptr Id, const MemInfoBlock &Block, MIBMapTy &Map) {
-  MIBMapTy::Handle h(&Map, static_cast<uptr>(Id), /*remove=*/false,
-                     /*create=*/true);
+void InsertOrMerge(const uptr AllocId, const MemInfoBlock &Block,
+                   MIBMapTy &Map, u32 StackId, uptr Addr) {
+  MIBMapTy::Handle h(&Map, AllocId, /*remove=*/false, /*create=*/true);
   if (h.created()) {
     LockedMemInfoBlock *lmib =
         (LockedMemInfoBlock *)InternalAlloc(sizeof(LockedMemInfoBlock));
     lmib->mutex.Init();
+    lmib->stack_id = StackId;
+    lmib->addr = Addr;
     lmib->mib = Block;
     *h = lmib;
   } else {
+    // AllocId is unique per allocation, so this branch should not be reached.
     LockedMemInfoBlock *lmib = *h;
     SpinMutexLock lock(&lmib->mutex);
     uintptr_t ShorterHistogram;
@@ -37,9 +40,6 @@ void InsertOrMerge(const uptr Id, const MemInfoBlock &Block, MIBMapTy &Map) {
       ShorterHistogram = Block.AccessHistogram;
 
     lmib->mib.Merge(Block);
-    // The larger histogram is kept and the shorter histogram is discarded after
-    // adding the counters to the larger historam. Free only the shorter
-    // Histogram
     if (Block.AccessHistogramSize > 0 || lmib->mib.AccessHistogramSize > 0)
       InternalFree((void *)ShorterHistogram);
   }

compiler-rt/lib/memprof/memprof_mibmap.h

@@ -11,16 +11,22 @@ namespace __memprof {
 
 struct LockedMemInfoBlock {
   __sanitizer::StaticSpinMutex mutex;
+  // Call stack id for this allocation, used to print the stack trace.
+  u32 stack_id;
+  // Starting address returned by the allocator. Primary key for joining with
+  // simulator access traces.
+  uptr addr;
   ::llvm::memprof::MemInfoBlock mib;
 };
 
-// The MIB map stores a mapping from stack ids to MemInfoBlocks.
+// The MIB map stores a mapping from monotonic alloc IDs to MemInfoBlocks.
+// Each entry represents one individual allocation (no merging).
 typedef __sanitizer::AddrHashMap<LockedMemInfoBlock *, 200003> MIBMapTy;
 
-// Insert a new MemInfoBlock or merge with an existing block identified by the
-// stack id.
-void InsertOrMerge(const uptr Id, const ::llvm::memprof::MemInfoBlock &Block,
-                   MIBMapTy &Map);
+// Insert a new MemInfoBlock keyed by a unique AllocId.
+// StackId is the call stack id; Addr is the starting address returned to the user.
+void InsertOrMerge(const uptr AllocId, const ::llvm::memprof::MemInfoBlock &Block,
+                   MIBMapTy &Map, u32 StackId, uptr Addr);
 
 } // namespace __memprof