4 files changed, 45 insertions, 6 deletions
diff --git a/clang/lib/CodeGen/CGBlocks.cpp b/clang/lib/CodeGen/CGBlocks.cpp
index b1717b38901..a75f341b3bf 100644
--- a/clang/lib/CodeGen/CGBlocks.cpp
+++ b/clang/lib/CodeGen/CGBlocks.cpp
@@ -1200,8 +1200,14 @@ CodeGenFunction::GenerateBlockFunction(GlobalDecl GD,
 
   if (IsLambdaConversionToBlock)
     EmitLambdaBlockInvokeBody();
-  else
+  else {
+    PGO.assignRegionCounters(blockDecl, fn);
+    RegionCounter Cnt = getPGORegionCounter(blockDecl->getBody());
+    Cnt.beginRegion(Builder);
     EmitStmt(blockDecl->getBody());
+    PGO.emitWriteoutFunction();
+    PGO.destroyRegionCounters();
+  }
 
   // Remember where we were...
   llvm::BasicBlock *resume = Builder.GetInsertBlock();
diff --git a/clang/lib/CodeGen/CodeGenPGO.cpp b/clang/lib/CodeGen/CodeGenPGO.cpp
index 349160ad352..cdd411cfa35 100644
--- a/clang/lib/CodeGen/CodeGenPGO.cpp
+++ b/clang/lib/CodeGen/CodeGenPGO.cpp
@@ -315,6 +315,10 @@ namespace {
       (*CounterMap)[S->getBody()] = NextCounter++;
       Visit(S->getBody());
     }
+    void VisitBlockDecl(const BlockDecl *S) {
+      (*CounterMap)[S->getBody()] = NextCounter++;
+      Visit(S->getBody());
+    }
     /// Assign a counter to track the block following a label.
     void VisitLabelStmt(const LabelStmt *S) {
       (*CounterMap)[S] = NextCounter++;
@@ -481,6 +485,13 @@ namespace {
       Visit(S->getBody());
     }
 
+    void VisitBlockDecl(const BlockDecl *S) {
+      RegionCounter Cnt(PGO, S->getBody());
+      Cnt.beginRegion();
+      (*CountMap)[S->getBody()] = PGO.getCurrentRegionCount();
+      Visit(S->getBody());
+    }
+
     void VisitReturnStmt(const ReturnStmt *S) {
       RecordStmtCount(S);
       if (S->getRetValue())
@@ -802,6 +813,8 @@ void CodeGenPGO::mapRegionCounters(const Decl *D) {
     Walker.VisitFunctionDecl(FD);
   else if (const ObjCMethodDecl *MD = dyn_cast_or_null<ObjCMethodDecl>(D))
     Walker.VisitObjCMethodDecl(MD);
+  else if (const BlockDecl *BD = dyn_cast_or_null<BlockDecl>(D))
+    Walker.VisitBlockDecl(BD);
   NumRegionCounters = Walker.NextCounter;
 }
 
@@ -812,6 +825,8 @@ void CodeGenPGO::computeRegionCounts(const Decl *D) {
     Walker.VisitFunctionDecl(FD);
   else if (const ObjCMethodDecl *MD = dyn_cast_or_null<ObjCMethodDecl>(D))
     Walker.VisitObjCMethodDecl(MD);
+  else if (const BlockDecl *BD = dyn_cast_or_null<BlockDecl>(D))
+    Walker.VisitBlockDecl(BD);
 }
 
 void CodeGenPGO::emitCounterVariables() {
diff --git a/clang/test/CodeGenObjC/Inputs/instr-profile.profdata b/clang/test/CodeGenObjC/Inputs/instr-profile.profdata
index f068640dbca..0fe046e5795 100644
--- a/clang/test/CodeGenObjC/Inputs/instr-profile.profdata
+++ b/clang/test/CodeGenObjC/Inputs/instr-profile.profdata
@@ -1,3 +1,7 @@
+instr-profile.m:__13+[A foreach:]_block_invoke 2
+2
+1
+
 instr-profile.m:+[A foreach:] 2
 1
 2
diff --git a/clang/test/CodeGenObjC/instr-profile.m b/clang/test/CodeGenObjC/instr-profile.m
index a7a0880eb60..c1bb50574de 100644
--- a/clang/test/CodeGenObjC/instr-profile.m
+++ b/clang/test/CodeGenObjC/instr-profile.m
@@ -4,8 +4,8 @@
 // of running the program generated by the -fprofile-instr-generate case. As
 // such, main() should call every function in this test.
 
-// RUN: %clang_cc1 -triple x86_64-apple-macosx10.9 -main-file-name instr-profile.m %s -o - -emit-llvm -fprofile-instr-generate | FileCheck -check-prefix=PGOGEN %s
-// RUN: %clang_cc1 -triple x86_64-apple-macosx10.9 -main-file-name instr-profile.m %s -o - -emit-llvm -fprofile-instr-use=%S/Inputs/instr-profile.profdata | FileCheck -check-prefix=PGOUSE %s
+// RUN: %clang_cc1 -triple x86_64-apple-macosx10.9 -main-file-name instr-profile.m %s -o - -emit-llvm -fblocks -fprofile-instr-generate | FileCheck -check-prefix=PGOGEN %s
+// RUN: %clang_cc1 -triple x86_64-apple-macosx10.9 -main-file-name instr-profile.m %s -o - -emit-llvm -fblocks -fprofile-instr-use=%S/Inputs/instr-profile.profdata | FileCheck -check-prefix=PGOUSE %s
 
 #ifdef HAVE_FOUNDATION
 
@@ -33,7 +33,9 @@ struct NSFastEnumerationState;
 @end;
 #endif
 
-// PGOGEN: @[[FOR:__llvm_pgo_ctr[0-9]*]] = private global [2 x i64] zeroinitializer
+// PGOGEN: @[[FRC:__llvm_pgo_ctr[0-9]*]] = private global [2 x i64] zeroinitializer
+// PGOGEN: @[[BLC:__llvm_pgo_ctr[0-9]*]] = private global [2 x i64] zeroinitializer
+// PGOGEN: @[[MAC:__llvm_pgo_ctr[0-9]*]] = private global [1 x i64] zeroinitializer
 
 @interface A : NSObject
 + (void)foreach: (NSArray *)array;
@@ -42,16 +44,28 @@ struct NSFastEnumerationState;
 @implementation A
 // PGOGEN: define {{.*}}+[A foreach:]
 // PGOUSE: define {{.*}}+[A foreach:]
-// PGOGEN: store {{.*}} @[[FOR]], i64 0, i64 0
+// PGOGEN: store {{.*}} @[[FRC]], i64 0, i64 0
 + (void)foreach: (NSArray *)array
 {
-  // PGOGEN: store {{.*}} @[[FOR]], i64 0, i64 1
+  __block id result;
+  // PGOGEN: store {{.*}} @[[FRC]], i64 0, i64 1
   // FIXME: We don't emit branch weights for this yet.
   for (id x in array) {
+    // PGOGEN: define {{.*}}_block_invoke
+    // PGOUSE: define {{.*}}_block_invoke
+    // PGOGEN: store {{.*}} @[[BLC]], i64 0, i64 0
+    ^{ static int init = 0;
+      // PGOGEN: store {{.*}} @[[BLC]], i64 0, i64 1
+      // PGOUSE: br {{.*}} !prof ![[BL1:[0-9]+]]
+       if (init)
+         result = x;
+       init = 1; }();
   }
 }
 @end
 
+// PGOUSE-DAG: ![[BL1]] = metadata !{metadata !"branch_weights", i32 2, i32 2}
+
 int main(int argc, const char *argv[]) {
   A *a = [[A alloc] init];
   NSArray *array = [NSArray arrayWithObjects: @"0", @"1", (void*)0];