1 files changed, 21 insertions, 7 deletions
diff --git a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
index 483809a8ed96..0e0eb8b8e00f 100644
--- a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
+++ b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
@@ -23,6 +23,7 @@
 // https://groups.google.com/d/msg/llvm-dev/RUegaMg-iqc/wFAVxa6fCgAJ
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/BasicBlockSectionUtils.h"
@@ -77,7 +78,7 @@ public:
 };
 } // end anonymous namespace
 
-static bool isColdBlock(MachineBasicBlock &MBB,
+static bool isColdBlock(const MachineBasicBlock &MBB,
                         const MachineBlockFrequencyInfo *MBFI,
                         ProfileSummaryInfo *PSI) {
   Optional<uint64_t> Count = MBFI->getBlockProfileCount(&MBB);
@@ -100,7 +101,8 @@ bool MachineFunctionSplitter::runOnMachineFunction(MachineFunction &MF) {
   // since the split part may not be placed in a contiguous region. It may also
   // be more beneficial to augment the linker to ensure contiguous layout of
   // split functions within the same section as specified by the attribute.
-  if (!MF.getFunction().getSection().empty())
+  if (MF.getFunction().hasSection() ||
+      MF.getFunction().hasFnAttribute("implicit-section-name"))
     return false;
 
   // We don't want to proceed further for cold functions
@@ -121,16 +123,28 @@ bool MachineFunctionSplitter::runOnMachineFunction(MachineFunction &MF) {
   auto *MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
   auto *PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
 
+  SmallVector<MachineBasicBlock *, 2> LandingPads;
   for (auto &MBB : MF) {
-    // FIXME: We retain the entry block and conservatively keep all landing pad
-    // blocks as part of the original function. Once D73739 is submitted, we can
-    // improve the handling of ehpads.
-    if ((MBB.pred_empty() || MBB.isEHPad()))
+    if (MBB.isEntryBlock())
       continue;
-    if (isColdBlock(MBB, MBFI, PSI))
+
+    if (MBB.isEHPad())
+      LandingPads.push_back(&MBB);
+    else if (isColdBlock(MBB, MBFI, PSI))
       MBB.setSectionID(MBBSectionID::ColdSectionID);
   }
 
+  // We only split out eh pads if all of them are cold.
+  bool HasHotLandingPads = false;
+  for (const MachineBasicBlock *LP : LandingPads) {
+    if (!isColdBlock(*LP, MBFI, PSI))
+      HasHotLandingPads = true;
+  }
+  if (!HasHotLandingPads) {
+    for (MachineBasicBlock *LP : LandingPads)
+      LP->setSectionID(MBBSectionID::ColdSectionID);
+  }
+
   auto Comparator = [](const MachineBasicBlock &X, const MachineBasicBlock &Y) {
     return X.getSectionID().Type < Y.getSectionID().Type;
   };