1 //===-- PGOInstrumentation.cpp - MST-based PGO Instrumentation ------------===//
2 //
3 //                      The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file implements PGO instrumentation using a minimum spanning tree based
11 // on the following paper:
12 //   [1] Donald E. Knuth, Francis R. Stevenson. Optimal measurement of points
13 //   for program frequency counts. BIT Numerical Mathematics 1973, Volume 13,
14 //   Issue 3, pp 313-322
15 // The idea of the algorithm based on the fact that for each node (except for
16 // the entry and exit), the sum of incoming edge counts equals the sum of
17 // outgoing edge counts. The count of edge on spanning tree can be derived from
18 // those edges not on the spanning tree. Knuth proves this method instruments
19 // the minimum number of edges.
20 //
21 // The minimal spanning tree here is actually a maximum weight tree -- on-tree
22 // edges have higher frequencies (more likely to execute). The idea is to
23 // instrument those less frequently executed edges to reduce the runtime
24 // overhead of instrumented binaries.
25 //
26 // This file contains two passes:
27 // (1) Pass PGOInstrumentationGen which instruments the IR to generate edge
28 // count profile, and
29 // (2) Pass PGOInstrumentationUse which reads the edge count profile and
30 // annotates the branch weights.
31 // To get the precise counter information, These two passes need to invoke at
32 // the same compilation point (so they see the same IR). For pass
33 // PGOInstrumentationGen, the real work is done in instrumentOneFunc(). For
34 // pass PGOInstrumentationUse, the real work in done in class PGOUseFunc and
35 // the profile is opened in module level and passed to each PGOUseFunc instance.
36 // The shared code for PGOInstrumentationGen and PGOInstrumentationUse is put
37 // in class FuncPGOInstrumentation.
38 //
39 // Class PGOEdge represents a CFG edge and some auxiliary information. Class
40 // BBInfo contains auxiliary information for each BB. These two classes are used
41 // in pass PGOInstrumentationGen. Class PGOUseEdge and UseBBInfo are the derived
42 // class of PGOEdge and BBInfo, respectively. They contains extra data structure
43 // used in populating profile counters.
44 // The MST implementation is in Class CFGMST (CFGMST.h).
45 //
46 //===----------------------------------------------------------------------===//
47 
48 #include "CFGMST.h"
49 #include "llvm/ADT/DenseMap.h"
50 #include "llvm/ADT/STLExtras.h"
51 #include "llvm/ADT/Statistic.h"
52 #include "llvm/Analysis/BlockFrequencyInfo.h"
53 #include "llvm/Analysis/BranchProbabilityInfo.h"
54 #include "llvm/Analysis/CFG.h"
55 #include "llvm/IR/CallSite.h"
56 #include "llvm/IR/DiagnosticInfo.h"
57 #include "llvm/IR/IRBuilder.h"
58 #include "llvm/IR/InstIterator.h"
59 #include "llvm/IR/InstVisitor.h"
60 #include "llvm/IR/Instructions.h"
61 #include "llvm/IR/IntrinsicInst.h"
62 #include "llvm/IR/MDBuilder.h"
63 #include "llvm/IR/Module.h"
64 #include "llvm/Pass.h"
65 #include "llvm/ProfileData/InstrProfReader.h"
66 #include "llvm/Support/BranchProbability.h"
67 #include "llvm/Support/Debug.h"
68 #include "llvm/Support/JamCRC.h"
69 #include "llvm/Transforms/Instrumentation.h"
70 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
71 #include <string>
72 #include <utility>
73 #include <vector>
74 
75 using namespace llvm;
76 
77 #define DEBUG_TYPE "pgo-instrumentation"
78 
79 STATISTIC(NumOfPGOInstrument, "Number of edges instrumented.");
80 STATISTIC(NumOfPGOEdge, "Number of edges.");
81 STATISTIC(NumOfPGOBB, "Number of basic-blocks.");
82 STATISTIC(NumOfPGOSplit, "Number of critical edge splits.");
83 STATISTIC(NumOfPGOFunc, "Number of functions having valid profile counts.");
84 STATISTIC(NumOfPGOMismatch, "Number of functions having mismatch profile.");
85 STATISTIC(NumOfPGOMissing, "Number of functions without profile.");
86 STATISTIC(NumOfPGOICall, "Number of indirect call value instrumentation.");
87 
88 // Command line option to specify the file to read profile from. This is
89 // mainly used for testing.
90 static cl::opt<std::string>
91     PGOTestProfileFile("pgo-test-profile-file", cl::init(""), cl::Hidden,
92                        cl::value_desc("filename"),
93                        cl::desc("Specify the path of profile data file. This is"
94                                 "mainly for test purpose."));
95 
96 // Command line options to disable value profiling. The default is false:
97 // i.e. vaule profiling is enabled by default. This is for debug purpose.
98 static cl::opt<bool>
99 DisableValueProfiling("disable-vp", cl::init(false),
100                                     cl::Hidden,
101                                     cl::desc("Disable Value Profiling"));
102 
103 namespace {
104 class PGOInstrumentationGen : public ModulePass {
105 public:
106   static char ID;
107 
108   PGOInstrumentationGen() : ModulePass(ID) {
109     initializePGOInstrumentationGenPass(*PassRegistry::getPassRegistry());
110   }
111 
112   const char *getPassName() const override {
113     return "PGOInstrumentationGenPass";
114   }
115 
116 private:
117   bool runOnModule(Module &M) override;
118 
119   void getAnalysisUsage(AnalysisUsage &AU) const override {
120     AU.addRequired<BlockFrequencyInfoWrapperPass>();
121   }
122 };
123 
124 class PGOInstrumentationUse : public ModulePass {
125 public:
126   static char ID;
127 
128   // Provide the profile filename as the parameter.
129   PGOInstrumentationUse(std::string Filename = "")
130       : ModulePass(ID), ProfileFileName(Filename) {
131     if (!PGOTestProfileFile.empty())
132       ProfileFileName = PGOTestProfileFile;
133     initializePGOInstrumentationUsePass(*PassRegistry::getPassRegistry());
134   }
135 
136   const char *getPassName() const override {
137     return "PGOInstrumentationUsePass";
138   }
139 
140 private:
141   std::string ProfileFileName;
142   std::unique_ptr<IndexedInstrProfReader> PGOReader;
143   bool runOnModule(Module &M) override;
144 
145   void getAnalysisUsage(AnalysisUsage &AU) const override {
146     AU.addRequired<BlockFrequencyInfoWrapperPass>();
147   }
148 };
149 } // end anonymous namespace
150 
151 char PGOInstrumentationGen::ID = 0;
152 INITIALIZE_PASS_BEGIN(PGOInstrumentationGen, "pgo-instr-gen",
153                       "PGO instrumentation.", false, false)
154 INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
155 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
156 INITIALIZE_PASS_END(PGOInstrumentationGen, "pgo-instr-gen",
157                     "PGO instrumentation.", false, false)
158 
159 ModulePass *llvm::createPGOInstrumentationGenPass() {
160   return new PGOInstrumentationGen();
161 }
162 
163 char PGOInstrumentationUse::ID = 0;
164 INITIALIZE_PASS_BEGIN(PGOInstrumentationUse, "pgo-instr-use",
165                       "Read PGO instrumentation profile.", false, false)
166 INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
167 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
168 INITIALIZE_PASS_END(PGOInstrumentationUse, "pgo-instr-use",
169                     "Read PGO instrumentation profile.", false, false)
170 
171 ModulePass *llvm::createPGOInstrumentationUsePass(StringRef Filename) {
172   return new PGOInstrumentationUse(Filename.str());
173 }
174 
175 namespace {
176 /// \brief An MST based instrumentation for PGO
177 ///
178 /// Implements a Minimum Spanning Tree (MST) based instrumentation for PGO
179 /// in the function level.
180 struct PGOEdge {
181   // This class implements the CFG edges. Note the CFG can be a multi-graph.
182   // So there might be multiple edges with same SrcBB and DestBB.
183   const BasicBlock *SrcBB;
184   const BasicBlock *DestBB;
185   uint64_t Weight;
186   bool InMST;
187   bool Removed;
188   bool IsCritical;
189   PGOEdge(const BasicBlock *Src, const BasicBlock *Dest, unsigned W = 1)
190       : SrcBB(Src), DestBB(Dest), Weight(W), InMST(false), Removed(false),
191         IsCritical(false) {}
192   // Return the information string of an edge.
193   const std::string infoString() const {
194     return (Twine(Removed ? "-" : " ") + (InMST ? " " : "*") +
195             (IsCritical ? "c" : " ") + "  W=" + Twine(Weight)).str();
196   }
197 };
198 
199 // This class stores the auxiliary information for each BB.
200 struct BBInfo {
201   BBInfo *Group;
202   uint32_t Index;
203   uint32_t Rank;
204 
205   BBInfo(unsigned IX) : Group(this), Index(IX), Rank(0) {}
206 
207   // Return the information string of this object.
208   const std::string infoString() const {
209     return (Twine("Index=") + Twine(Index)).str();
210   }
211 };
212 
213 // This class implements the CFG edges. Note the CFG can be a multi-graph.
214 template <class Edge, class BBInfo> class FuncPGOInstrumentation {
215 private:
216   Function &F;
217   void computeCFGHash();
218 
219 public:
220   std::string FuncName;
221   GlobalVariable *FuncNameVar;
222   // CFG hash value for this function.
223   uint64_t FunctionHash;
224 
225   // The Minimum Spanning Tree of function CFG.
226   CFGMST<Edge, BBInfo> MST;
227 
228   // Give an edge, find the BB that will be instrumented.
229   // Return nullptr if there is no BB to be instrumented.
230   BasicBlock *getInstrBB(Edge *E);
231 
232   // Return the auxiliary BB information.
233   BBInfo &getBBInfo(const BasicBlock *BB) const { return MST.getBBInfo(BB); }
234 
235   // Dump edges and BB information.
236   void dumpInfo(std::string Str = "") const {
237     MST.dumpEdges(dbgs(), Twine("Dump Function ") + FuncName + " Hash: " +
238                               Twine(FunctionHash) + "\t" + Str);
239   }
240 
241   FuncPGOInstrumentation(Function &Func, bool CreateGlobalVar = false,
242                          BranchProbabilityInfo *BPI = nullptr,
243                          BlockFrequencyInfo *BFI = nullptr)
244       : F(Func), FunctionHash(0), MST(F, BPI, BFI) {
245     FuncName = getPGOFuncName(F);
246     computeCFGHash();
247     DEBUG(dumpInfo("after CFGMST"));
248 
249     NumOfPGOBB += MST.BBInfos.size();
250     for (auto &E : MST.AllEdges) {
251       if (E->Removed)
252         continue;
253       NumOfPGOEdge++;
254       if (!E->InMST)
255         NumOfPGOInstrument++;
256     }
257 
258     if (CreateGlobalVar)
259       FuncNameVar = createPGOFuncNameVar(F, FuncName);
260   }
261 };
262 
263 // Compute Hash value for the CFG: the lower 32 bits are CRC32 of the index
264 // value of each BB in the CFG. The higher 32 bits record the number of edges.
265 template <class Edge, class BBInfo>
266 void FuncPGOInstrumentation<Edge, BBInfo>::computeCFGHash() {
267   std::vector<char> Indexes;
268   JamCRC JC;
269   for (auto &BB : F) {
270     const TerminatorInst *TI = BB.getTerminator();
271     for (unsigned I = 0, E = TI->getNumSuccessors(); I != E; ++I) {
272       BasicBlock *Succ = TI->getSuccessor(I);
273       uint32_t Index = getBBInfo(Succ).Index;
274       for (int J = 0; J < 4; J++)
275         Indexes.push_back((char)(Index >> (J * 8)));
276     }
277   }
278   JC.update(Indexes);
279   FunctionHash = (uint64_t)MST.AllEdges.size() << 32 | JC.getCRC();
280 }
281 
282 // Given a CFG E to be instrumented, find which BB to place the instrumented
283 // code. The function will split the critical edge if necessary.
284 template <class Edge, class BBInfo>
285 BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
286   if (E->InMST || E->Removed)
287     return nullptr;
288 
289   BasicBlock *SrcBB = const_cast<BasicBlock *>(E->SrcBB);
290   BasicBlock *DestBB = const_cast<BasicBlock *>(E->DestBB);
291   // For a fake edge, instrument the real BB.
292   if (SrcBB == nullptr)
293     return DestBB;
294   if (DestBB == nullptr)
295     return SrcBB;
296 
297   // Instrument the SrcBB if it has a single successor,
298   // otherwise, the DestBB if this is not a critical edge.
299   TerminatorInst *TI = SrcBB->getTerminator();
300   if (TI->getNumSuccessors() <= 1)
301     return SrcBB;
302   if (!E->IsCritical)
303     return DestBB;
304 
305   // For a critical edge, we have to split. Instrument the newly
306   // created BB.
307   NumOfPGOSplit++;
308   DEBUG(dbgs() << "Split critical edge: " << getBBInfo(SrcBB).Index << " --> "
309                << getBBInfo(DestBB).Index << "\n");
310   unsigned SuccNum = GetSuccessorNumber(SrcBB, DestBB);
311   BasicBlock *InstrBB = SplitCriticalEdge(TI, SuccNum);
312   assert(InstrBB && "Critical edge is not split");
313 
314   E->Removed = true;
315   return InstrBB;
316 }
317 
318 // Visitor class that finds all indirect call sites.
319 struct PGOIndirectCallSiteVisitor
320     : public InstVisitor<PGOIndirectCallSiteVisitor> {
321   std::vector<CallInst *> IndirectCallInsts;
322   PGOIndirectCallSiteVisitor() {}
323 
324   void visitCallInst(CallInst &I) {
325     CallSite CS(&I);
326     if (CS.getCalledFunction() || !CS.getCalledValue())
327       return;
328     IndirectCallInsts.push_back(&I);
329   }
330 };
331 
332 // Visit all edge and instrument the edges not in MST, and do value profiling.
333 // Critical edges will be split.
334 static void instrumentOneFunc(Function &F, Module *M,
335                               BranchProbabilityInfo *BPI,
336                               BlockFrequencyInfo *BFI) {
337   unsigned NumCounters = 0;
338   FuncPGOInstrumentation<PGOEdge, BBInfo> FuncInfo(F, true, BPI, BFI);
339   for (auto &E : FuncInfo.MST.AllEdges) {
340     if (!E->InMST && !E->Removed)
341       NumCounters++;
342   }
343 
344   uint32_t I = 0;
345   Type *I8PtrTy = Type::getInt8PtrTy(M->getContext());
346   for (auto &E : FuncInfo.MST.AllEdges) {
347     BasicBlock *InstrBB = FuncInfo.getInstrBB(E.get());
348     if (!InstrBB)
349       continue;
350 
351     IRBuilder<> Builder(InstrBB, InstrBB->getFirstInsertionPt());
352     assert(Builder.GetInsertPoint() != InstrBB->end() &&
353            "Cannot get the Instrumentation point");
354     Builder.CreateCall(
355         Intrinsic::getDeclaration(M, Intrinsic::instrprof_increment),
356         {llvm::ConstantExpr::getBitCast(FuncInfo.FuncNameVar, I8PtrTy),
357          Builder.getInt64(FuncInfo.FunctionHash), Builder.getInt32(NumCounters),
358          Builder.getInt32(I++)});
359   }
360 
361   if (DisableValueProfiling)
362     return;
363 
364   unsigned NumIndirectCallSites = 0;
365   PGOIndirectCallSiteVisitor ICV;
366   ICV.visit(F);
367   for (auto &I : ICV.IndirectCallInsts) {
368     CallSite CS(I);
369     Value *Callee = CS.getCalledValue();
370     DEBUG(dbgs() << "Instrument one indirect call: CallSite Index = "
371                  << NumIndirectCallSites << "\n");
372     IRBuilder<> Builder(I);
373     assert(Builder.GetInsertPoint() != I->getParent()->end() &&
374            "Cannot get the Instrumentation point");
375     Builder.CreateCall(
376         Intrinsic::getDeclaration(M, Intrinsic::instrprof_value_profile),
377         {llvm::ConstantExpr::getBitCast(FuncInfo.FuncNameVar, I8PtrTy),
378          Builder.getInt64(FuncInfo.FunctionHash),
379          Builder.CreatePtrToInt(Callee, Builder.getInt64Ty()),
380          Builder.getInt32(llvm::InstrProfValueKind::IPVK_IndirectCallTarget),
381          Builder.getInt32(NumIndirectCallSites++)});
382   }
383   NumOfPGOICall += NumIndirectCallSites;
384 }
385 
386 // This class represents a CFG edge in profile use compilation.
387 struct PGOUseEdge : public PGOEdge {
388   bool CountValid;
389   uint64_t CountValue;
390   PGOUseEdge(const BasicBlock *Src, const BasicBlock *Dest, unsigned W = 1)
391       : PGOEdge(Src, Dest, W), CountValid(false), CountValue(0) {}
392 
393   // Set edge count value
394   void setEdgeCount(uint64_t Value) {
395     CountValue = Value;
396     CountValid = true;
397   }
398 
399   // Return the information string for this object.
400   const std::string infoString() const {
401     if (!CountValid)
402       return PGOEdge::infoString();
403     return (Twine(PGOEdge::infoString()) + "  Count=" + Twine(CountValue)).str();
404   }
405 };
406 
407 typedef SmallVector<PGOUseEdge *, 2> DirectEdges;
408 
409 // This class stores the auxiliary information for each BB.
410 struct UseBBInfo : public BBInfo {
411   uint64_t CountValue;
412   bool CountValid;
413   int32_t UnknownCountInEdge;
414   int32_t UnknownCountOutEdge;
415   DirectEdges InEdges;
416   DirectEdges OutEdges;
417   UseBBInfo(unsigned IX)
418       : BBInfo(IX), CountValue(0), CountValid(false), UnknownCountInEdge(0),
419         UnknownCountOutEdge(0) {}
420   UseBBInfo(unsigned IX, uint64_t C)
421       : BBInfo(IX), CountValue(C), CountValid(true), UnknownCountInEdge(0),
422         UnknownCountOutEdge(0) {}
423 
424   // Set the profile count value for this BB.
425   void setBBInfoCount(uint64_t Value) {
426     CountValue = Value;
427     CountValid = true;
428   }
429 
430   // Return the information string of this object.
431   const std::string infoString() const {
432     if (!CountValid)
433       return BBInfo::infoString();
434     return (Twine(BBInfo::infoString()) + "  Count=" + Twine(CountValue)).str();
435   }
436 };
437 
438 // Sum up the count values for all the edges.
439 static uint64_t sumEdgeCount(const ArrayRef<PGOUseEdge *> Edges) {
440   uint64_t Total = 0;
441   for (auto &E : Edges) {
442     if (E->Removed)
443       continue;
444     Total += E->CountValue;
445   }
446   return Total;
447 }
448 
449 class PGOUseFunc {
450 private:
451   Function &F;
452   Module *M;
453   // This member stores the shared information with class PGOGenFunc.
454   FuncPGOInstrumentation<PGOUseEdge, UseBBInfo> FuncInfo;
455 
456   // Return the auxiliary BB information.
457   UseBBInfo &getBBInfo(const BasicBlock *BB) const {
458     return FuncInfo.getBBInfo(BB);
459   }
460 
461   // The maximum count value in the profile. This is only used in PGO use
462   // compilation.
463   uint64_t ProgramMaxCount;
464 
465   // Find the Instrumented BB and set the value.
466   void setInstrumentedCounts(const std::vector<uint64_t> &CountFromProfile);
467 
468   // Set the edge counter value for the unknown edge -- there should be only
469   // one unknown edge.
470   void setEdgeCount(DirectEdges &Edges, uint64_t Value);
471 
472   // Return FuncName string;
473   const std::string getFuncName() const { return FuncInfo.FuncName; }
474 
475   // Set the hot/cold inline hints based on the count values.
476   // FIXME: This function should be removed once the functionality in
477   // the inliner is implemented.
478   void applyFunctionAttributes(uint64_t EntryCount, uint64_t MaxCount) {
479     if (ProgramMaxCount == 0)
480       return;
481     // Threshold of the hot functions.
482     const BranchProbability HotFunctionThreshold(1, 100);
483     // Threshold of the cold functions.
484     const BranchProbability ColdFunctionThreshold(2, 10000);
485     if (EntryCount >= HotFunctionThreshold.scale(ProgramMaxCount))
486       F.addFnAttr(llvm::Attribute::InlineHint);
487     else if (MaxCount <= ColdFunctionThreshold.scale(ProgramMaxCount))
488       F.addFnAttr(llvm::Attribute::Cold);
489   }
490 
491 public:
492   PGOUseFunc(Function &Func, Module *Modu, BranchProbabilityInfo *BPI = nullptr,
493              BlockFrequencyInfo *BFI = nullptr)
494       : F(Func), M(Modu), FuncInfo(Func, false, BPI, BFI) {}
495 
496   // Read counts for the instrumented BB from profile.
497   bool readCounters(IndexedInstrProfReader *PGOReader);
498 
499   // Populate the counts for all BBs.
500   void populateCounters();
501 
502   // Set the branch weights based on the count values.
503   void setBranchWeights();
504 };
505 
506 // Visit all the edges and assign the count value for the instrumented
507 // edges and the BB.
508 void PGOUseFunc::setInstrumentedCounts(
509     const std::vector<uint64_t> &CountFromProfile) {
510 
511   // Use a worklist as we will update the vector during the iteration.
512   std::vector<PGOUseEdge *> WorkList;
513   for (auto &E : FuncInfo.MST.AllEdges)
514     WorkList.push_back(E.get());
515 
516   uint32_t I = 0;
517   for (auto &E : WorkList) {
518     BasicBlock *InstrBB = FuncInfo.getInstrBB(E);
519     if (!InstrBB)
520       continue;
521     uint64_t CountValue = CountFromProfile[I++];
522     if (!E->Removed) {
523       getBBInfo(InstrBB).setBBInfoCount(CountValue);
524       E->setEdgeCount(CountValue);
525       continue;
526     }
527 
528     // Need to add two new edges.
529     BasicBlock *SrcBB = const_cast<BasicBlock *>(E->SrcBB);
530     BasicBlock *DestBB = const_cast<BasicBlock *>(E->DestBB);
531     // Add new edge of SrcBB->InstrBB.
532     PGOUseEdge &NewEdge = FuncInfo.MST.addEdge(SrcBB, InstrBB, 0);
533     NewEdge.setEdgeCount(CountValue);
534     // Add new edge of InstrBB->DestBB.
535     PGOUseEdge &NewEdge1 = FuncInfo.MST.addEdge(InstrBB, DestBB, 0);
536     NewEdge1.setEdgeCount(CountValue);
537     NewEdge1.InMST = true;
538     getBBInfo(InstrBB).setBBInfoCount(CountValue);
539   }
540 }
541 
542 // Set the count value for the unknown edge. There should be one and only one
543 // unknown edge in Edges vector.
544 void PGOUseFunc::setEdgeCount(DirectEdges &Edges, uint64_t Value) {
545   for (auto &E : Edges) {
546     if (E->CountValid)
547       continue;
548     E->setEdgeCount(Value);
549 
550     getBBInfo(E->SrcBB).UnknownCountOutEdge--;
551     getBBInfo(E->DestBB).UnknownCountInEdge--;
552     return;
553   }
554   llvm_unreachable("Cannot find the unknown count edge");
555 }
556 
557 // Read the profile from ProfileFileName and assign the value to the
558 // instrumented BB and the edges. This function also updates ProgramMaxCount.
559 // Return true if the profile are successfully read, and false on errors.
560 bool PGOUseFunc::readCounters(IndexedInstrProfReader *PGOReader) {
561   auto &Ctx = M->getContext();
562   ErrorOr<InstrProfRecord> Result =
563       PGOReader->getInstrProfRecord(FuncInfo.FuncName, FuncInfo.FunctionHash);
564   if (std::error_code EC = Result.getError()) {
565     if (EC == instrprof_error::unknown_function)
566       NumOfPGOMissing++;
567     else if (EC == instrprof_error::hash_mismatch ||
568              EC == llvm::instrprof_error::malformed)
569       NumOfPGOMismatch++;
570 
571     std::string Msg = EC.message() + std::string(" ") + F.getName().str();
572     Ctx.diagnose(
573         DiagnosticInfoPGOProfile(M->getName().data(), Msg, DS_Warning));
574     return false;
575   }
576   std::vector<uint64_t> &CountFromProfile = Result.get().Counts;
577 
578   NumOfPGOFunc++;
579   DEBUG(dbgs() << CountFromProfile.size() << " counts\n");
580   uint64_t ValueSum = 0;
581   for (unsigned I = 0, S = CountFromProfile.size(); I < S; I++) {
582     DEBUG(dbgs() << "  " << I << ": " << CountFromProfile[I] << "\n");
583     ValueSum += CountFromProfile[I];
584   }
585 
586   DEBUG(dbgs() << "SUM =  " << ValueSum << "\n");
587 
588   getBBInfo(nullptr).UnknownCountOutEdge = 2;
589   getBBInfo(nullptr).UnknownCountInEdge = 2;
590 
591   setInstrumentedCounts(CountFromProfile);
592   ProgramMaxCount = PGOReader->getMaximumFunctionCount();
593   return true;
594 }
595 
596 // Populate the counters from instrumented BBs to all BBs.
597 // In the end of this operation, all BBs should have a valid count value.
598 void PGOUseFunc::populateCounters() {
599   // First set up Count variable for all BBs.
600   for (auto &E : FuncInfo.MST.AllEdges) {
601     if (E->Removed)
602       continue;
603 
604     const BasicBlock *SrcBB = E->SrcBB;
605     const BasicBlock *DestBB = E->DestBB;
606     UseBBInfo &SrcInfo = getBBInfo(SrcBB);
607     UseBBInfo &DestInfo = getBBInfo(DestBB);
608     SrcInfo.OutEdges.push_back(E.get());
609     DestInfo.InEdges.push_back(E.get());
610     SrcInfo.UnknownCountOutEdge++;
611     DestInfo.UnknownCountInEdge++;
612 
613     if (!E->CountValid)
614       continue;
615     DestInfo.UnknownCountInEdge--;
616     SrcInfo.UnknownCountOutEdge--;
617   }
618 
619   bool Changes = true;
620   unsigned NumPasses = 0;
621   while (Changes) {
622     NumPasses++;
623     Changes = false;
624 
625     // For efficient traversal, it's better to start from the end as most
626     // of the instrumented edges are at the end.
627     for (auto &BB : reverse(F)) {
628       UseBBInfo &Count = getBBInfo(&BB);
629       if (!Count.CountValid) {
630         if (Count.UnknownCountOutEdge == 0) {
631           Count.CountValue = sumEdgeCount(Count.OutEdges);
632           Count.CountValid = true;
633           Changes = true;
634         } else if (Count.UnknownCountInEdge == 0) {
635           Count.CountValue = sumEdgeCount(Count.InEdges);
636           Count.CountValid = true;
637           Changes = true;
638         }
639       }
640       if (Count.CountValid) {
641         if (Count.UnknownCountOutEdge == 1) {
642           uint64_t Total = Count.CountValue - sumEdgeCount(Count.OutEdges);
643           setEdgeCount(Count.OutEdges, Total);
644           Changes = true;
645         }
646         if (Count.UnknownCountInEdge == 1) {
647           uint64_t Total = Count.CountValue - sumEdgeCount(Count.InEdges);
648           setEdgeCount(Count.InEdges, Total);
649           Changes = true;
650         }
651       }
652     }
653   }
654 
655   DEBUG(dbgs() << "Populate counts in " << NumPasses << " passes.\n");
656   // Assert every BB has a valid counter.
657   uint64_t FuncEntryCount = getBBInfo(&*F.begin()).CountValue;
658   uint64_t FuncMaxCount = FuncEntryCount;
659   for (auto &BB : F) {
660     assert(getBBInfo(&BB).CountValid && "BB count is not valid");
661     uint64_t Count = getBBInfo(&BB).CountValue;
662     if (Count > FuncMaxCount)
663       FuncMaxCount = Count;
664   }
665   applyFunctionAttributes(FuncEntryCount, FuncMaxCount);
666 
667   DEBUG(FuncInfo.dumpInfo("after reading profile."));
668 }
669 
670 // Assign the scaled count values to the BB with multiple out edges.
671 void PGOUseFunc::setBranchWeights() {
672   // Generate MD_prof metadata for every branch instruction.
673   DEBUG(dbgs() << "\nSetting branch weights.\n");
674   MDBuilder MDB(M->getContext());
675   for (auto &BB : F) {
676     TerminatorInst *TI = BB.getTerminator();
677     if (TI->getNumSuccessors() < 2)
678       continue;
679     if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI))
680       continue;
681     if (getBBInfo(&BB).CountValue == 0)
682       continue;
683 
684     // We have a non-zero Branch BB.
685     const UseBBInfo &BBCountInfo = getBBInfo(&BB);
686     unsigned Size = BBCountInfo.OutEdges.size();
687     SmallVector<unsigned, 2> EdgeCounts(Size, 0);
688     uint64_t MaxCount = 0;
689     for (unsigned s = 0; s < Size; s++) {
690       const PGOUseEdge *E = BBCountInfo.OutEdges[s];
691       const BasicBlock *SrcBB = E->SrcBB;
692       const BasicBlock *DestBB = E->DestBB;
693       if (DestBB == nullptr)
694         continue;
695       unsigned SuccNum = GetSuccessorNumber(SrcBB, DestBB);
696       uint64_t EdgeCount = E->CountValue;
697       if (EdgeCount > MaxCount)
698         MaxCount = EdgeCount;
699       EdgeCounts[SuccNum] = EdgeCount;
700     }
701     assert(MaxCount > 0 && "Bad max count");
702     uint64_t Scale = calculateCountScale(MaxCount);
703     SmallVector<unsigned, 4> Weights;
704     for (const auto &ECI : EdgeCounts)
705       Weights.push_back(scaleBranchCount(ECI, Scale));
706 
707     TI->setMetadata(llvm::LLVMContext::MD_prof,
708                     MDB.createBranchWeights(Weights));
709     DEBUG(dbgs() << "Weight is: ";
710           for (const auto &W : Weights) { dbgs() << W << " "; }
711           dbgs() << "\n";);
712   }
713 }
714 } // end anonymous namespace
715 
716 bool PGOInstrumentationGen::runOnModule(Module &M) {
717   for (auto &F : M) {
718     if (F.isDeclaration())
719       continue;
720     BranchProbabilityInfo *BPI =
721         &(getAnalysis<BranchProbabilityInfoWrapperPass>(F).getBPI());
722     BlockFrequencyInfo *BFI =
723         &(getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI());
724     instrumentOneFunc(F, &M, BPI, BFI);
725   }
726   return true;
727 }
728 
729 static void setPGOCountOnFunc(PGOUseFunc &Func,
730                               IndexedInstrProfReader *PGOReader) {
731   if (Func.readCounters(PGOReader)) {
732     Func.populateCounters();
733     Func.setBranchWeights();
734   }
735 }
736 
737 bool PGOInstrumentationUse::runOnModule(Module &M) {
738   DEBUG(dbgs() << "Read in profile counters: ");
739   auto &Ctx = M.getContext();
740   // Read the counter array from file.
741   auto ReaderOrErr = IndexedInstrProfReader::create(ProfileFileName);
742   if (std::error_code EC = ReaderOrErr.getError()) {
743     Ctx.diagnose(
744         DiagnosticInfoPGOProfile(ProfileFileName.data(), EC.message()));
745     return false;
746   }
747 
748   PGOReader = std::move(ReaderOrErr.get());
749   if (!PGOReader) {
750     Ctx.diagnose(DiagnosticInfoPGOProfile(ProfileFileName.data(),
751                                           "Cannot get PGOReader"));
752     return false;
753   }
754 
755   for (auto &F : M) {
756     if (F.isDeclaration())
757       continue;
758     BranchProbabilityInfo *BPI =
759         &(getAnalysis<BranchProbabilityInfoWrapperPass>(F).getBPI());
760     BlockFrequencyInfo *BFI =
761         &(getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI());
762     PGOUseFunc Func(F, &M, BPI, BFI);
763     setPGOCountOnFunc(Func, PGOReader.get());
764   }
765   return true;
766 }
767