1 files changed, 39 insertions, 20 deletions

diff --git a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
index d6ac7183f77..49b59504738 100644
--- a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
+++ b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
@@ -153,15 +153,19 @@ static bool isProfitableToOutline(const BlockSequence &Region,
   return false;
 }
 
-/// Mark \p F cold. Return true if it's changed.
-static bool markEntireFunctionCold(Function &F) {
+/// Mark \p F cold. Based on this assumption, also optimize it for minimum size.
+/// Return true if the function is changed.
+static bool markFunctionCold(Function &F) {
   assert(!F.hasFnAttribute(Attribute::OptimizeNone) && "Can't mark this cold");
   bool Changed = false;
+  if (!F.hasFnAttribute(Attribute::Cold)) {
+    F.addFnAttr(Attribute::Cold);
+    Changed = true;
+  }
   if (!F.hasFnAttribute(Attribute::MinSize)) {
     F.addFnAttr(Attribute::MinSize);
     Changed = true;
   }
-  // TODO: Move this function into a cold section.
   return Changed;
 }
 
@@ -175,6 +179,7 @@ public:
   bool run(Module &M);
 
 private:
+  bool isFunctionCold(const Function &F) const;
   bool shouldOutlineFrom(const Function &F) const;
   bool outlineColdRegions(Function &F, ProfileSummaryInfo &PSI,
                           BlockFrequencyInfo *BFI, TargetTransformInfo &TTI,
@@ -208,28 +213,29 @@ public:
 
 } // end anonymous namespace
 
-// Returns false if the function should not be considered for hot-cold split
-// optimization.
-bool HotColdSplitting::shouldOutlineFrom(const Function &F) const {
-  if (F.size() <= 2)
-    return false;
+/// Check whether \p F is inherently cold.
+bool HotColdSplitting::isFunctionCold(const Function &F) const {
+  if (F.hasFnAttribute(Attribute::Cold))
+    return true;
+
+  if (F.getCallingConv() == CallingConv::Cold)
+    return true;
 
-  // TODO: Consider only skipping functions marked `optnone` or `cold`.
+  if (PSI->isFunctionEntryCold(&F))
+    return true;
 
-  if (F.hasAddressTaken())
-    return false;
+  return false;
+}
 
+// Returns false if the function should not be considered for hot-cold split
+// optimization.
+bool HotColdSplitting::shouldOutlineFrom(const Function &F) const {
   if (F.hasFnAttribute(Attribute::AlwaysInline))
     return false;
 
   if (F.hasFnAttribute(Attribute::NoInline))
     return false;
 
-  if (F.getCallingConv() == CallingConv::Cold)
-    return false;
-
-  if (PSI->isFunctionEntryCold(&F))
-    return false;
   return true;
 }
 
@@ -259,9 +265,7 @@ Function *HotColdSplitting::extractColdRegion(const BlockSequence &Region,
     }
     CI->setIsNoInline();
 
-    // Try to make the outlined code as small as possible on the assumption
-    // that it's cold.
-    markEntireFunctionCold(*OutF);
+    markFunctionCold(*OutF);
 
     LLVM_DEBUG(llvm::dbgs() << "Outlined Region: " << *OutF);
     ORE.emit([&]() {
@@ -492,7 +496,7 @@ bool HotColdSplitting::outlineColdRegions(Function &F, ProfileSummaryInfo &PSI,
 
     if (Region.isEntireFunctionCold()) {
       LLVM_DEBUG(dbgs() << "Entire function is cold\n");
-      return markEntireFunctionCold(F);
+      return markFunctionCold(F);
     }
 
     // If this outlining region intersects with another, drop the new region.
@@ -548,10 +552,25 @@ bool HotColdSplitting::run(Module &M) {
   for (auto It = M.begin(), End = M.end(); It != End; ++It) {
     Function &F = *It;
 
+    // Do not touch declarations.
+    if (F.isDeclaration())
+      continue;
+
+    // Do not modify `optnone` functions.
+    if (F.hasFnAttribute(Attribute::OptimizeNone))
+      continue;
+
+    // Detect inherently cold functions and mark them as such.
+    if (isFunctionCold(F)) {
+      Changed |= markFunctionCold(F);
+      continue;
+    }
+
     if (!shouldOutlineFrom(F)) {
       LLVM_DEBUG(llvm::dbgs() << "Skipping " << F.getName() << "\n");
       continue;
     }
+
     LLVM_DEBUG(llvm::dbgs() << "Outlining in " << F.getName() << "\n");
     DominatorTree DT(F);
     PostDomTree PDT(F);