[WinEH] Permit branch folding in the face of funclets

Track which basic blocks belong to which funclets.  Permit branch
folding to fire but only if it can prove that doing so will not cause
code in one funclet to be reused in another.

llvm-svn: 249257

include/llvm/CodeGen/Analysis.h

@@ -15,6 +15,7 @@
 #define LLVM_CODEGEN_ANALYSIS_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/ISDOpcodes.h"
 #include "llvm/IR/CallSite.h"
@@ -23,6 +24,8 @@
 
 namespace llvm {
 class GlobalValue;
+class MachineBasicBlock;
+class MachineFunction;
 class TargetLoweringBase;
 class TargetLowering;
 class TargetMachine;
@@ -115,6 +118,9 @@ bool returnTypeIsEligibleForTailCall(const Function *F,
 // or we are in LTO.
 bool canBeOmittedFromSymbolTable(const GlobalValue *GV);
 
+DenseMap<const MachineBasicBlock *, int>
+getFuncletMembership(const MachineFunction &MF);
+
 } // End llvm namespace
 
 #endif

lib/CodeGen/Analysis.cpp

@@ -14,6 +14,7 @@
 #include "llvm/CodeGen/Analysis.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -25,6 +26,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Target/TargetLowering.h"
+#include "llvm/Target/TargetInstrInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 #include "llvm/Transforms/Utils/GlobalStatus.h"
 
@@ -643,3 +645,86 @@ bool llvm::canBeOmittedFromSymbolTable(const GlobalValue *GV) {
 
   return !GS.IsCompared;
 }
+
+static void collectFuncletMembers(
+    DenseMap<const MachineBasicBlock *, int> &FuncletMembership, int Funclet,
+    const MachineBasicBlock *MBB) {
+  // Don't revisit blocks.
+  if (FuncletMembership.count(MBB) > 0) {
+    if (FuncletMembership[MBB] != Funclet) {
+      assert(false && "MBB is part of two funclets!");
+      report_fatal_error("MBB is part of two funclets!");
+    }
+    return;
+  }
+
+  // Add this MBB to our funclet.
+  FuncletMembership[MBB] = Funclet;
+
+  bool IsReturn = false;
+  int NumTerminators = 0;
+  for (const MachineInstr &MI : MBB->terminators()) {
+    IsReturn |= MI.isReturn();
+    ++NumTerminators;
+  }
+  assert((!IsReturn || NumTerminators == 1) &&
+         "Expected only one terminator when a return is present!");
+
+  // Returns are boundaries where funclet transfer can occur, don't follow
+  // successors.
+  if (IsReturn)
+    return;
+
+  for (const MachineBasicBlock *SMBB : MBB->successors())
+    if (!SMBB->isEHPad())
+      collectFuncletMembers(FuncletMembership, Funclet, SMBB);
+}
+
+DenseMap<const MachineBasicBlock *, int>
+llvm::getFuncletMembership(const MachineFunction &MF) {
+  DenseMap<const MachineBasicBlock *, int> FuncletMembership;
+
+  // We don't have anything to do if there aren't any EH pads.
+  if (!MF.getMMI().hasEHFunclets())
+    return FuncletMembership;
+
+  bool IsSEH = isAsynchronousEHPersonality(
+      classifyEHPersonality(MF.getFunction()->getPersonalityFn()));
+
+  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+  SmallVector<const MachineBasicBlock *, 16> FuncletBlocks;
+  SmallVector<std::pair<const MachineBasicBlock *, int>, 16> CatchRetSuccessors;
+  for (const MachineBasicBlock &MBB : MF) {
+    if (MBB.isEHFuncletEntry())
+      FuncletBlocks.push_back(&MBB);
+
+    MachineBasicBlock::const_iterator MBBI = MBB.getFirstTerminator();
+    // CatchPads are not funclets for SEH so do not consider CatchRet to
+    // transfer control to another funclet.
+    if (IsSEH || MBBI->getOpcode() != TII->getCatchReturnOpcode())
+      continue;
+
+    const MachineBasicBlock *Successor = MBBI->getOperand(0).getMBB();
+    const MachineBasicBlock *SuccessorColor = MBBI->getOperand(1).getMBB();
+    CatchRetSuccessors.push_back({Successor, SuccessorColor->getNumber()});
+  }
+
+  // We don't have anything to do if there aren't any EH pads.
+  if (FuncletBlocks.empty())
+    return FuncletMembership;
+
+  // Identify all the basic blocks reachable from the function entry.
+  collectFuncletMembers(FuncletMembership, MF.front().getNumber(), MF.begin());
+  // Next, identify all the blocks inside the funclets.
+  for (const MachineBasicBlock *MBB : FuncletBlocks)
+    collectFuncletMembers(FuncletMembership, MBB->getNumber(), MBB);
+  // Finally, identify all the targets of a catchret.
+  for (std::pair<const MachineBasicBlock *, int> CatchRetPair :
+       CatchRetSuccessors)
+    collectFuncletMembers(FuncletMembership, CatchRetPair.second,
+                          CatchRetPair.first);
+  // All blocks not part of a funclet are in the parent function.
+  for (const MachineBasicBlock &MBB : MF)
+    FuncletMembership.insert({&MBB, MF.front().getNumber()});
+  return FuncletMembership;
+}

lib/CodeGen/BranchFolding.cpp

@@ -21,6 +21,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/Analysis.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
@@ -94,10 +95,8 @@ bool BranchFolderPass::runOnMachineFunction(MachineFunction &MF) {
 
   TargetPassConfig *PassConfig = &getAnalysis<TargetPassConfig>();
   // TailMerge can create jump into if branches that make CFG irreducible for
-  // HW that requires structurized CFG.  It can also cause BBs to get shared
-  // between funclets.
+  // HW that requires structurized CFG.
   bool EnableTailMerge = !MF.getTarget().requiresStructuredCFG() &&
-                         !MF.getMMI().hasEHFunclets() &&
                          PassConfig->getEnableTailMerge();
   BranchFolder Folder(EnableTailMerge, /*CommonHoist=*/true,
                       getAnalysis<MachineBlockFrequencyInfo>(),
@@ -135,6 +134,7 @@ void BranchFolder::RemoveDeadBlock(MachineBasicBlock *MBB) {
 
   // Remove the block.
   MF->erase(MBB);
+  FuncletMembership.erase(MBB);
 }
 
 /// OptimizeImpDefsBlock - If a basic block is just a bunch of implicit_def
@@ -222,6 +222,9 @@ bool BranchFolder::OptimizeFunction(MachineFunction &MF,
     MadeChange |= OptimizeImpDefsBlock(&MBB);
   }
 
+  // Recalculate funclet membership.
+  FuncletMembership = getFuncletMembership(MF);
+
   bool MadeChangeThisIteration = true;
   while (MadeChangeThisIteration) {
     MadeChangeThisIteration    = TailMergeBlocks(MF);
@@ -448,6 +451,11 @@ MachineBasicBlock *BranchFolder::SplitMBBAt(MachineBasicBlock &CurMBB,
   // For targets that use the register scavenger, we must maintain LiveIns.
   MaintainLiveIns(&CurMBB, NewMBB);
 
+  // Add the new block to the funclet.
+  const auto &FuncletI = FuncletMembership.find(&CurMBB);
+  if (FuncletI != FuncletMembership.end())
+    FuncletMembership[NewMBB] = FuncletI->second;
+
   return NewMBB;
 }
 
@@ -552,14 +560,23 @@ static unsigned CountTerminators(MachineBasicBlock *MBB,
 /// and decide if it would be profitable to merge those tails.  Return the
 /// length of the common tail and iterators to the first common instruction
 /// in each block.
-static bool ProfitableToMerge(MachineBasicBlock *MBB1,
-                              MachineBasicBlock *MBB2,
-                              unsigned minCommonTailLength,
-                              unsigned &CommonTailLen,
+static bool
+ProfitableToMerge(MachineBasicBlock *MBB1, MachineBasicBlock *MBB2,
+                  unsigned minCommonTailLength, unsigned &CommonTailLen,
                   MachineBasicBlock::iterator &I1,
-                              MachineBasicBlock::iterator &I2,
-                              MachineBasicBlock *SuccBB,
-                              MachineBasicBlock *PredBB) {
+                  MachineBasicBlock::iterator &I2, MachineBasicBlock *SuccBB,
+                  MachineBasicBlock *PredBB,
+                  DenseMap<const MachineBasicBlock *, int> &FuncletMembership) {
+  // It is never profitable to tail-merge blocks from two different funclets.
+  if (!FuncletMembership.empty()) {
+    auto Funclet1 = FuncletMembership.find(MBB1);
+    assert(Funclet1 != FuncletMembership.end());
+    auto Funclet2 = FuncletMembership.find(MBB2);
+    assert(Funclet2 != FuncletMembership.end());
+    if (Funclet1->second != Funclet2->second)
+      return false;
+  }
+
   CommonTailLen = ComputeCommonTailLength(MBB1, MBB2, I1, I2);
   if (CommonTailLen == 0)
     return false;
@@ -636,7 +653,8 @@ unsigned BranchFolder::ComputeSameTails(unsigned CurHash,
       if (ProfitableToMerge(CurMPIter->getBlock(), I->getBlock(),
                             minCommonTailLength,
                             CommonTailLen, TrialBBI1, TrialBBI2,
-                            SuccBB, PredBB)) {
+                            SuccBB, PredBB,
+                            FuncletMembership)) {
         if (CommonTailLen > maxCommonTailLength) {
           SameTails.clear();
           maxCommonTailLength = CommonTailLen;
@@ -1099,6 +1117,8 @@ bool BranchFolder::OptimizeBranches(MachineFunction &MF) {
 
   // Make sure blocks are numbered in order
   MF.RenumberBlocks();
+  // Renumbering blocks alters funclet membership, recalculate it.
+  FuncletMembership = getFuncletMembership(MF);
 
   for (MachineFunction::iterator I = std::next(MF.begin()), E = MF.end();
        I != E; ) {
@@ -1112,6 +1132,7 @@ bool BranchFolder::OptimizeBranches(MachineFunction &MF) {
       ++NumDeadBlocks;
     }
   }
+
   return MadeChange;
 }
 
@@ -1171,11 +1192,22 @@ ReoptimizeBlock:
   MachineFunction::iterator FallThrough = MBB;
   ++FallThrough;
 
+  // Make sure MBB and FallThrough belong to the same funclet.
+  bool SameFunclet = true;
+  if (!FuncletMembership.empty() && FallThrough != MF.end()) {
+    auto MBBFunclet = FuncletMembership.find(MBB);
+    assert(MBBFunclet != FuncletMembership.end());
+    auto FallThroughFunclet = FuncletMembership.find(FallThrough);
+    assert(FallThroughFunclet != FuncletMembership.end());
+    SameFunclet = MBBFunclet->second == FallThroughFunclet->second;
+  }
+
   // If this block is empty, make everyone use its fall-through, not the block
   // explicitly.  Landing pads should not do this since the landing-pad table
   // points to this block.  Blocks with their addresses taken shouldn't be
   // optimized away.
-  if (IsEmptyBlock(MBB) && !MBB->isEHPad() && !MBB->hasAddressTaken()) {
+  if (IsEmptyBlock(MBB) && !MBB->isEHPad() && !MBB->hasAddressTaken() &&
+      SameFunclet) {
     // Dead block?  Leave for cleanup later.
     if (MBB->pred_empty()) return MadeChange;
 

lib/CodeGen/BranchFolding.h

@@ -54,6 +54,7 @@ namespace llvm {
     typedef std::vector<MergePotentialsElt>::iterator MPIterator;
     std::vector<MergePotentialsElt> MergePotentials;
     SmallPtrSet<const MachineBasicBlock*, 2> TriedMerging;
+    DenseMap<const MachineBasicBlock *, int> FuncletMembership;
 
     class SameTailElt {
       MPIterator MPIter;

lib/CodeGen/FuncletLayout.cpp

@@ -12,12 +12,9 @@
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/CodeGen/Passes.h"
-#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/Analysis.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
-#include "llvm/CodeGen/MachineModuleInfo.h"
-#include "llvm/Target/TargetInstrInfo.h"
-#include "llvm/Target/TargetSubtargetInfo.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "funclet-layout"
@@ -34,81 +31,17 @@ public:
 };
 }
 
-static void
-collectFuncletMembers(DenseMap<MachineBasicBlock *, int> &FuncletMembership,
-                      int Funclet, MachineBasicBlock *MBB) {
-  // Don't revisit blocks.
-  if (FuncletMembership.count(MBB) > 0) {
-    if (FuncletMembership[MBB] != Funclet) {
-      assert(false && "MBB is part of two funclets!");
-      report_fatal_error("MBB is part of two funclets!");
-    }
-    return;
-  }
-
-  // Add this MBB to our funclet.
-  FuncletMembership[MBB] = Funclet;
-
-  bool IsReturn = false;
-  int NumTerminators = 0;
-  for (MachineInstr &MI : MBB->terminators()) {
-    IsReturn |= MI.isReturn();
-    ++NumTerminators;
-  }
-  assert((!IsReturn || NumTerminators == 1) &&
-         "Expected only one terminator when a return is present!");
-
-  // Returns are boundaries where funclet transfer can occur, don't follow
-  // successors.
-  if (IsReturn)
-    return;
-
-  for (MachineBasicBlock *SMBB : MBB->successors())
-    if (!SMBB->isEHPad())
-      collectFuncletMembers(FuncletMembership, Funclet, SMBB);
-}
-
 char FuncletLayout::ID = 0;
 char &llvm::FuncletLayoutID = FuncletLayout::ID;
 INITIALIZE_PASS(FuncletLayout, "funclet-layout",
                 "Contiguously Lay Out Funclets", false, false)
 
 bool FuncletLayout::runOnMachineFunction(MachineFunction &F) {
-  // We don't have anything to do if there aren't any EH pads.
-  if (!F.getMMI().hasEHFunclets())
+  DenseMap<const MachineBasicBlock *, int> FuncletMembership =
+      getFuncletMembership(F);
+  if (FuncletMembership.empty())
     return false;
 
-  const TargetInstrInfo *TII = F.getSubtarget().getInstrInfo();
-  SmallVector<MachineBasicBlock *, 16> FuncletBlocks;
-  SmallVector<std::pair<MachineBasicBlock *, int>, 16> CatchRetSuccessors;
-  for (MachineBasicBlock &MBB : F) {
-    if (MBB.isEHFuncletEntry())
-      FuncletBlocks.push_back(&MBB);
-
-    MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
-    if (MBBI->getOpcode() != TII->getCatchReturnOpcode())
-      continue;
-
-    MachineBasicBlock *Successor = MBBI->getOperand(0).getMBB();
-    MachineBasicBlock *SuccessorColor = MBBI->getOperand(1).getMBB();
-    CatchRetSuccessors.push_back({Successor, SuccessorColor->getNumber()});
-  }
-
-  // We don't have anything to do if there aren't any EH pads.
-  if (FuncletBlocks.empty())
-    return false;
-
-  DenseMap<MachineBasicBlock *, int> FuncletMembership;
-  // Identify all the basic blocks reachable from the function entry.
-  collectFuncletMembers(FuncletMembership, F.front().getNumber(), F.begin());
-  // Next, identify all the blocks inside the funclets.
-  for (MachineBasicBlock *MBB : FuncletBlocks)
-    collectFuncletMembers(FuncletMembership, MBB->getNumber(), MBB);
-  // Finally, identify all the targets of a catchret.
-  for (std::pair<MachineBasicBlock *, int> CatchRetPair : CatchRetSuccessors)
-    collectFuncletMembers(FuncletMembership, CatchRetPair.second,
-                          CatchRetPair.first);
-
   F.sort([&](MachineBasicBlock &x, MachineBasicBlock &y) {
     return FuncletMembership[&x] < FuncletMembership[&y];
   });

test/CodeGen/X86/funclet-layout.ll

@@ -103,7 +103,7 @@ unreachable:                                      ; preds = %catch, %entry
 ; CHECK: retq
 
 
-define void @test3() #0 personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+define void @test3(i1 %V) #0 personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
 entry:
   invoke void @g()
           to label %try.cont unwind label %catch.dispatch
@@ -128,7 +128,15 @@ catchendblock:                                    ; preds = %catch.dispatch.1
   catchendpad unwind to caller
 
 try.cont:                                         ; preds = %entry
-  ret void
+  br i1 %V, label %exit_one, label %exit_two
+
+exit_one:
+  tail call void @exit(i32 0)
+  unreachable
+
+exit_two:
+  tail call void @exit(i32 0)
+  unreachable
 }
 
 ; CHECK-LABEL: test3:
@@ -136,15 +144,20 @@ try.cont:                                         ; preds = %entry
 ; The entry funclet contains %entry and %try.cont
 ; CHECK: # %entry
 ; CHECK: # %try.cont
-; CHECK: retq
+; CHECK: callq exit
+; CHECK-NOT: # exit_one
+; CHECK-NOT: # exit_two
+; CHECK: ud2
 
 ; The catch(int) funclet contains %catch.2
 ; CHECK: # %catch.2
 ; CHECK: callq exit
+; CHECK: ud2
 
 ; The catch(...) funclet contains %catch
 ; CHECK: # %catch{{$}}
 ; CHECK: callq exit
+; CHECK: ud2
 
 declare void @exit(i32) noreturn nounwind
 declare void @_CxxThrowException(i8*, %eh.ThrowInfo*)