[PGO] Add simplified branch weights for Objective-C for-collection loops.

Conceptually one of these loops is just a while-loop, but the actual code-gen
is more complicated. We don't instrument all the different control flow edges
to get accurate counts for each conditional branch, nor do I think it makes
sense to do so. Instead, make the simplifying assumption that the loop behaves
like a while-loop. Use the same branch weights for the first check for an
empty collection as would be used for the back-edge of a while loop, and use
that same weighting for the innermost loop, ignoring the possibility that there
may be some extra code to go fetch more elements.

llvm-svn: 204767

2 files changed, 15 insertions, 9 deletions

diff --git a/clang/lib/CodeGen/CGObjC.cpp b/clang/lib/CodeGen/CGObjC.cpp
index 02327c51b8f..f78bb0b3106 100644
--- a/clang/lib/CodeGen/CGObjC.cpp
+++ b/clang/lib/CodeGen/CGObjC.cpp
@@ -1509,9 +1509,13 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S){
   llvm::Value *zero = llvm::Constant::getNullValue(UnsignedLongLTy);
 
   // If the limit pointer was zero to begin with, the collection is
-  // empty; skip all this.
+  // empty; skip all this. Set the branch weight assuming this has the same
+  // probability of exiting the loop as any other loop exit.
+  uint64_t EntryCount = PGO.getCurrentRegionCount();
+  RegionCounter Cnt = getPGORegionCounter(&S);
   Builder.CreateCondBr(Builder.CreateICmpEQ(initialBufferLimit, zero, "iszero"),
-                       EmptyBB, LoopInitBB);
+                       EmptyBB, LoopInitBB,
+                       PGO.createBranchWeights(EntryCount, Cnt.getCount()));
 
   // Otherwise, initialize the loop.
   EmitBlock(LoopInitBB);
@@ -1540,7 +1544,6 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S){
   llvm::PHINode *count = Builder.CreatePHI(UnsignedLongLTy, 3, "forcoll.count");
   count->addIncoming(initialBufferLimit, LoopInitBB);
 
-  RegionCounter Cnt = getPGORegionCounter(&S);
   Cnt.beginRegion(Builder);
 
   // Check whether the mutations value has changed from where it was
@@ -1649,10 +1652,13 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S){
   llvm::Value *indexPlusOne
     = Builder.CreateAdd(index, llvm::ConstantInt::get(UnsignedLongLTy, 1));
 
-  // TODO: We should probably model this as a "continue" for PGO
   // If we haven't overrun the buffer yet, we can continue.
+  // Set the branch weights based on the simplifying assumption that this is
+  // like a while-loop, i.e., ignoring that the false branch fetches more
+  // elements and then returns to the loop.
   Builder.CreateCondBr(Builder.CreateICmpULT(indexPlusOne, count),
-                       LoopBodyBB, FetchMoreBB);
+                       LoopBodyBB, FetchMoreBB,
+                       PGO.createBranchWeights(Cnt.getCount(), EntryCount));
 
   index->addIncoming(indexPlusOne, AfterBody.getBlock());
   count->addIncoming(count, AfterBody.getBlock());
@@ -1673,8 +1679,6 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S){
   index->addIncoming(zero, Builder.GetInsertBlock());
   count->addIncoming(refetchCount, Builder.GetInsertBlock());
 
-  // TODO: We should be applying PGO weights here, but this needs to handle the
-  // branch before FetchMoreBB or we risk getting the numbers wrong.
   Builder.CreateCondBr(Builder.CreateICmpEQ(refetchCount, zero),
                        EmptyBB, LoopBodyBB);
 
@@ -1697,7 +1701,6 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S){
     PopCleanupBlock();
 
   EmitBlock(LoopEnd.getBlock());
-  // TODO: Once we calculate PGO weights above, set the region count here
 }
 
 void CodeGenFunction::EmitObjCAtTryStmt(const ObjCAtTryStmt &S) {
diff --git a/clang/test/Profile/objc-general.m b/clang/test/Profile/objc-general.m
index fe792bb2025..568e5e4182d 100644
--- a/clang/test/Profile/objc-general.m
+++ b/clang/test/Profile/objc-general.m
@@ -45,7 +45,8 @@ struct NSFastEnumerationState;
 {
   __block id result;
   // PGOGEN: store {{.*}} @[[FRC]], i64 0, i64 1
-  // FIXME: We don't emit branch weights for this yet.
+  // PGOUSE: br {{.*}} !prof ![[FR1:[0-9]+]]
+  // PGOUSE: br {{.*}} !prof ![[FR2:[0-9]+]]
   for (id x in array) {
     // PGOGEN: define {{.*}}_block_invoke
     // PGOUSE: define {{.*}}_block_invoke
@@ -60,6 +61,8 @@ struct NSFastEnumerationState;
 }
 @end
 
+// PGOUSE-DAG: ![[FR1]] = metadata !{metadata !"branch_weights", i32 2, i32 3}
+// PGOUSE-DAG: ![[FR2]] = metadata !{metadata !"branch_weights", i32 3, i32 2}
 // PGOUSE-DAG: ![[BL1]] = metadata !{metadata !"branch_weights", i32 2, i32 2}
 
 int main(int argc, const char *argv[]) {