[PGO] Don't set the function hotness attribute when populating counters

[xur]

Summary

Don't set the function hotness attribute on the fly. This changes the CFG branch probability of the caller function, which leads to inconsistent BB ordering. This patch moves the attribute setting to a separated loop after the counts in all functions are populated.

Fixes PR27024 - PGO instrumentation profile data is not reflected in correct basic blocks.

[davidxl]

The logic to set InlineHint should be removed and handled by inline cost analysis (FE based pgo removed it already) -- but that is a separate patch.

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

463 ↗

(On Diff #51715)

Initialize FreqAttr to Normal

481 ↗

(On Diff #51715)

FFA_Normal,
FFA_Cold,
FFA_Hot

519 ↗

(On Diff #51715)

Move this closer to other private data members.

[xur]

Integrated David's review comments.

Thanks,

-Rong

[davidxl]

lgtm

llvm/trunk/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

for (auto &E : Edges) {
     if (E->Removed)
       continue;
     Total += E->CountValue;
   }
   return Total;
 }
 
 class PGOUseFunc {
+public:
+  PGOUseFunc(Function &Func, Module *Modu, BranchProbabilityInfo *BPI = nullptr,
+             BlockFrequencyInfo *BFI = nullptr)
+      : F(Func), M(Modu), FuncInfo(Func, false, BPI, BFI),
+        FreqAttr(FFA_Normal) {}
+
+  // Read counts for the instrumented BB from profile.
+  bool readCounters(IndexedInstrProfReader *PGOReader);
+
+  // Populate the counts for all BBs.
+  void populateCounters();
+
+  // Set the branch weights based on the count values.
+  void setBranchWeights();
+
+  // Annotate the indirect call sites.
+  void annotateIndirectCallSites();
+
+  // The hotness of the function from the profile count.
+  enum FuncFreqAttr { FFA_Normal, FFA_Cold, FFA_Hot };
+
+  // Return the funtion hotness from the profile.
+  FuncFreqAttr getFuncFreqAttr() const {
+    return FreqAttr;
+  }
+
 private:
   Function &F;
   Module *M;
   // This member stores the shared information with class PGOGenFunc.
   FuncPGOInstrumentation<PGOUseEdge, UseBBInfo> FuncInfo;
 
   // Return the auxiliary BB information.
   UseBBInfo &getBBInfo(const BasicBlock *BB) const {
     return FuncInfo.getBBInfo(BB);
   }
 
   // The maximum count value in the profile. This is only used in PGO use
   // compilation.
   uint64_t ProgramMaxCount;
 
   // ProfileRecord for this function.
   InstrProfRecord ProfileRecord;
 
+  // Function hotness info derived from profile.
+  FuncFreqAttr FreqAttr;
+
   // Find the Instrumented BB and set the value.
   void setInstrumentedCounts(const std::vector<uint64_t> &CountFromProfile);
 
   // Set the edge counter value for the unknown edge -- there should be only
   // one unknown edge.
   void setEdgeCount(DirectEdges &Edges, uint64_t Value);
 
   // Return FuncName string;
   const std::string getFuncName() const { return FuncInfo.FuncName; }
 
   // Set the hot/cold inline hints based on the count values.
   // FIXME: This function should be removed once the functionality in
   // the inliner is implemented.
-  void applyFunctionAttributes(uint64_t EntryCount, uint64_t MaxCount) {
+  void markFunctionAttributes(uint64_t EntryCount, uint64_t MaxCount) {
     if (ProgramMaxCount == 0)
       return;
     // Threshold of the hot functions.
     const BranchProbability HotFunctionThreshold(1, 100);
     // Threshold of the cold functions.
     const BranchProbability ColdFunctionThreshold(2, 10000);
     if (EntryCount >= HotFunctionThreshold.scale(ProgramMaxCount))
-      F.addFnAttr(llvm::Attribute::InlineHint);
+      FreqAttr = FFA_Hot;
     else if (MaxCount <= ColdFunctionThreshold.scale(ProgramMaxCount))
-      F.addFnAttr(llvm::Attribute::Cold);
+      FreqAttr = FFA_Cold;
   }
-
-public:
-  PGOUseFunc(Function &Func, Module *Modu, BranchProbabilityInfo *BPI = nullptr,
-             BlockFrequencyInfo *BFI = nullptr)
-      : F(Func), M(Modu), FuncInfo(Func, false, BPI, BFI) {}
-
-  // Read counts for the instrumented BB from profile.
-  bool readCounters(IndexedInstrProfReader *PGOReader);
-
-  // Populate the counts for all BBs.
-  void populateCounters();
-
-  // Set the branch weights based on the count values.
-  void setBranchWeights();
-
-  // Annotate the indirect call sites.
-  void annotateIndirectCallSites();
 };
 
 // Visit all the edges and assign the count value for the instrumented
 // edges and the BB.
 void PGOUseFunc::setInstrumentedCounts(
     const std::vector<uint64_t> &CountFromProfile) {
 
   // Use a worklist as we will update the vector during the iteration.
void PGOUseFunc::populateCounters() {
   uint64_t FuncEntryCount = getBBInfo(&*F.begin()).CountValue;
   uint64_t FuncMaxCount = FuncEntryCount;
   for (auto &BB : F) {
     assert(getBBInfo(&BB).CountValid && "BB count is not valid");
     uint64_t Count = getBBInfo(&BB).CountValue;
     if (Count > FuncMaxCount)
       FuncMaxCount = Count;
   }
-  applyFunctionAttributes(FuncEntryCount, FuncMaxCount);
+  markFunctionAttributes(FuncEntryCount, FuncMaxCount);
 
   DEBUG(FuncInfo.dumpInfo("after reading profile."));
 }
 
 // Assign the scaled count values to the BB with multiple out edges.
 void PGOUseFunc::setBranchWeights() {
   // Generate MD_prof metadata for every branch instruction.
   DEBUG(dbgs() << "\nSetting branch weights.\n");
bool PGOInstrumentationUse::runOnModule(Module &M) {
   }
   // TODO: might need to change the warning once the clang option is finalized.
   if (!PGOReader->isIRLevelProfile()) {
     Ctx.diagnose(DiagnosticInfoPGOProfile(
         ProfileFileName.data(), "Not an IR level instrumentation profile"));
     return false;
   }
 
+  std::vector<Function *> HotFunctions;
+  std::vector<Function *> ColdFunctions;
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
     BranchProbabilityInfo *BPI =
         &(getAnalysis<BranchProbabilityInfoWrapperPass>(F).getBPI());
     BlockFrequencyInfo *BFI =
         &(getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI());
     PGOUseFunc Func(F, &M, BPI, BFI);
     setPGOCountOnFunc(Func, PGOReader.get());
+    PGOUseFunc::FuncFreqAttr FreqAttr = Func.getFuncFreqAttr();
+    if (FreqAttr == PGOUseFunc::FFA_Cold)
+      ColdFunctions.push_back(&F);
+    else if (FreqAttr == PGOUseFunc::FFA_Hot)
+      HotFunctions.push_back(&F);
   }
+
+  // Set function hotness attribute from the profile.
+  for (auto &F : HotFunctions) {
+    F->addFnAttr(llvm::Attribute::InlineHint);
+    DEBUG(dbgs() << "Set inline attribute to function: " << F->getName()
+                 << "\n");
+  }
+  for (auto &F : ColdFunctions) {
+    F->addFnAttr(llvm::Attribute::Cold);
+    DEBUG(dbgs() << "Set cold attribute to function: " << F->getName() << "\n");
+  }
+
   return true;
 }