[PGO] Fix select instruction annotation

[xur]

Summary

Select instruction annotation in IR PGO uses the edge count to infer the
branch count. It's currently placed in setInstrumentedCounts() where
no all the BB counts have been computed. This leads to wrong branch weights.
Move the annotation after all BB counts are populated.

[davidxl]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

685

that being read out --> that remains to be read.

686

initialize new members.

[xur]

Integrated David's review comments.

[davidxl]

lgtm

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

private:
   Module *M;
   // This member stores the shared information with class PGOGenFunc.
   FuncPGOInstrumentation<PGOUseEdge, UseBBInfo> FuncInfo;
 
   // The maximum count value in the profile. This is only used in PGO use
   // compilation.
   uint64_t ProgramMaxCount;
 
+  // Position of counter that being read out.
+  uint32_t CountPosition;
+
+  // Total size of the profile count for this function.
+  uint32_t ProfileCountSize;
+
   // 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);
for (auto &E : WorkList) {
     PGOUseEdge &NewEdge = FuncInfo.MST.addEdge(SrcBB, InstrBB, 0);
     NewEdge.setEdgeCount(CountValue);
     // Add new edge of InstrBB->DestBB.
     PGOUseEdge &NewEdge1 = FuncInfo.MST.addEdge(InstrBB, DestBB, 0);
     NewEdge1.setEdgeCount(CountValue);
     NewEdge1.InMST = true;
     getBBInfo(InstrBB).setBBInfoCount(CountValue);
   }
-  // Now annotate select instructions
-  FuncInfo.SIVisitor.annotateSelects(F, this, &I);
-  assert(I == CountFromProfile.size());
+  ProfileCountSize =  CountFromProfile.size();
+  CountPosition = I;
 }
 
 // Set the count value for the unknown edge. There should be one and only one
 // unknown edge in Edges vector.
 void PGOUseFunc::setEdgeCount(DirectEdges &Edges, uint64_t Value) {
   for (auto &E : Edges) {
     if (E->CountValid)
       continue;

 #endif
   uint64_t FuncEntryCount = getBBInfo(&*F.begin()).CountValue;
   F.setEntryCount(FuncEntryCount);
   uint64_t FuncMaxCount = FuncEntryCount;
   for (auto &BB : F)
     FuncMaxCount = std::max(FuncMaxCount, getBBInfo(&BB).CountValue);
   markFunctionAttributes(FuncEntryCount, FuncMaxCount);
 
+  // Now annotate select instructions
+  FuncInfo.SIVisitor.annotateSelects(F, this, &CountPosition);
+  assert(CountPosition == ProfileCountSize);
+
   DEBUG(FuncInfo.dumpInfo("after reading profile."));
 }
 
 static void setProfMetadata(Module *M, Instruction *TI,
                             ArrayRef<uint64_t> EdgeCounts, uint64_t MaxCount) {
   MDBuilder MDB(M->getContext());
   assert(MaxCount > 0 && "Bad max count");
   uint64_t Scale = calculateCountScale(MaxCount);

test/Transforms/PGOProfile/Inputs/select2.proftext

+# IR level Instrumentation Flag
+:ir
+foo
+# Func Hash:
+72057628175588252
+# Num Counters:
+3
+# Counter Values:
+800
+3
+300

test/Transforms/PGOProfile/select2.ll

+; RUN: llvm-profdata merge %S/Inputs/select2.proftext -o %t.profdata
+; RUN: opt < %s -pgo-instr-use -pgo-test-profile-file=%t.profdata -pgo-instr-select=true -S | FileCheck %s --check-prefix=USE
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -pgo-instr-select=true -S | FileCheck %s --check-prefix=USE
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @foo(i32 %n) {
+;USE: define i32 @foo(i32 %n) !prof ![[ENTRY_COUNT:[0-9]+]] {
+entry:
+  br label %for.cond
+
+for.cond:
+  %i.0 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
+  %sum.0 = phi i32 [ 0, %entry ], [ %add, %for.inc ]
+  %cmp = icmp slt i32 %i.0, %n
+  br i1 %cmp, label %for.body, label %for.end
+;USE:  br i1 %cmp, label %for.body, label %for.end, !prof ![[BW_FOR_BR:[0-9]+]]
+
+for.body:
+  %cmp1 = icmp sgt i32 %sum.0, 10
+  %cond = select i1 %cmp1, i32 20, i32 -10
+;USE:  %cond = select i1 %cmp1, i32 20, i32 -10, !prof ![[BW_FOR_SELECT:[0-9]+]]
+  %add = add nsw i32 %sum.0, %cond
+  br label %for.inc
+
+for.inc:
+  %inc = add nsw i32 %i.0, 1
+  br label %for.cond
+
+for.end:
+  ret i32 %sum.0
+}
+
+;USE: ![[ENTRY_COUNT]] = !{!"function_entry_count", i64 3}
+;USE: ![[BW_FOR_BR]] = !{!"branch_weights", i32 800, i32 3}
+;USE: ![[BW_FOR_SELECT]] = !{!"branch_weights", i32 300, i32 500}