[compiler-rt] [llvm] [Memprof] Adds instrumentation support for memprof with histograms. (PR #100834)

[Matthew Weingarten via llvm-commits]

https://github.com/mattweingarten created https://github.com/llvm/llvm-project/pull/100834

This patch allows running `-fmemory-profile` without the flag `-memprof-use-callbacks`, meaning the `RecordAccessesHistogram`  is injected into IR as a sequence of instructions.  This significantly increases performance of the instrumented binary. 

Changes the `HISTOGRAM_GRANULARITY` from `8U` to `8ULL` to correct the shadow address translation.

>From 68725fd958ec6d480eab20faca4612148aeb1fdf Mon Sep 17 00:00:00 2001
From: Matthew Weingarten <matt at weingarten.org>
Date: Mon, 22 Jul 2024 22:56:58 +0000
Subject: [PATCH] [Memprof] Adds instrumentation support for running memprof
 with histograms.

This patch allows running `-fmemory-profile` without `-memprof-use-callbacks`. This significantly increases
performance of the instrumented binary.

Additionally, changes the `HISTOGRAM_GRANULARITY` from 8U to 8ULL to correct shadow address translation.
---
 compiler-rt/lib/memprof/memprof_mapping.h     |  2 +-
 .../Instrumentation/MemProfiler.cpp           | 52 +++++++++++++------
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/compiler-rt/lib/memprof/memprof_mapping.h b/compiler-rt/lib/memprof/memprof_mapping.h
index fef8acfcfc921..6da385ab3d6e2 100644
--- a/compiler-rt/lib/memprof/memprof_mapping.h
+++ b/compiler-rt/lib/memprof/memprof_mapping.h
@@ -55,7 +55,7 @@ extern uptr kHighMemEnd; // Initialized in __memprof_init.
 // computed by summing up all individual 1 byte counters. This can incur an
 // accuracy penalty.
 
-#define HISTOGRAM_GRANULARITY 8U
+#define HISTOGRAM_GRANULARITY 8ULL
 
 #define HISTOGRAM_MAX_COUNTER 255U
 
diff --git a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
index 2c5d749d4a67a..0e4d21671d2a7 100644
--- a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
@@ -61,6 +61,9 @@ constexpr int LLVM_MEM_PROFILER_VERSION = 1;
 // Size of memory mapped to a single shadow location.
 constexpr uint64_t DefaultMemGranularity = 64;
 
+// Size of memory mapped to a single histogram bucket.
+constexpr uint64_t HistogramGranularity = 8;
+
 // Scale from granularity down to shadow size.
 constexpr uint64_t DefaultShadowScale = 3;
 
@@ -192,7 +195,8 @@ namespace {
 struct ShadowMapping {
   ShadowMapping() {
     Scale = ClMappingScale;
-    Granularity = ClMappingGranularity;
+    // Histogram Granularity is always 8U for now.
+    Granularity = ClHistogram ? HistogramGranularity : ClMappingGranularity;
     Mask = ~(Granularity - 1);
   }
 
@@ -288,10 +292,6 @@ ModuleMemProfilerPass::ModuleMemProfilerPass() = default;
 PreservedAnalyses ModuleMemProfilerPass::run(Module &M,
                                              AnalysisManager<Module> &AM) {
 
-  assert((!ClHistogram || (ClHistogram && ClUseCalls)) &&
-         "Cannot use -memprof-histogram without Callbacks. Set "
-         "memprof-use-callbacks");
-
   ModuleMemProfiler Profiler(M);
   if (Profiler.instrumentModule(M))
     return PreservedAnalyses::none();
@@ -489,16 +489,38 @@ void MemProfiler::instrumentAddress(Instruction *OrigIns,
     return;
   }
 
-  // Create an inline sequence to compute shadow location, and increment the
-  // value by one.
-  Type *ShadowTy = Type::getInt64Ty(*C);
-  Type *ShadowPtrTy = PointerType::get(ShadowTy, 0);
-  Value *ShadowPtr = memToShadow(AddrLong, IRB);
-  Value *ShadowAddr = IRB.CreateIntToPtr(ShadowPtr, ShadowPtrTy);
-  Value *ShadowValue = IRB.CreateLoad(ShadowTy, ShadowAddr);
-  Value *Inc = ConstantInt::get(Type::getInt64Ty(*C), 1);
-  ShadowValue = IRB.CreateAdd(ShadowValue, Inc);
-  IRB.CreateStore(ShadowValue, ShadowAddr);
+  if (ClHistogram) {
+    Type *ShadowTy = Type::getInt8Ty(*C);
+    Type *ShadowPtrTy = PointerType::get(ShadowTy, 0);
+    Value *Shadow = IRB.CreateAnd(AddrLong, Mapping.Mask);
+    Shadow = IRB.CreateLShr(Shadow, Mapping.Scale);
+    assert(DynamicShadowOffset);
+    Value *ShadowInteger = IRB.CreateAdd(Shadow, DynamicShadowOffset);
+
+    Value *ShadowAddr = IRB.CreateIntToPtr(ShadowInteger, ShadowPtrTy);
+    Value *ShadowValue = IRB.CreateLoad(ShadowTy, ShadowAddr);
+    Value *MaxCount = ConstantInt::get(Type::getInt8Ty(*C), 255);
+    Value *Cmp = IRB.CreateICmpULT(ShadowValue, MaxCount);
+    // Insert Basic block that increments histogram bucket if value is less
+    // than 255.
+    Instruction *IncBlock =
+        SplitBlockAndInsertIfThen(Cmp, InsertBefore, /*Unreachable=*/false);
+    IRB.SetInsertPoint(IncBlock);
+    Value *Inc = ConstantInt::get(Type::getInt8Ty(*C), 1);
+    ShadowValue = IRB.CreateAdd(ShadowValue, Inc);
+    IRB.CreateStore(ShadowValue, ShadowAddr);
+  } else {
+    // Create an inline sequence to compute shadow location, and increment the
+    // value by one.
+    Type *ShadowTy = Type::getInt64Ty(*C);
+    Type *ShadowPtrTy = PointerType::get(ShadowTy, 0);
+    Value *ShadowPtr = memToShadow(AddrLong, IRB);
+    Value *ShadowAddr = IRB.CreateIntToPtr(ShadowPtr, ShadowPtrTy);
+    Value *ShadowValue = IRB.CreateLoad(ShadowTy, ShadowAddr);
+    Value *Inc = ConstantInt::get(Type::getInt64Ty(*C), 1);
+    ShadowValue = IRB.CreateAdd(ShadowValue, Inc);
+    IRB.CreateStore(ShadowValue, ShadowAddr);
+  }
 }
 
 // Create the variable for the profile file name.