Index: llvm/lib/Transforms/IPO/HotColdSplitting.cpp
===================================================================
--- llvm/lib/Transforms/IPO/HotColdSplitting.cpp
+++ llvm/lib/Transforms/IPO/HotColdSplitting.cpp
@@ -101,14 +101,19 @@
   return true;
 }
 
+// Same as blockEndsInUnreachable in CodeGen/BranchFolding.cpp. Do not modify
+// this function unless you modify the MBB version as well.
+//
+/// A no successor, non-return block probably ends in unreachable and is cold.
+/// Also consider a block that ends in an indirect branch to be a return block,
+/// since many targets use plain indirect branches to return.
 bool blockEndsInUnreachable(const BasicBlock &BB) {
+  if (!succ_empty(&BB))
+    return false;
   if (BB.empty())
     return true;
   const TerminatorInst *I = BB.getTerminator();
-  if (isa<ReturnInst>(I) || isa<IndirectBrInst>(I))
-    return true;
-  // Unreachable blocks do not have any successor.
-  return succ_empty(&BB);
+  return !(isa<ReturnInst>(I) || isa<IndirectBrInst>(I));
 }
 
 static bool exceptionHandlingFunctions(const CallInst *CI) {
@@ -123,8 +128,7 @@
          FName == "__cxa_end_catch";
 }
 
-static
-bool unlikelyExecuted(const BasicBlock &BB) {
+static bool unlikelyExecuted(const BasicBlock &BB) {
   if (blockEndsInUnreachable(BB))
     return true;
   // Exception handling blocks are unlikely executed.
@@ -145,13 +149,32 @@
   return false;
 }
 
+static bool returnsOrHasSideEffects(const BasicBlock &BB) {
+  const TerminatorInst *I = BB.getTerminator();
+  if (isa<ReturnInst>(I) || isa<IndirectBrInst>(I) || isa<InvokeInst>(I))
+    return true;
+
+  for (const Instruction &I : BB)
+    if (const CallInst *CI = dyn_cast<CallInst>(&I)) {
+      if (CI->hasFnAttr(Attribute::NoReturn))
+        return true;
+
+      if (isa<InlineAsm>(CI->getCalledValue()))
+        return true;
+    }
+
+  return false;
+}
+
 static DenseSetBB getHotBlocks(Function &F) {
 
   // Mark all cold basic blocks.
   DenseSetBB ColdBlocks;
   for (BasicBlock &BB : F)
-    if (unlikelyExecuted(BB))
+    if (unlikelyExecuted(BB)) {
+      LLVM_DEBUG(llvm::dbgs() << "\nForward propagation marks cold: " << BB);
       ColdBlocks.insert((const BasicBlock *)&BB);
+    }
 
   // Forward propagation: basic blocks are hot when they are reachable from the
   // beginning of the function through a path that does not contain cold blocks.
@@ -203,7 +226,12 @@
     if (ColdBlocks.count(It))
       continue;
 
+    // Do not back-propagate to blocks that return or have side effects.
+    if (returnsOrHasSideEffects(*It))
+      continue;
+
     // Move the block from HotBlocks to ColdBlocks.
+    LLVM_DEBUG(llvm::dbgs() << "\nBack propagation marks cold: " << *It);
     HotBlocks.erase(It);
     ColdBlocks.insert(It);
 
@@ -353,6 +381,12 @@
   // Walking the dominator tree allows us to find the largest
   // cold region.
   BasicBlock *Begin = DT->getRootNode()->getBlock();
+
+  // Early return if the beginning of the function has been marked cold,
+  // otherwise all the function gets outlined.
+  if (PSI->isColdBB(Begin, BFI) || !HotBlocks.count(Begin))
+    return nullptr;
+
   for (auto I = df_begin(Begin), E = df_end(Begin); I != E; ++I) {
     BasicBlock *BB = *I;
     if (PSI->isColdBB(BB, BFI) || !HotBlocks.count(BB)) {
Index: llvm/test/Transforms/HotColdSplit/split-cold-1.ll
===================================================================
--- llvm/test/Transforms/HotColdSplit/split-cold-1.ll
+++ llvm/test/Transforms/HotColdSplit/split-cold-1.ll
@@ -1,9 +1,11 @@
 ; RUN: opt -hotcoldsplit -S < %s | FileCheck %s
 ; RUN: opt -passes=hotcoldsplit -S < %s | FileCheck %s
 
-; Outlined function is called from a basic block named codeRepl
-; CHECK: codeRepl:
-; CHECK-NEXT: call void @foo
+; Check that the function is not split. Outlined function is called from a
+; basic block named codeRepl.
+
+; CHECK-LABEL: @foo
+; CHECK-NOT: codeRepl
 define void @foo() {
 entry:
   br i1 undef, label %if.then, label %if.end
@@ -23,3 +25,19 @@
 return:                                           ; preds = %cleanup40
   ret void
 }
+
+; Check that the function is not split. We used to outline the full function.
+
+; CHECK-LABEL: @fun
+; CHECK-NOT: codeRepl
+
+define void @fun() {
+entry:
+  br i1 undef, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  br label %if.end
+
+if.end:                                           ; preds = %entry
+  ret void
+}
Index: llvm/test/Transforms/HotColdSplit/split-cold-2.ll
===================================================================
--- llvm/test/Transforms/HotColdSplit/split-cold-2.ll
+++ llvm/test/Transforms/HotColdSplit/split-cold-2.ll
@@ -4,6 +4,10 @@
 ; Make sure this compiles. This test used to fail with an invalid phi node: the
 ; two predecessors were outlined and the SSA representation was invalid.
 
+; CHECK-LABEL: @fun
+; CHECK: codeRepl:
+; CHECK-NEXT: call void @fun_if.else
+
 define void @fun() {
 entry:
   br i1 undef, label %if.then, label %if.else