[InstrProf][NFC] Save profile bias to function map

[ellis]

Summary

Add a map from functions to load instructions that compute the profile bias. Previously we assumed that if the first instruction in the function was a load instruction, then it must be computing the bias. This was likely to work out because functions usually start with the llvm.instrprof.increment instruction, but optimizations could change this. For example, inlining into a non-profiled function.

[ellis]

Run clang-format.

[ellis]

Rebase.

[phosek]

Thanks, alternative solution we've implemented in D114383 is to name the load instruction and then check the name before reusing it, but this solution is more precise.

[phosek]

LGTM

@ellis do you have time to land this change? This change should address https://github.com/llvm/llvm-project/issues/55125 so I'd like to see it landed.

llvm/include/llvm/Transforms/Instrumentation/InstrProfiling.h

61

I'd prefer FunctionToProfileBiasMap to be consistent with other variables like ProfileDataMap.

[ellis]

In D114319#3512308, @phosek wrote:

LGTM

@ellis do you have time to land this change? This change should address https://github.com/llvm/llvm-project/issues/55125 so I'd like to see it landed.

Sure, I'll try to land this now!

[ellis]

Rename function map and rebase

[ellis]

Rebase

[ellis]

@phosek I've landed this, but it does not seem to fix https://github.com/llvm/llvm-project/issues/55125

I believe the error is related to -do-counter-promotion since setting that to false makes the crash go away. I'll comment in the github issue.

llvm/include/llvm/Transforms/Instrumentation/InstrProfiling.h

struct PerFunctionProfileData {
     GlobalVariable *RegionCounters = nullptr;
     GlobalVariable *DataVar = nullptr;
 
     PerFunctionProfileData() {
       memset(NumValueSites, 0, sizeof(uint32_t) * (IPVK_Last + 1));
     }
   };
   DenseMap<GlobalVariable *, PerFunctionProfileData> ProfileDataMap;
+  /// If runtime relocation is enabled, this maps functions to the load
+  /// instruction that produces the profile relocation bias.
+  DenseMap<const Function *, LoadInst *> FunctionToProfileBiasMap;
   std::vector<GlobalValue *> CompilerUsedVars;
   std::vector<GlobalValue *> UsedVars;
   std::vector<GlobalVariable *> ReferencedNames;
   GlobalVariable *NamesVar;
   size_t NamesSize;
 
   // Is this lowering for the context-sensitive instrumentation.
   bool IsCS;

llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp

Value *InstrProfiling::getCounterAddress(InstrProfInstBase *I) {
   auto *Addr = Builder.CreateConstInBoundsGEP2_32(
       Counters->getValueType(), Counters, 0, I->getIndex()->getZExtValue());
 
   if (!isRuntimeCounterRelocationEnabled())
     return Addr;
 
   Type *Int64Ty = Type::getInt64Ty(M->getContext());
   Function *Fn = I->getParent()->getParent();
-  Instruction &EntryI = Fn->getEntryBlock().front();
-  LoadInst *LI = dyn_cast<LoadInst>(&EntryI);
-  if (!LI) {
-    IRBuilder<> EntryBuilder(&EntryI);
+  LoadInst *&BiasLI = FunctionToProfileBiasMap[Fn];
+  if (!BiasLI) {
+    IRBuilder<> EntryBuilder(&Fn->getEntryBlock().front());
     auto *Bias = M->getGlobalVariable(getInstrProfCounterBiasVarName());
     if (!Bias) {
       // Compiler must define this variable when runtime counter relocation
       // is being used. Runtime has a weak external reference that is used
       // to check whether that's the case or not.
       Bias = new GlobalVariable(
           *M, Int64Ty, false, GlobalValue::LinkOnceODRLinkage,
           Constant::getNullValue(Int64Ty), getInstrProfCounterBiasVarName());
       Bias->setVisibility(GlobalVariable::HiddenVisibility);
       // A definition that's weak (linkonce_odr) without being in a COMDAT
       // section wouldn't lead to link errors, but it would lead to a dead
       // data word from every TU but one. Putting it in COMDAT ensures there
       // will be exactly one data slot in the link.
       if (TT.supportsCOMDAT())
         Bias->setComdat(M->getOrInsertComdat(Bias->getName()));
     }
-    LI = EntryBuilder.CreateLoad(Int64Ty, Bias);
+    BiasLI = EntryBuilder.CreateLoad(Int64Ty, Bias);
   }
-  auto *Add = Builder.CreateAdd(Builder.CreatePtrToInt(Addr, Int64Ty), LI);
+  auto *Add = Builder.CreateAdd(Builder.CreatePtrToInt(Addr, Int64Ty), BiasLI);
   return Builder.CreateIntToPtr(Add, Addr->getType());
 }
 
 void InstrProfiling::lowerCover(InstrProfCoverInst *CoverInstruction) {
   auto *Addr = getCounterAddress(CoverInstruction);
   IRBuilder<> Builder(CoverInstruction);
   // We store zero to represent that this block is covered.
   Builder.CreateStore(Builder.getInt8(0), Addr);