2 files changed, 22 insertions, 1 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() {