llvm/lib/Transforms/Instrumentation/MemProfiler.cpp

   1 //===- MemProfiler.cpp - memory allocation and access profiler ------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file is a part of MemProfiler. Memory accesses are instrumented
  10 // to increment the access count held in a shadow memory location, or
  11 // alternatively to call into the runtime. Memory intrinsic calls (memmove,
  12 // memcpy, memset) are changed to call the memory profiling runtime version
  13 // instead.
  14 //
  15 //===----------------------------------------------------------------------===//
  16
  17 #include "llvm/Transforms/Instrumentation/MemProfiler.h"
  18 #include "llvm/ADT/SmallVector.h"
  19 #include "llvm/ADT/Statistic.h"
  20 #include "llvm/ADT/StringRef.h"
  21 #include "llvm/Analysis/MemoryBuiltins.h"
  22 #include "llvm/Analysis/MemoryProfileInfo.h"
  23 #include "llvm/Analysis/ValueTracking.h"
  24 #include "llvm/IR/Constant.h"
  25 #include "llvm/IR/DataLayout.h"
  26 #include "llvm/IR/DiagnosticInfo.h"
  27 #include "llvm/IR/Function.h"
  28 #include "llvm/IR/GlobalValue.h"
  29 #include "llvm/IR/IRBuilder.h"
  30 #include "llvm/IR/Instruction.h"
  31 #include "llvm/IR/IntrinsicInst.h"
  32 #include "llvm/IR/Module.h"
  33 #include "llvm/IR/Type.h"
  34 #include "llvm/IR/Value.h"
  35 #include "llvm/ProfileData/InstrProf.h"
  36 #include "llvm/ProfileData/InstrProfReader.h"
  37 #include "llvm/Support/BLAKE3.h"
  38 #include "llvm/Support/CommandLine.h"
  39 #include "llvm/Support/Debug.h"
  40 #include "llvm/Support/HashBuilder.h"
  41 #include "llvm/Support/VirtualFileSystem.h"
  42 #include "llvm/TargetParser/Triple.h"
  43 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  44 #include "llvm/Transforms/Utils/ModuleUtils.h"
  45 #include <map>
  46 #include <set>
  47
  48 using namespace llvm;
  49 using namespace llvm::memprof;
  50
  51 #define DEBUG_TYPE "memprof"
  52
  53 namespace llvm {
  54 extern cl::opt<bool> PGOWarnMissing;
  55 extern cl::opt<bool> NoPGOWarnMismatch;
  56 extern cl::opt<bool> NoPGOWarnMismatchComdatWeak;
  57 } // namespace llvm
  58
  59 constexpr int LLVM_MEM_PROFILER_VERSION = 1;
  60
  61 // Size of memory mapped to a single shadow location.
  62 constexpr uint64_t DefaultShadowGranularity = 64;
  63
  64 // Scale from granularity down to shadow size.
  65 constexpr uint64_t DefaultShadowScale = 3;
  66
  67 constexpr char MemProfModuleCtorName[] = "memprof.module_ctor";
  68 constexpr uint64_t MemProfCtorAndDtorPriority = 1;
  69 // On Emscripten, the system needs more than one priorities for constructors.
  70 constexpr uint64_t MemProfEmscriptenCtorAndDtorPriority = 50;
  71 constexpr char MemProfInitName[] = "__memprof_init";
  72 constexpr char MemProfVersionCheckNamePrefix[] =
  73     "__memprof_version_mismatch_check_v";
  74
  75 constexpr char MemProfShadowMemoryDynamicAddress[] =
  76     "__memprof_shadow_memory_dynamic_address";
  77
  78 constexpr char MemProfFilenameVar[] = "__memprof_profile_filename";
  79
  80 // Command-line flags.
  81
  82 static cl::opt<bool> ClInsertVersionCheck(
  83     "memprof-guard-against-version-mismatch",
  84     cl::desc("Guard against compiler/runtime version mismatch."), cl::Hidden,
  85     cl::init(true));
  86
  87 // This flag may need to be replaced with -f[no-]memprof-reads.
  88 static cl::opt<bool> ClInstrumentReads("memprof-instrument-reads",
  89                                        cl::desc("instrument read instructions"),
  90                                        cl::Hidden, cl::init(true));
  91
  92 static cl::opt<bool>
  93     ClInstrumentWrites("memprof-instrument-writes",
  94                        cl::desc("instrument write instructions"), cl::Hidden,
  95                        cl::init(true));
  96
  97 static cl::opt<bool> ClInstrumentAtomics(
  98     "memprof-instrument-atomics",
  99     cl::desc("instrument atomic instructions (rmw, cmpxchg)"), cl::Hidden,
 100     cl::init(true));
 101
 102 static cl::opt<bool> ClUseCalls(
 103     "memprof-use-callbacks",
 104     cl::desc("Use callbacks instead of inline instrumentation sequences."),
 105     cl::Hidden, cl::init(false));
 106
 107 static cl::opt<std::string>
 108     ClMemoryAccessCallbackPrefix("memprof-memory-access-callback-prefix",
 109                                  cl::desc("Prefix for memory access callbacks"),
 110                                  cl::Hidden, cl::init("__memprof_"));
 111
 112 // These flags allow to change the shadow mapping.
 113 // The shadow mapping looks like
 114 //    Shadow = ((Mem & mask) >> scale) + offset
 115
 116 static cl::opt<int> ClMappingScale("memprof-mapping-scale",
 117                                    cl::desc("scale of memprof shadow mapping"),
 118                                    cl::Hidden, cl::init(DefaultShadowScale));
 119
 120 static cl::opt<int>
 121     ClMappingGranularity("memprof-mapping-granularity",
 122                          cl::desc("granularity of memprof shadow mapping"),
 123                          cl::Hidden, cl::init(DefaultShadowGranularity));
 124
 125 static cl::opt<bool> ClStack("memprof-instrument-stack",
 126                              cl::desc("Instrument scalar stack variables"),
 127                              cl::Hidden, cl::init(false));
 128
 129 // Debug flags.
 130
 131 static cl::opt<int> ClDebug("memprof-debug", cl::desc("debug"), cl::Hidden,
 132                             cl::init(0));
 133
 134 static cl::opt<std::string> ClDebugFunc("memprof-debug-func", cl::Hidden,
 135                                         cl::desc("Debug func"));
 136
 137 static cl::opt<int> ClDebugMin("memprof-debug-min", cl::desc("Debug min inst"),
 138                                cl::Hidden, cl::init(-1));
 139
 140 static cl::opt<int> ClDebugMax("memprof-debug-max", cl::desc("Debug max inst"),
 141                                cl::Hidden, cl::init(-1));
 142
 143 STATISTIC(NumInstrumentedReads, "Number of instrumented reads");
 144 STATISTIC(NumInstrumentedWrites, "Number of instrumented writes");
 145 STATISTIC(NumSkippedStackReads, "Number of non-instrumented stack reads");
 146 STATISTIC(NumSkippedStackWrites, "Number of non-instrumented stack writes");
 147 STATISTIC(NumOfMemProfMissing, "Number of functions without memory profile.");
 148
 149 namespace {
 150
 151 /// This struct defines the shadow mapping using the rule:
 152 ///   shadow = ((mem & mask) >> Scale) ADD DynamicShadowOffset.
 153 struct ShadowMapping {
 154   ShadowMapping() {
 155     Scale = ClMappingScale;
 156     Granularity = ClMappingGranularity;
 157     Mask = ~(Granularity - 1);
 158   }
 159
 160   int Scale;
 161   int Granularity;
 162   uint64_t Mask; // Computed as ~(Granularity-1)
 163 };
 164
 165 static uint64_t getCtorAndDtorPriority(Triple &TargetTriple) {
 166   return TargetTriple.isOSEmscripten() ? MemProfEmscriptenCtorAndDtorPriority
 167                                        : MemProfCtorAndDtorPriority;
 168 }
 169
 170 struct InterestingMemoryAccess {
 171   Value *Addr = nullptr;
 172   bool IsWrite;
 173   Type *AccessTy;
 174   uint64_t TypeSize;
 175   Value *MaybeMask = nullptr;
 176 };
 177
 178 /// Instrument the code in module to profile memory accesses.
 179 class MemProfiler {
 180 public:
 181   MemProfiler(Module &M) {
 182     C = &(M.getContext());
 183     LongSize = M.getDataLayout().getPointerSizeInBits();
 184     IntptrTy = Type::getIntNTy(*C, LongSize);
 185   }
 186
 187   /// If it is an interesting memory access, populate information
 188   /// about the access and return a InterestingMemoryAccess struct.
 189   /// Otherwise return std::nullopt.
 190   std::optional<InterestingMemoryAccess>
 191   isInterestingMemoryAccess(Instruction *I) const;
 192
 193   void instrumentMop(Instruction *I, const DataLayout &DL,
 194                      InterestingMemoryAccess &Access);
 195   void instrumentAddress(Instruction *OrigIns, Instruction *InsertBefore,
 196                          Value *Addr, uint32_t TypeSize, bool IsWrite);
 197   void instrumentMaskedLoadOrStore(const DataLayout &DL, Value *Mask,
 198                                    Instruction *I, Value *Addr, Type *AccessTy,
 199                                    bool IsWrite);
 200   void instrumentMemIntrinsic(MemIntrinsic *MI);
 201   Value *memToShadow(Value *Shadow, IRBuilder<> &IRB);
 202   bool instrumentFunction(Function &F);
 203   bool maybeInsertMemProfInitAtFunctionEntry(Function &F);
 204   bool insertDynamicShadowAtFunctionEntry(Function &F);
 205
 206 private:
 207   void initializeCallbacks(Module &M);
 208
 209   LLVMContext *C;
 210   int LongSize;
 211   Type *IntptrTy;
 212   ShadowMapping Mapping;
 213
 214   // These arrays is indexed by AccessIsWrite
 215   FunctionCallee MemProfMemoryAccessCallback[2];
 216   FunctionCallee MemProfMemoryAccessCallbackSized[2];
 217
 218   FunctionCallee MemProfMemmove, MemProfMemcpy, MemProfMemset;
 219   Value *DynamicShadowOffset = nullptr;
 220 };
 221
 222 class ModuleMemProfiler {
 223 public:
 224   ModuleMemProfiler(Module &M) { TargetTriple = Triple(M.getTargetTriple()); }
 225
 226   bool instrumentModule(Module &);
 227
 228 private:
 229   Triple TargetTriple;
 230   ShadowMapping Mapping;
 231   Function *MemProfCtorFunction = nullptr;
 232 };
 233
 234 } // end anonymous namespace
 235
 236 MemProfilerPass::MemProfilerPass() = default;
 237
 238 PreservedAnalyses MemProfilerPass::run(Function &F,
 239                                        AnalysisManager<Function> &AM) {
 240   Module &M = *F.getParent();
 241   MemProfiler Profiler(M);
 242   if (Profiler.instrumentFunction(F))
 243     return PreservedAnalyses::none();
 244   return PreservedAnalyses::all();
 245 }
 246
 247 ModuleMemProfilerPass::ModuleMemProfilerPass() = default;
 248
 249 PreservedAnalyses ModuleMemProfilerPass::run(Module &M,
 250                                              AnalysisManager<Module> &AM) {
 251   ModuleMemProfiler Profiler(M);
 252   if (Profiler.instrumentModule(M))
 253     return PreservedAnalyses::none();
 254   return PreservedAnalyses::all();
 255 }
 256
 257 Value *MemProfiler::memToShadow(Value *Shadow, IRBuilder<> &IRB) {
 258   // (Shadow & mask) >> scale
 259   Shadow = IRB.CreateAnd(Shadow, Mapping.Mask);
 260   Shadow = IRB.CreateLShr(Shadow, Mapping.Scale);
 261   // (Shadow >> scale) | offset
 262   assert(DynamicShadowOffset);
 263   return IRB.CreateAdd(Shadow, DynamicShadowOffset);
 264 }
 265
 266 // Instrument memset/memmove/memcpy
 267 void MemProfiler::instrumentMemIntrinsic(MemIntrinsic *MI) {
 268   IRBuilder<> IRB(MI);
 269   if (isa<MemTransferInst>(MI)) {
 270     IRB.CreateCall(
 271         isa<MemMoveInst>(MI) ? MemProfMemmove : MemProfMemcpy,
 272         {IRB.CreatePointerCast(MI->getOperand(0), IRB.getInt8PtrTy()),
 273          IRB.CreatePointerCast(MI->getOperand(1), IRB.getInt8PtrTy()),
 274          IRB.CreateIntCast(MI->getOperand(2), IntptrTy, false)});
 275   } else if (isa<MemSetInst>(MI)) {
 276     IRB.CreateCall(
 277         MemProfMemset,
 278         {IRB.CreatePointerCast(MI->getOperand(0), IRB.getInt8PtrTy()),
 279          IRB.CreateIntCast(MI->getOperand(1), IRB.getInt32Ty(), false),
 280          IRB.CreateIntCast(MI->getOperand(2), IntptrTy, false)});
 281   }
 282   MI->eraseFromParent();
 283 }
 284
 285 std::optional<InterestingMemoryAccess>
 286 MemProfiler::isInterestingMemoryAccess(Instruction *I) const {
 287   // Do not instrument the load fetching the dynamic shadow address.
 288   if (DynamicShadowOffset == I)
 289     return std::nullopt;
 290
 291   InterestingMemoryAccess Access;
 292
 293   if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
 294     if (!ClInstrumentReads)
 295       return std::nullopt;
 296     Access.IsWrite = false;
 297     Access.AccessTy = LI->getType();
 298     Access.Addr = LI->getPointerOperand();
 299   } else if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
 300     if (!ClInstrumentWrites)
 301       return std::nullopt;
 302     Access.IsWrite = true;
 303     Access.AccessTy = SI->getValueOperand()->getType();
 304     Access.Addr = SI->getPointerOperand();
 305   } else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(I)) {
 306     if (!ClInstrumentAtomics)
 307       return std::nullopt;
 308     Access.IsWrite = true;
 309     Access.AccessTy = RMW->getValOperand()->getType();
 310     Access.Addr = RMW->getPointerOperand();
 311   } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(I)) {
 312     if (!ClInstrumentAtomics)
 313       return std::nullopt;
 314     Access.IsWrite = true;
 315     Access.AccessTy = XCHG->getCompareOperand()->getType();
 316     Access.Addr = XCHG->getPointerOperand();
 317   } else if (auto *CI = dyn_cast<CallInst>(I)) {
 318     auto *F = CI->getCalledFunction();
 319     if (F && (F->getIntrinsicID() == Intrinsic::masked_load ||
 320               F->getIntrinsicID() == Intrinsic::masked_store)) {
 321       unsigned OpOffset = 0;
 322       if (F->getIntrinsicID() == Intrinsic::masked_store) {
 323         if (!ClInstrumentWrites)
 324           return std::nullopt;
 325         // Masked store has an initial operand for the value.
 326         OpOffset = 1;
 327         Access.AccessTy = CI->getArgOperand(0)->getType();
 328         Access.IsWrite = true;
 329       } else {
 330         if (!ClInstrumentReads)
 331           return std::nullopt;
 332         Access.AccessTy = CI->getType();
 333         Access.IsWrite = false;
 334       }
 335
 336       auto *BasePtr = CI->getOperand(0 + OpOffset);
 337       Access.MaybeMask = CI->getOperand(2 + OpOffset);
 338       Access.Addr = BasePtr;
 339     }
 340   }
 341
 342   if (!Access.Addr)
 343     return std::nullopt;
 344
 345   // Do not instrument accesses from different address spaces; we cannot deal
 346   // with them.
 347   Type *PtrTy = cast<PointerType>(Access.Addr->getType()->getScalarType());
 348   if (PtrTy->getPointerAddressSpace() != 0)
 349     return std::nullopt;
 350
 351   // Ignore swifterror addresses.
 352   // swifterror memory addresses are mem2reg promoted by instruction
 353   // selection. As such they cannot have regular uses like an instrumentation
 354   // function and it makes no sense to track them as memory.
 355   if (Access.Addr->isSwiftError())
 356     return std::nullopt;
 357
 358   // Peel off GEPs and BitCasts.
 359   auto *Addr = Access.Addr->stripInBoundsOffsets();
 360
 361   if (GlobalVariable *GV = dyn_cast<GlobalVariable>(Addr)) {
 362     // Do not instrument PGO counter updates.
 363     if (GV->hasSection()) {
 364       StringRef SectionName = GV->getSection();
 365       // Check if the global is in the PGO counters section.
 366       auto OF = Triple(I->getModule()->getTargetTriple()).getObjectFormat();
 367       if (SectionName.endswith(
 368               getInstrProfSectionName(IPSK_cnts, OF, /*AddSegmentInfo=*/false)))
 369         return std::nullopt;
 370     }
 371
 372     // Do not instrument accesses to LLVM internal variables.
 373     if (GV->getName().startswith("__llvm"))
 374       return std::nullopt;
 375   }
 376
 377   const DataLayout &DL = I->getModule()->getDataLayout();
 378   Access.TypeSize = DL.getTypeStoreSizeInBits(Access.AccessTy);
 379   return Access;
 380 }
 381
 382 void MemProfiler::instrumentMaskedLoadOrStore(const DataLayout &DL, Value *Mask,
 383                                               Instruction *I, Value *Addr,
 384                                               Type *AccessTy, bool IsWrite) {
 385   auto *VTy = cast<FixedVectorType>(AccessTy);
 386   uint64_t ElemTypeSize = DL.getTypeStoreSizeInBits(VTy->getScalarType());
 387   unsigned Num = VTy->getNumElements();
 388   auto *Zero = ConstantInt::get(IntptrTy, 0);
 389   for (unsigned Idx = 0; Idx < Num; ++Idx) {
 390     Value *InstrumentedAddress = nullptr;
 391     Instruction *InsertBefore = I;
 392     if (auto *Vector = dyn_cast<ConstantVector>(Mask)) {
 393       // dyn_cast as we might get UndefValue
 394       if (auto *Masked = dyn_cast<ConstantInt>(Vector->getOperand(Idx))) {
 395         if (Masked->isZero())
 396           // Mask is constant false, so no instrumentation needed.
 397           continue;
 398         // If we have a true or undef value, fall through to instrumentAddress.
 399         // with InsertBefore == I
 400       }
 401     } else {
 402       IRBuilder<> IRB(I);
 403       Value *MaskElem = IRB.CreateExtractElement(Mask, Idx);
 404       Instruction *ThenTerm = SplitBlockAndInsertIfThen(MaskElem, I, false);
 405       InsertBefore = ThenTerm;
 406     }
 407
 408     IRBuilder<> IRB(InsertBefore);
 409     InstrumentedAddress =
 410         IRB.CreateGEP(VTy, Addr, {Zero, ConstantInt::get(IntptrTy, Idx)});
 411     instrumentAddress(I, InsertBefore, InstrumentedAddress, ElemTypeSize,
 412                       IsWrite);
 413   }
 414 }
 415
 416 void MemProfiler::instrumentMop(Instruction *I, const DataLayout &DL,
 417                                 InterestingMemoryAccess &Access) {
 418   // Skip instrumentation of stack accesses unless requested.
 419   if (!ClStack && isa<AllocaInst>(getUnderlyingObject(Access.Addr))) {
 420     if (Access.IsWrite)
 421       ++NumSkippedStackWrites;
 422     else
 423       ++NumSkippedStackReads;
 424     return;
 425   }
 426
 427   if (Access.IsWrite)
 428     NumInstrumentedWrites++;
 429   else
 430     NumInstrumentedReads++;
 431
 432   if (Access.MaybeMask) {
 433     instrumentMaskedLoadOrStore(DL, Access.MaybeMask, I, Access.Addr,
 434                                 Access.AccessTy, Access.IsWrite);
 435   } else {
 436     // Since the access counts will be accumulated across the entire allocation,
 437     // we only update the shadow access count for the first location and thus
 438     // don't need to worry about alignment and type size.
 439     instrumentAddress(I, I, Access.Addr, Access.TypeSize, Access.IsWrite);
 440   }
 441 }
 442
 443 void MemProfiler::instrumentAddress(Instruction *OrigIns,
 444                                     Instruction *InsertBefore, Value *Addr,
 445                                     uint32_t TypeSize, bool IsWrite) {
 446   IRBuilder<> IRB(InsertBefore);
 447   Value *AddrLong = IRB.CreatePointerCast(Addr, IntptrTy);
 448
 449   if (ClUseCalls) {
 450     IRB.CreateCall(MemProfMemoryAccessCallback[IsWrite], AddrLong);
 451     return;
 452   }
 453
 454   // Create an inline sequence to compute shadow location, and increment the
 455   // value by one.
 456   Type *ShadowTy = Type::getInt64Ty(*C);
 457   Type *ShadowPtrTy = PointerType::get(ShadowTy, 0);
 458   Value *ShadowPtr = memToShadow(AddrLong, IRB);
 459   Value *ShadowAddr = IRB.CreateIntToPtr(ShadowPtr, ShadowPtrTy);
 460   Value *ShadowValue = IRB.CreateLoad(ShadowTy, ShadowAddr);
 461   Value *Inc = ConstantInt::get(Type::getInt64Ty(*C), 1);
 462   ShadowValue = IRB.CreateAdd(ShadowValue, Inc);
 463   IRB.CreateStore(ShadowValue, ShadowAddr);
 464 }
 465
 466 // Create the variable for the profile file name.
 467 void createProfileFileNameVar(Module &M) {
 468   const MDString *MemProfFilename =
 469       dyn_cast_or_null<MDString>(M.getModuleFlag("MemProfProfileFilename"));
 470   if (!MemProfFilename)
 471     return;
 472   assert(!MemProfFilename->getString().empty() &&
 473          "Unexpected MemProfProfileFilename metadata with empty string");
 474   Constant *ProfileNameConst = ConstantDataArray::getString(
 475       M.getContext(), MemProfFilename->getString(), true);
 476   GlobalVariable *ProfileNameVar = new GlobalVariable(
 477       M, ProfileNameConst->getType(), /*isConstant=*/true,
 478       GlobalValue::WeakAnyLinkage, ProfileNameConst, MemProfFilenameVar);
 479   Triple TT(M.getTargetTriple());
 480   if (TT.supportsCOMDAT()) {
 481     ProfileNameVar->setLinkage(GlobalValue::ExternalLinkage);
 482     ProfileNameVar->setComdat(M.getOrInsertComdat(MemProfFilenameVar));
 483   }
 484 }
 485
 486 bool ModuleMemProfiler::instrumentModule(Module &M) {
 487   // Create a module constructor.
 488   std::string MemProfVersion = std::to_string(LLVM_MEM_PROFILER_VERSION);
 489   std::string VersionCheckName =
 490       ClInsertVersionCheck ? (MemProfVersionCheckNamePrefix + MemProfVersion)
 491                            : "";
 492   std::tie(MemProfCtorFunction, std::ignore) =
 493       createSanitizerCtorAndInitFunctions(M, MemProfModuleCtorName,
 494                                           MemProfInitName, /*InitArgTypes=*/{},
 495                                           /*InitArgs=*/{}, VersionCheckName);
 496
 497   const uint64_t Priority = getCtorAndDtorPriority(TargetTriple);
 498   appendToGlobalCtors(M, MemProfCtorFunction, Priority);
 499
 500   createProfileFileNameVar(M);
 501
 502   return true;
 503 }
 504
 505 void MemProfiler::initializeCallbacks(Module &M) {
 506   IRBuilder<> IRB(*C);
 507
 508   for (size_t AccessIsWrite = 0; AccessIsWrite <= 1; AccessIsWrite++) {
 509     const std::string TypeStr = AccessIsWrite ? "store" : "load";
 510
 511     SmallVector<Type *, 3> Args2 = {IntptrTy, IntptrTy};
 512     SmallVector<Type *, 2> Args1{1, IntptrTy};
 513     MemProfMemoryAccessCallbackSized[AccessIsWrite] =
 514         M.getOrInsertFunction(ClMemoryAccessCallbackPrefix + TypeStr + "N",
 515                               FunctionType::get(IRB.getVoidTy(), Args2, false));
 516
 517     MemProfMemoryAccessCallback[AccessIsWrite] =
 518         M.getOrInsertFunction(ClMemoryAccessCallbackPrefix + TypeStr,
 519                               FunctionType::get(IRB.getVoidTy(), Args1, false));
 520   }
 521   MemProfMemmove = M.getOrInsertFunction(
 522       ClMemoryAccessCallbackPrefix + "memmove", IRB.getInt8PtrTy(),
 523       IRB.getInt8PtrTy(), IRB.getInt8PtrTy(), IntptrTy);
 524   MemProfMemcpy = M.getOrInsertFunction(ClMemoryAccessCallbackPrefix + "memcpy",
 525                                         IRB.getInt8PtrTy(), IRB.getInt8PtrTy(),
 526                                         IRB.getInt8PtrTy(), IntptrTy);
 527   MemProfMemset = M.getOrInsertFunction(ClMemoryAccessCallbackPrefix + "memset",
 528                                         IRB.getInt8PtrTy(), IRB.getInt8PtrTy(),
 529                                         IRB.getInt32Ty(), IntptrTy);
 530 }
 531
 532 bool MemProfiler::maybeInsertMemProfInitAtFunctionEntry(Function &F) {
 533   // For each NSObject descendant having a +load method, this method is invoked
 534   // by the ObjC runtime before any of the static constructors is called.
 535   // Therefore we need to instrument such methods with a call to __memprof_init
 536   // at the beginning in order to initialize our runtime before any access to
 537   // the shadow memory.
 538   // We cannot just ignore these methods, because they may call other
 539   // instrumented functions.
 540   if (F.getName().find(" load]") != std::string::npos) {
 541     FunctionCallee MemProfInitFunction =
 542         declareSanitizerInitFunction(*F.getParent(), MemProfInitName, {});
 543     IRBuilder<> IRB(&F.front(), F.front().begin());
 544     IRB.CreateCall(MemProfInitFunction, {});
 545     return true;
 546   }
 547   return false;
 548 }
 549
 550 bool MemProfiler::insertDynamicShadowAtFunctionEntry(Function &F) {
 551   IRBuilder<> IRB(&F.front().front());
 552   Value *GlobalDynamicAddress = F.getParent()->getOrInsertGlobal(
 553       MemProfShadowMemoryDynamicAddress, IntptrTy);
 554   if (F.getParent()->getPICLevel() == PICLevel::NotPIC)
 555     cast<GlobalVariable>(GlobalDynamicAddress)->setDSOLocal(true);
 556   DynamicShadowOffset = IRB.CreateLoad(IntptrTy, GlobalDynamicAddress);
 557   return true;
 558 }
 559
 560 bool MemProfiler::instrumentFunction(Function &F) {
 561   if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage)
 562     return false;
 563   if (ClDebugFunc == F.getName())
 564     return false;
 565   if (F.getName().startswith("__memprof_"))
 566     return false;
 567
 568   bool FunctionModified = false;
 569
 570   // If needed, insert __memprof_init.
 571   // This function needs to be called even if the function body is not
 572   // instrumented.
 573   if (maybeInsertMemProfInitAtFunctionEntry(F))
 574     FunctionModified = true;
 575
 576   LLVM_DEBUG(dbgs() << "MEMPROF instrumenting:\n" << F << "\n");
 577
 578   initializeCallbacks(*F.getParent());
 579
 580   SmallVector<Instruction *, 16> ToInstrument;
 581
 582   // Fill the set of memory operations to instrument.
 583   for (auto &BB : F) {
 584     for (auto &Inst : BB) {
 585       if (isInterestingMemoryAccess(&Inst) || isa<MemIntrinsic>(Inst))
 586         ToInstrument.push_back(&Inst);
 587     }
 588   }
 589
 590   if (ToInstrument.empty()) {
 591     LLVM_DEBUG(dbgs() << "MEMPROF done instrumenting: " << FunctionModified
 592                       << " " << F << "\n");
 593
 594     return FunctionModified;
 595   }
 596
 597   FunctionModified |= insertDynamicShadowAtFunctionEntry(F);
 598
 599   int NumInstrumented = 0;
 600   for (auto *Inst : ToInstrument) {
 601     if (ClDebugMin < 0 || ClDebugMax < 0 ||
 602         (NumInstrumented >= ClDebugMin && NumInstrumented <= ClDebugMax)) {
 603       std::optional<InterestingMemoryAccess> Access =
 604           isInterestingMemoryAccess(Inst);
 605       if (Access)
 606         instrumentMop(Inst, F.getParent()->getDataLayout(), *Access);
 607       else
 608         instrumentMemIntrinsic(cast<MemIntrinsic>(Inst));
 609     }
 610     NumInstrumented++;
 611   }
 612
 613   if (NumInstrumented > 0)
 614     FunctionModified = true;
 615
 616   LLVM_DEBUG(dbgs() << "MEMPROF done instrumenting: " << FunctionModified << " "
 617                     << F << "\n");
 618
 619   return FunctionModified;
 620 }
 621
 622 static void addCallsiteMetadata(Instruction &I,
 623                                 std::vector<uint64_t> &InlinedCallStack,
 624                                 LLVMContext &Ctx) {
 625   I.setMetadata(LLVMContext::MD_callsite,
 626                 buildCallstackMetadata(InlinedCallStack, Ctx));
 627 }
 628
 629 static uint64_t computeStackId(GlobalValue::GUID Function, uint32_t LineOffset,
 630                                uint32_t Column) {
 631   llvm::HashBuilder<llvm::TruncatedBLAKE3<8>, llvm::support::endianness::little>
 632       HashBuilder;
 633   HashBuilder.add(Function, LineOffset, Column);
 634   llvm::BLAKE3Result<8> Hash = HashBuilder.final();
 635   uint64_t Id;
 636   std::memcpy(&Id, Hash.data(), sizeof(Hash));
 637   return Id;
 638 }
 639
 640 static uint64_t computeStackId(const memprof::Frame &Frame) {
 641   return computeStackId(Frame.Function, Frame.LineOffset, Frame.Column);
 642 }
 643
 644 static void addCallStack(CallStackTrie &AllocTrie,
 645                          const AllocationInfo *AllocInfo) {
 646   SmallVector<uint64_t> StackIds;
 647   for (const auto &StackFrame : AllocInfo->CallStack)
 648     StackIds.push_back(computeStackId(StackFrame));
 649   auto AllocType = getAllocType(AllocInfo->Info.getTotalLifetimeAccessDensity(),
 650                                 AllocInfo->Info.getAllocCount(),
 651                                 AllocInfo->Info.getTotalLifetime());
 652   AllocTrie.addCallStack(AllocType, StackIds);
 653 }
 654
 655 // Helper to compare the InlinedCallStack computed from an instruction's debug
 656 // info to a list of Frames from profile data (either the allocation data or a
 657 // callsite). For callsites, the StartIndex to use in the Frame array may be
 658 // non-zero.
 659 static bool
 660 stackFrameIncludesInlinedCallStack(ArrayRef<Frame> ProfileCallStack,
 661                                    ArrayRef<uint64_t> InlinedCallStack,
 662                                    unsigned StartIndex = 0) {
 663   auto StackFrame = ProfileCallStack.begin() + StartIndex;
 664   auto InlCallStackIter = InlinedCallStack.begin();
 665   for (; StackFrame != ProfileCallStack.end() &&
 666          InlCallStackIter != InlinedCallStack.end();
 667        ++StackFrame, ++InlCallStackIter) {
 668     uint64_t StackId = computeStackId(*StackFrame);
 669     if (StackId != *InlCallStackIter)
 670       return false;
 671   }
 672   // Return true if we found and matched all stack ids from the call
 673   // instruction.
 674   return InlCallStackIter == InlinedCallStack.end();
 675 }
 676
 677 static void readMemprof(Module &M, Function &F,
 678                         IndexedInstrProfReader *MemProfReader,
 679                         const TargetLibraryInfo &TLI) {
 680   auto &Ctx = M.getContext();
 681
 682   auto FuncName = getIRPGOFuncName(F);
 683   auto FuncGUID = Function::getGUID(FuncName);
 684   std::optional<memprof::MemProfRecord> MemProfRec;
 685   auto Err = MemProfReader->getMemProfRecord(FuncGUID).moveInto(MemProfRec);
 686   if (Err) {
 687     // If we don't find getIRPGOFuncName(), try getPGOFuncName() to handle
 688     // profiles built by older compilers
 689     Err = handleErrors(std::move(Err), [&](const InstrProfError &IE) -> Error {
 690       if (IE.get() != instrprof_error::unknown_function)
 691         return make_error<InstrProfError>(IE);
 692       auto FuncName = getPGOFuncName(F);
 693       auto FuncGUID = Function::getGUID(FuncName);
 694       if (auto Err =
 695               MemProfReader->getMemProfRecord(FuncGUID).moveInto(MemProfRec))
 696         return Err;
 697       return Error::success();
 698     });
 699   }
 700   if (Err) {
 701     handleAllErrors(std::move(Err), [&](const InstrProfError &IPE) {
 702       auto Err = IPE.get();
 703       bool SkipWarning = false;
 704       LLVM_DEBUG(dbgs() << "Error in reading profile for Func " << FuncName
 705                         << ": ");
 706       if (Err == instrprof_error::unknown_function) {
 707         NumOfMemProfMissing++;
 708         SkipWarning = !PGOWarnMissing;
 709         LLVM_DEBUG(dbgs() << "unknown function");
 710       } else if (Err == instrprof_error::hash_mismatch) {
 711         SkipWarning =
 712             NoPGOWarnMismatch ||
 713             (NoPGOWarnMismatchComdatWeak &&
 714              (F.hasComdat() ||
 715               F.getLinkage() == GlobalValue::AvailableExternallyLinkage));
 716         LLVM_DEBUG(dbgs() << "hash mismatch (skip=" << SkipWarning << ")");
 717       }
 718
 719       if (SkipWarning)
 720         return;
 721
 722       std::string Msg = (IPE.message() + Twine(" ") + F.getName().str() +
 723                          Twine(" Hash = ") + std::to_string(FuncGUID))
 724                             .str();
 725
 726       Ctx.diagnose(
 727           DiagnosticInfoPGOProfile(M.getName().data(), Msg, DS_Warning));
 728     });
 729     return;
 730   }
 731
 732   // Build maps of the location hash to all profile data with that leaf location
 733   // (allocation info and the callsites).
 734   std::map<uint64_t, std::set<const AllocationInfo *>> LocHashToAllocInfo;
 735   // For the callsites we need to record the index of the associated frame in
 736   // the frame array (see comments below where the map entries are added).
 737   std::map<uint64_t, std::set<std::pair<const SmallVector<Frame> *, unsigned>>>
 738       LocHashToCallSites;
 739   for (auto &AI : MemProfRec->AllocSites) {
 740     // Associate the allocation info with the leaf frame. The later matching
 741     // code will match any inlined call sequences in the IR with a longer prefix
 742     // of call stack frames.
 743     uint64_t StackId = computeStackId(AI.CallStack[0]);
 744     LocHashToAllocInfo[StackId].insert(&AI);
 745   }
 746   for (auto &CS : MemProfRec->CallSites) {
 747     // Need to record all frames from leaf up to and including this function,
 748     // as any of these may or may not have been inlined at this point.
 749     unsigned Idx = 0;
 750     for (auto &StackFrame : CS) {
 751       uint64_t StackId = computeStackId(StackFrame);
 752       LocHashToCallSites[StackId].insert(std::make_pair(&CS, Idx++));
 753       // Once we find this function, we can stop recording.
 754       if (StackFrame.Function == FuncGUID)
 755         break;
 756     }
 757     assert(Idx <= CS.size() && CS[Idx - 1].Function == FuncGUID);
 758   }
 759
 760   auto GetOffset = [](const DILocation *DIL) {
 761     return (DIL->getLine() - DIL->getScope()->getSubprogram()->getLine()) &
 762            0xffff;
 763   };
 764
 765   // Now walk the instructions, looking up the associated profile data using
 766   // dbug locations.
 767   for (auto &BB : F) {
 768     for (auto &I : BB) {
 769       if (I.isDebugOrPseudoInst())
 770         continue;
 771       // We are only interested in calls (allocation or interior call stack
 772       // context calls).
 773       auto *CI = dyn_cast<CallBase>(&I);
 774       if (!CI)
 775         continue;
 776       auto *CalledFunction = CI->getCalledFunction();
 777       if (CalledFunction && CalledFunction->isIntrinsic())
 778         continue;
 779       // List of call stack ids computed from the location hashes on debug
 780       // locations (leaf to inlined at root).
 781       std::vector<uint64_t> InlinedCallStack;
 782       // Was the leaf location found in one of the profile maps?
 783       bool LeafFound = false;
 784       // If leaf was found in a map, iterators pointing to its location in both
 785       // of the maps. It might exist in neither, one, or both (the latter case
 786       // can happen because we don't currently have discriminators to
 787       // distinguish the case when a single line/col maps to both an allocation
 788       // and another callsite).
 789       std::map<uint64_t, std::set<const AllocationInfo *>>::iterator
 790           AllocInfoIter;
 791       std::map<uint64_t, std::set<std::pair<const SmallVector<Frame> *,
 792                                             unsigned>>>::iterator CallSitesIter;
 793       for (const DILocation *DIL = I.getDebugLoc(); DIL != nullptr;
 794            DIL = DIL->getInlinedAt()) {
 795         // Use C++ linkage name if possible. Need to compile with
 796         // -fdebug-info-for-profiling to get linkage name.
 797         StringRef Name = DIL->getScope()->getSubprogram()->getLinkageName();
 798         if (Name.empty())
 799           Name = DIL->getScope()->getSubprogram()->getName();
 800         auto CalleeGUID = Function::getGUID(Name);
 801         auto StackId =
 802             computeStackId(CalleeGUID, GetOffset(DIL), DIL->getColumn());
 803         // LeafFound will only be false on the first iteration, since we either
 804         // set it true or break out of the loop below.
 805         if (!LeafFound) {
 806           AllocInfoIter = LocHashToAllocInfo.find(StackId);
 807           CallSitesIter = LocHashToCallSites.find(StackId);
 808           // Check if the leaf is in one of the maps. If not, no need to look
 809           // further at this call.
 810           if (AllocInfoIter == LocHashToAllocInfo.end() &&
 811               CallSitesIter == LocHashToCallSites.end())
 812             break;
 813           LeafFound = true;
 814         }
 815         InlinedCallStack.push_back(StackId);
 816       }
 817       // If leaf not in either of the maps, skip inst.
 818       if (!LeafFound)
 819         continue;
 820
 821       // First add !memprof metadata from allocation info, if we found the
 822       // instruction's leaf location in that map, and if the rest of the
 823       // instruction's locations match the prefix Frame locations on an
 824       // allocation context with the same leaf.
 825       if (AllocInfoIter != LocHashToAllocInfo.end()) {
 826         // Only consider allocations via new, to reduce unnecessary metadata,
 827         // since those are the only allocations that will be targeted initially.
 828         if (!isNewLikeFn(CI, &TLI))
 829           continue;
 830         // We may match this instruction's location list to multiple MIB
 831         // contexts. Add them to a Trie specialized for trimming the contexts to
 832         // the minimal needed to disambiguate contexts with unique behavior.
 833         CallStackTrie AllocTrie;
 834         for (auto *AllocInfo : AllocInfoIter->second) {
 835           // Check the full inlined call stack against this one.
 836           // If we found and thus matched all frames on the call, include
 837           // this MIB.
 838           if (stackFrameIncludesInlinedCallStack(AllocInfo->CallStack,
 839                                                  InlinedCallStack))
 840             addCallStack(AllocTrie, AllocInfo);
 841         }
 842         // We might not have matched any to the full inlined call stack.
 843         // But if we did, create and attach metadata, or a function attribute if
 844         // all contexts have identical profiled behavior.
 845         if (!AllocTrie.empty()) {
 846           // MemprofMDAttached will be false if a function attribute was
 847           // attached.
 848           bool MemprofMDAttached = AllocTrie.buildAndAttachMIBMetadata(CI);
 849           assert(MemprofMDAttached == I.hasMetadata(LLVMContext::MD_memprof));
 850           if (MemprofMDAttached) {
 851             // Add callsite metadata for the instruction's location list so that
 852             // it simpler later on to identify which part of the MIB contexts
 853             // are from this particular instruction (including during inlining,
 854             // when the callsite metdata will be updated appropriately).
 855             // FIXME: can this be changed to strip out the matching stack
 856             // context ids from the MIB contexts and not add any callsite
 857             // metadata here to save space?
 858             addCallsiteMetadata(I, InlinedCallStack, Ctx);
 859           }
 860         }
 861         continue;
 862       }
 863
 864       // Otherwise, add callsite metadata. If we reach here then we found the
 865       // instruction's leaf location in the callsites map and not the allocation
 866       // map.
 867       assert(CallSitesIter != LocHashToCallSites.end());
 868       for (auto CallStackIdx : CallSitesIter->second) {
 869         // If we found and thus matched all frames on the call, create and
 870         // attach call stack metadata.
 871         if (stackFrameIncludesInlinedCallStack(
 872                 *CallStackIdx.first, InlinedCallStack, CallStackIdx.second)) {
 873           addCallsiteMetadata(I, InlinedCallStack, Ctx);
 874           // Only need to find one with a matching call stack and add a single
 875           // callsite metadata.
 876           break;
 877         }
 878       }
 879     }
 880   }
 881 }
 882
 883 MemProfUsePass::MemProfUsePass(std::string MemoryProfileFile,
 884                                IntrusiveRefCntPtr<vfs::FileSystem> FS)
 885     : MemoryProfileFileName(MemoryProfileFile), FS(FS) {
 886   if (!FS)
 887     this->FS = vfs::getRealFileSystem();
 888 }
 889
 890 PreservedAnalyses MemProfUsePass::run(Module &M, ModuleAnalysisManager &AM) {
 891   LLVM_DEBUG(dbgs() << "Read in memory profile:");
 892   auto &Ctx = M.getContext();
 893   auto ReaderOrErr = IndexedInstrProfReader::create(MemoryProfileFileName, *FS);
 894   if (Error E = ReaderOrErr.takeError()) {
 895     handleAllErrors(std::move(E), [&](const ErrorInfoBase &EI) {
 896       Ctx.diagnose(
 897           DiagnosticInfoPGOProfile(MemoryProfileFileName.data(), EI.message()));
 898     });
 899     return PreservedAnalyses::all();
 900   }
 901
 902   std::unique_ptr<IndexedInstrProfReader> MemProfReader =
 903       std::move(ReaderOrErr.get());
 904   if (!MemProfReader) {
 905     Ctx.diagnose(DiagnosticInfoPGOProfile(
 906         MemoryProfileFileName.data(), StringRef("Cannot get MemProfReader")));
 907     return PreservedAnalyses::all();
 908   }
 909
 910   if (!MemProfReader->hasMemoryProfile()) {
 911     Ctx.diagnose(DiagnosticInfoPGOProfile(MemoryProfileFileName.data(),
 912                                           "Not a memory profile"));
 913     return PreservedAnalyses::all();
 914   }
 915
 916   auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
 917
 918   for (auto &F : M) {
 919     if (F.isDeclaration())
 920       continue;
 921
 922     const TargetLibraryInfo &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
 923     readMemprof(M, F, MemProfReader.get(), TLI);
 924   }
 925
 926   return PreservedAnalyses::none();
 927 }