From 33308f92ebef0a0848997f912659c50ddf647383 Mon Sep 17 00:00:00 2001
From: Rong Xu <xur@google.com>
Date: Tue, 25 Oct 2016 21:47:24 +0000
Subject: [PATCH] [PGO] Fix select instruction annotation

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.

Reviewers: davidxl

Subscribers: llvm-commits

Differential Revision: https://reviews.llvm.org/D25961

llvm-svn: 285128
---
 .../Instrumentation/PGOInstrumentation.cpp    | 17 +++++++--
 .../PGOProfile/Inputs/select2.proftext        | 11 ++++++
 llvm/test/Transforms/PGOProfile/select2.ll    | 37 +++++++++++++++++++
 3 files changed, 61 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Transforms/PGOProfile/Inputs/select2.proftext
 create mode 100644 llvm/test/Transforms/PGOProfile/select2.ll

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index c2e421776610..833a8d8536f1 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -642,7 +642,7 @@ public:
              BranchProbabilityInfo *BPI = nullptr,
              BlockFrequencyInfo *BFI = nullptr)
       : F(Func), M(Modu), FuncInfo(Func, ComdatMembers, false, BPI, BFI),
-        FreqAttr(FFA_Normal) {}
+        CountPosition(0), ProfileCountSize(0), FreqAttr(FFA_Normal) {}
 
   // Read counts for the instrumented BB from profile.
   bool readCounters(IndexedInstrProfReader *PGOReader);
@@ -682,6 +682,12 @@ private:
   // compilation.
   uint64_t ProgramMaxCount;
 
+  // Position of counter that remains to be read.
+  uint32_t CountPosition;
+
+  // Total size of the profile count for this function.
+  uint32_t ProfileCountSize;
+
   // ProfileRecord for this function.
   InstrProfRecord ProfileRecord;
 
@@ -750,9 +756,8 @@ void PGOUseFunc::setInstrumentedCounts(
     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
@@ -892,6 +897,10 @@ void PGOUseFunc::populateCounters() {
     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."));
 }
 
diff --git a/llvm/test/Transforms/PGOProfile/Inputs/select2.proftext b/llvm/test/Transforms/PGOProfile/Inputs/select2.proftext
new file mode 100644
index 000000000000..56d78387ec2e
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/Inputs/select2.proftext
@@ -0,0 +1,11 @@
+# IR level Instrumentation Flag
+:ir
+foo
+# Func Hash:
+72057628175588252
+# Num Counters:
+3
+# Counter Values:
+800
+3
+300
diff --git a/llvm/test/Transforms/PGOProfile/select2.ll b/llvm/test/Transforms/PGOProfile/select2.ll
new file mode 100644
index 000000000000..b19fa1ff798f
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/select2.ll
@@ -0,0 +1,37 @@
+; 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}