Index: lib/Transforms/Instrumentation/PGOInstrumentation.cpp
===================================================================
--- lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -459,6 +459,35 @@
 }
 
 class PGOUseFunc {
+public:
+  PGOUseFunc(Function &Func, Module *Modu, BranchProbabilityInfo *BPI = nullptr,
+             BlockFrequencyInfo *BFI = nullptr)
+      : F(Func), M(Modu), FuncInfo(Func, false, BPI, BFI) {}
+
+  // Read counts for the instrumented BB from profile.
+  bool readCounters(IndexedInstrProfReader *PGOReader);
+
+  // Populate the counts for all BBs.
+  void populateCounters();
+
+  // Set the branch weights based on the count values.
+  void setBranchWeights();
+
+  // Annotate the indirect call sites.
+  void annotateIndirectCallSites();
+
+  // The hotness of the function from the profile count.
+  enum FuncFreqAttr {
+    FunctionFrequencyNormal,
+    FunctionFrequencyCold,
+    FunctionFrequencyHot
+  };
+
+  // Return the funtion hotness from the profile.
+  FuncFreqAttr getFuncFreqAttr() const {
+    return FreqAttr;
+  }
+
 private:
   Function &F;
   Module *M;
@@ -487,10 +516,13 @@
   // Return FuncName string;
   const std::string getFuncName() const { return FuncInfo.FuncName; }
 
+  FuncFreqAttr FreqAttr;
+
   // Set the hot/cold inline hints based on the count values.
   // FIXME: This function should be removed once the functionality in
   // the inliner is implemented.
-  void applyFunctionAttributes(uint64_t EntryCount, uint64_t MaxCount) {
+  void markFunctionAttributes(uint64_t EntryCount, uint64_t MaxCount) {
+    FreqAttr = FunctionFrequencyNormal;
     if (ProgramMaxCount == 0)
       return;
     // Threshold of the hot functions.
@@ -498,27 +530,10 @@
     // Threshold of the cold functions.
     const BranchProbability ColdFunctionThreshold(2, 10000);
     if (EntryCount >= HotFunctionThreshold.scale(ProgramMaxCount))
-      F.addFnAttr(llvm::Attribute::InlineHint);
+      FreqAttr = FunctionFrequencyHot;
     else if (MaxCount <= ColdFunctionThreshold.scale(ProgramMaxCount))
-      F.addFnAttr(llvm::Attribute::Cold);
+      FreqAttr = FunctionFrequencyCold;
   }
-
-public:
-  PGOUseFunc(Function &Func, Module *Modu, BranchProbabilityInfo *BPI = nullptr,
-             BlockFrequencyInfo *BFI = nullptr)
-      : F(Func), M(Modu), FuncInfo(Func, false, BPI, BFI) {}
-
-  // Read counts for the instrumented BB from profile.
-  bool readCounters(IndexedInstrProfReader *PGOReader);
-
-  // Populate the counts for all BBs.
-  void populateCounters();
-
-  // Set the branch weights based on the count values.
-  void setBranchWeights();
-
-  // Annotate the indirect call sites.
-  void annotateIndirectCallSites();
 };
 
 // Visit all the edges and assign the count value for the instrumented
@@ -681,7 +696,7 @@
     if (Count > FuncMaxCount)
       FuncMaxCount = Count;
   }
-  applyFunctionAttributes(FuncEntryCount, FuncMaxCount);
+  markFunctionAttributes(FuncEntryCount, FuncMaxCount);
 
   DEBUG(FuncInfo.dumpInfo("after reading profile."));
 }
@@ -827,6 +842,8 @@
     return false;
   }
 
+  std::vector<Function *> HotFunctions;
+  std::vector<Function *> ColdFunctions;
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
@@ -836,6 +853,23 @@
         &(getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI());
     PGOUseFunc Func(F, &M, BPI, BFI);
     setPGOCountOnFunc(Func, PGOReader.get());
+    PGOUseFunc::FuncFreqAttr FreqAttr = Func.getFuncFreqAttr();
+    if (FreqAttr == PGOUseFunc::FunctionFrequencyCold)
+      ColdFunctions.push_back(&F);
+    else if (FreqAttr == PGOUseFunc::FunctionFrequencyHot)
+      HotFunctions.push_back(&F);
   }
+
+  // Set function hotness attribute from the profile.
+  for (auto &F : HotFunctions) {
+    F->addFnAttr(llvm::Attribute::InlineHint);
+    DEBUG(dbgs() << "Set inline attribute to function: " << F->getName()
+                 << "\n");
+  }
+  for (auto &F : ColdFunctions) {
+    F->addFnAttr(llvm::Attribute::Cold);
+    DEBUG(dbgs() << "Set cold attribute to function: " << F->getName() << "\n");
+  }
+
   return true;
 }