[WinEH] Use operand bundles to describe call sites

SimplifyCFG allows tail merging with code which terminates in
unreachable which, in turn, makes it possible for an invoke to end up in
a funclet which it was not originally part of.

Using operand bundles on invokes allows us to determine whether or not
an invoke was part of a funclet in the source program.

Furthermore, it allows us to unambiguously answer questions about the
legality of inlining into call sites which the personality may have
trouble with.

Differential Revision: http://reviews.llvm.org/D15517

llvm-svn: 255674

llvm/docs/LangRef.rst

@@ -1571,6 +1571,15 @@ caller's deoptimization state to the callee's deoptimization state is
 semantically equivalent to composing the caller's deoptimization
 continuation after the callee's deoptimization continuation.
 
+Funclet Operand Bundles
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Funclet operand bundles are characterized by the ``"funclet"``
+operand bundle tag.  These operand bundles indicate that a call site
+is within a particular funclet.  There can be at most one
+``"funclet"`` operand bundle attached to a call site and it must have
+exactly one bundle operand.
+
 .. _moduleasm:
 
 Module-Level Inline Assembly

llvm/include/llvm/IR/InstrTypes.h

@@ -1219,6 +1219,11 @@ struct OperandBundleUse {
     return getTagID() == LLVMContext::OB_deopt;
   }
 
+  /// \brief Return true if this is a "funclet" operand bundle.
+  bool isFuncletOperandBundle() const {
+    return getTagID() == LLVMContext::OB_funclet;
+  }
+
 private:
   /// \brief Pointer to an entry in LLVMContextImpl::getOrInsertBundleTag.
   StringMapEntry<uint32_t> *Tag;
@@ -1237,6 +1242,8 @@ template <typename InputTy> class OperandBundleDefT {
 public:
   explicit OperandBundleDefT(std::string Tag, std::vector<InputTy> Inputs)
       : Tag(std::move(Tag)), Inputs(std::move(Inputs)) {}
+  explicit OperandBundleDefT(std::string Tag, ArrayRef<InputTy> Inputs)
+      : Tag(std::move(Tag)), Inputs(Inputs) {}
 
   explicit OperandBundleDefT(const OperandBundleUse &OBU) {
     Tag = OBU.getTagName();
@@ -1430,11 +1437,12 @@ template <typename InstrTy, typename OpIteratorTy> class OperandBundleUser {
   /// may write to the heap.
   bool hasClobberingOperandBundles() const {
     for (auto &BOI : bundle_op_infos()) {
-      if (BOI.Tag->second == LLVMContext::OB_deopt)
+      if (BOI.Tag->second == LLVMContext::OB_deopt ||
+          BOI.Tag->second == LLVMContext::OB_funclet)
         continue;
 
-      // This instruction has an operand bundle that is not a "deopt" operand
-      // bundle.  Assume the worst.
+      // This instruction has an operand bundle that is not known to us.
+      // Assume the worst.
       return true;
     }
 

llvm/include/llvm/IR/LLVMContext.h

@@ -72,7 +72,8 @@ class LLVMContext {
   /// Additionally, this scheme allows LLVM to efficiently check for specific
   /// operand bundle tags without comparing strings.
   enum {
-    OB_deopt = 0,  // "deopt"
+    OB_deopt = 0,   // "deopt"
+    OB_funclet = 1, // "funclet"
   };
 
   /// getMDKindID - Return a unique non-zero ID for the specified metadata kind.

llvm/include/llvm/Transforms/Utils/Local.h

@@ -288,6 +288,10 @@ bool replaceDbgDeclare(Value *Address, Value *NewAddress,
 bool replaceDbgDeclareForAlloca(AllocaInst *AI, Value *NewAllocaAddress,
                                 DIBuilder &Builder, bool Deref, int Offset = 0);
 
+/// \brief Insert an unreachable instruction before the specified
+/// instruction, making it and the rest of the code in the block dead.
+void changeToUnreachable(Instruction *I, bool UseLLVMTrap);
+
 /// Replace 'BB's terminator with one that does not have an unwind successor
 /// block.  Rewrites `invoke` to `call`, etc.  Updates any PHIs in unwind
 /// successor.

llvm/lib/CodeGen/WinEHPrepare.cpp

@@ -74,7 +74,7 @@ class WinEHPrepare : public FunctionPass {
 
   void demotePHIsOnFunclets(Function &F);
   void cloneCommonBlocks(Function &F);
-  void removeImplausibleTerminators(Function &F);
+  void removeImplausibleInstructions(Function &F);
   void cleanupPreparedFunclets(Function &F);
   void verifyPreparedFunclets(Function &F);
 
@@ -765,19 +765,56 @@ void WinEHPrepare::cloneCommonBlocks(Function &F) {
   }
 }
 
-void WinEHPrepare::removeImplausibleTerminators(Function &F) {
+void WinEHPrepare::removeImplausibleInstructions(Function &F) {
   // Remove implausible terminators and replace them with UnreachableInst.
   for (auto &Funclet : FuncletBlocks) {
     BasicBlock *FuncletPadBB = Funclet.first;
     std::vector<BasicBlock *> &BlocksInFunclet = Funclet.second;
-    Instruction *FuncletPadInst = FuncletPadBB->getFirstNonPHI();
-    auto *CatchPad = dyn_cast<CatchPadInst>(FuncletPadInst);
-    auto *CleanupPad = dyn_cast<CleanupPadInst>(FuncletPadInst);
+    Instruction *FirstNonPHI = FuncletPadBB->getFirstNonPHI();
+    auto *FuncletPad = dyn_cast<FuncletPadInst>(FirstNonPHI);
+    auto *CatchPad = dyn_cast_or_null<CatchPadInst>(FuncletPad);
+    auto *CleanupPad = dyn_cast_or_null<CleanupPadInst>(FuncletPad);
 
     for (BasicBlock *BB : BlocksInFunclet) {
+      for (Instruction &I : *BB) {
+        CallSite CS(&I);
+        if (!CS)
+          continue;
+
+        Value *FuncletBundleOperand = nullptr;
+        if (auto BU = CS.getOperandBundle(LLVMContext::OB_funclet))
+          FuncletBundleOperand = BU->Inputs.front();
+
+        if (FuncletBundleOperand == FuncletPad)
+          continue;
+
+        // Skip call sites which are nounwind intrinsics.
+        auto *CalledFn =
+            dyn_cast<Function>(CS.getCalledValue()->stripPointerCasts());
+        if (CalledFn && CalledFn->isIntrinsic() && CS.doesNotThrow())
+          continue;
+
+        // This call site was not part of this funclet, remove it.
+        if (CS.isInvoke()) {
+          // Remove the unwind edge if it was an invoke.
+          removeUnwindEdge(BB);
+          // Get a pointer to the new call.
+          BasicBlock::iterator CallI =
+              std::prev(BB->getTerminator()->getIterator());
+          auto *CI = cast<CallInst>(&*CallI);
+          changeToUnreachable(CI, /*UseLLVMTrap=*/false);
+        } else {
+          changeToUnreachable(&I, /*UseLLVMTrap=*/false);
+        }
+
+        // There are no more instructions in the block (except for unreachable),
+        // we are done.
+        break;
+      }
+
       TerminatorInst *TI = BB->getTerminator();
       // CatchPadInst and CleanupPadInst can't transfer control to a ReturnInst.
-      bool IsUnreachableRet = isa<ReturnInst>(TI) && (CatchPad || CleanupPad);
+      bool IsUnreachableRet = isa<ReturnInst>(TI) && FuncletPad;
       // The token consumed by a CatchReturnInst must match the funclet token.
       bool IsUnreachableCatchret = false;
       if (auto *CRI = dyn_cast<CatchReturnInst>(TI))
@@ -788,19 +825,14 @@ void WinEHPrepare::removeImplausibleTerminators(Function &F) {
         IsUnreachableCleanupret = CRI->getCleanupPad() != CleanupPad;
       if (IsUnreachableRet || IsUnreachableCatchret ||
           IsUnreachableCleanupret) {
-        // Loop through all of our successors and make sure they know that one
-        // of their predecessors is going away.
-        for (BasicBlock *SuccBB : TI->successors())
-          SuccBB->removePredecessor(BB);
-
-        new UnreachableInst(BB->getContext(), TI);
-        TI->eraseFromParent();
+        changeToUnreachable(TI, /*UseLLVMTrap=*/false);
       } else if (isa<InvokeInst>(TI)) {
-        // Invokes within a cleanuppad for the MSVC++ personality never
-        // transfer control to their unwind edge: the personality will
-        // terminate the program.
-        if (Personality == EHPersonality::MSVC_CXX && CleanupPad)
+        if (Personality == EHPersonality::MSVC_CXX && CleanupPad) {
+          // Invokes within a cleanuppad for the MSVC++ personality never
+          // transfer control to their unwind edge: the personality will
+          // terminate the program.
           removeUnwindEdge(BB);
+        }
       }
     }
   }
@@ -854,7 +886,7 @@ bool WinEHPrepare::prepareExplicitEH(Function &F) {
     demotePHIsOnFunclets(F);
 
   if (!DisableCleanups) {
-    removeImplausibleTerminators(F);
+    removeImplausibleInstructions(F);
 
     cleanupPreparedFunclets(F);
   }

llvm/lib/IR/LLVMContext.cpp

@@ -132,6 +132,11 @@ LLVMContext::LLVMContext() : pImpl(new LLVMContextImpl(*this)) {
   assert(DeoptEntry->second == LLVMContext::OB_deopt &&
          "deopt operand bundle id drifted!");
   (void)DeoptEntry;
+
+  auto *FuncletEntry = pImpl->getOrInsertBundleTag("funclet");
+  assert(FuncletEntry->second == LLVMContext::OB_funclet &&
+         "funclet operand bundle id drifted!");
+  (void)FuncletEntry;
 }
 LLVMContext::~LLVMContext() { delete pImpl; }
 

llvm/lib/IR/Verifier.cpp

@@ -2381,13 +2381,25 @@ void Verifier::VerifyCallSite(CallSite CS) {
     if (Intrinsic::ID ID = (Intrinsic::ID)F->getIntrinsicID())
       visitIntrinsicCallSite(ID, CS);
 
-  // Verify that a callsite has at most one "deopt" operand bundle.
-  bool FoundDeoptBundle = false;
+  // Verify that a callsite has at most one "deopt" and one "funclet" operand
+  // bundle.
+  bool FoundDeoptBundle = false, FoundFuncletBundle = false;
   for (unsigned i = 0, e = CS.getNumOperandBundles(); i < e; ++i) {
-    if (CS.getOperandBundleAt(i).getTagID() == LLVMContext::OB_deopt) {
+    OperandBundleUse BU = CS.getOperandBundleAt(i);
+    uint32_t Tag = BU.getTagID();
+    if (Tag == LLVMContext::OB_deopt) {
       Assert(!FoundDeoptBundle, "Multiple deopt operand bundles", I);
       FoundDeoptBundle = true;
     }
+    if (Tag == LLVMContext::OB_funclet) {
+      Assert(!FoundFuncletBundle, "Multiple funclet operand bundles", I);
+      FoundFuncletBundle = true;
+      Assert(BU.Inputs.size() == 1,
+             "Expected exactly one funclet bundle operand", I);
+      Assert(isa<FuncletPadInst>(BU.Inputs.front()),
+             "Funclet bundle operands should correspond to a FuncletPadInst",
+             I);
+    }
   }
 
   visitInstruction(*I);
@@ -3002,6 +3014,8 @@ void Verifier::visitCleanupPadInst(CleanupPadInst &CPI) {
       UnwindDest = CRI->getUnwindDest();
     } else if (isa<CleanupPadInst>(U) || isa<CatchSwitchInst>(U)) {
       continue;
+    } else if (CallSite(U)) {
+      continue;
     } else {
       Assert(false, "bogus cleanuppad use", &CPI);
     }

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -1035,12 +1035,17 @@ bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
   // The inliner does not know how to inline through calls with operand bundles
   // in general ...
   if (CS.hasOperandBundles()) {
-    // ... but it knows how to inline through "deopt" operand bundles.
-    bool CanInline =
-        CS.getNumOperandBundles() == 1 &&
-        CS.getOperandBundleAt(0).getTagID() == LLVMContext::OB_deopt;
-    if (!CanInline)
+    for (int i = 0, e = CS.getNumOperandBundles(); i != e; ++i) {
+      uint32_t Tag = CS.getOperandBundleAt(i).getTagID();
+      // ... but it knows how to inline through "deopt" operand bundles ...
+      if (Tag == LLVMContext::OB_deopt)
+        continue;
+      // ... and "funclet" operand bundles.
+      if (Tag == LLVMContext::OB_funclet)
+        continue;
+
       return false;
+    }
   }
 
   // If the call to the callee cannot throw, set the 'nounwind' flag on any
@@ -1088,23 +1093,13 @@ bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
   // We need to figure out which funclet the callsite was in so that we may
   // properly nest the callee.
   Instruction *CallSiteEHPad = nullptr;
-  if (CalledPersonality && CallerPersonality) {
-    EHPersonality Personality = classifyEHPersonality(CalledPersonality);
+  if (CallerPersonality) {
+    EHPersonality Personality = classifyEHPersonality(CallerPersonality);
     if (isFuncletEHPersonality(Personality)) {
-      DenseMap<BasicBlock *, ColorVector> CallerBlockColors =
-          colorEHFunclets(*Caller);
-      ColorVector &CallSiteColors = CallerBlockColors[OrigBB];
-      size_t NumColors = CallSiteColors.size();
-      // There is no single parent, inlining will not succeed.
-      if (NumColors > 1)
-        return false;
-      if (NumColors == 1) {
-        BasicBlock *CallSiteFuncletBB = CallSiteColors.front();
-        if (CallSiteFuncletBB != Caller->begin()) {
-          CallSiteEHPad = CallSiteFuncletBB->getFirstNonPHI();
-          assert(CallSiteEHPad->isEHPad() && "Expected an EHPad!");
-        }
-      }
+      Optional<OperandBundleUse> ParentFunclet =
+          CS.getOperandBundle(LLVMContext::OB_funclet);
+      if (ParentFunclet)
+        CallSiteEHPad = cast<FuncletPadInst>(ParentFunclet->Inputs.front());
 
       // OK, the inlining site is legal.  What about the target function?
 
@@ -1116,7 +1111,7 @@ bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
             // Ok, the call site is within a cleanuppad.  Let's check the callee
             // for catchpads.
             for (const BasicBlock &CalledBB : *CalledFunc) {
-              if (isa<CatchPadInst>(CalledBB.getFirstNonPHI()))
+              if (isa<CatchSwitchInst>(CalledBB.getFirstNonPHI()))
                 return false;
             }
           }
@@ -1195,11 +1190,9 @@ bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
       HandleByValArgumentInit(Init.first, Init.second, Caller->getParent(),
                               &*FirstNewBlock, IFI);
 
-    if (CS.hasOperandBundles()) {
-      auto ParentDeopt = CS.getOperandBundleAt(0);
-      assert(ParentDeopt.getTagID() == LLVMContext::OB_deopt &&
-             "Checked on entry!");
-
+    Optional<OperandBundleUse> ParentDeopt =
+        CS.getOperandBundle(LLVMContext::OB_deopt);
+    if (ParentDeopt) {
       SmallVector<OperandBundleDef, 2> OpDefs;
 
       for (auto &VH : InlinedFunctionInfo.OperandBundleCallSites) {
@@ -1225,12 +1218,12 @@ bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
           // Prepend the parent's deoptimization continuation to the newly
           // inlined call's deoptimization continuation.
           std::vector<Value *> MergedDeoptArgs;
-          MergedDeoptArgs.reserve(ParentDeopt.Inputs.size() +
+          MergedDeoptArgs.reserve(ParentDeopt->Inputs.size() +
                                   ChildOB.Inputs.size());
 
           MergedDeoptArgs.insert(MergedDeoptArgs.end(),
-                                 ParentDeopt.Inputs.begin(),
-                                 ParentDeopt.Inputs.end());
+                                 ParentDeopt->Inputs.begin(),
+                                 ParentDeopt->Inputs.end());
           MergedDeoptArgs.insert(MergedDeoptArgs.end(), ChildOB.Inputs.begin(),
                                  ChildOB.Inputs.end());
 
@@ -1423,12 +1416,48 @@ bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
     }
   }
 
-  // Update the lexical scopes of the new funclets.  Anything that had 'none' as
-  // its parent is now nested inside the callsite's EHPad.
+  // Update the lexical scopes of the new funclets and callsites.
+  // Anything that had 'none' as its parent is now nested inside the callsite's
+  // EHPad.
+
   if (CallSiteEHPad) {
     for (Function::iterator BB = FirstNewBlock->getIterator(),
                             E = Caller->end();
          BB != E; ++BB) {
+      // Add bundle operands to any top-level call sites.
+      SmallVector<OperandBundleDef, 1> OpBundles;
+      for (BasicBlock::iterator BBI = BB->begin(), E = BB->end(); BBI != E;) {
+        Instruction *I = &*BBI++;
+        CallSite CS(I);
+        if (!CS)
+          continue;
+
+        // Skip call sites which are nounwind intrinsics.
+        auto *CalledFn =
+            dyn_cast<Function>(CS.getCalledValue()->stripPointerCasts());
+        if (CalledFn && CalledFn->isIntrinsic() && CS.doesNotThrow())
+          continue;
+
+        // Skip call sites which already have a "funclet" bundle.
+        if (CS.getOperandBundle(LLVMContext::OB_funclet))
+          continue;
+
+        CS.getOperandBundlesAsDefs(OpBundles);
+        OpBundles.emplace_back("funclet", CallSiteEHPad);
+
+        Instruction *NewInst;
+        if (CS.isCall())
+          NewInst = CallInst::Create(cast<CallInst>(I), OpBundles, I);
+        else
+          NewInst = InvokeInst::Create(cast<InvokeInst>(I), OpBundles, I);
+        NewInst->setDebugLoc(I->getDebugLoc());
+        NewInst->takeName(I);
+        I->replaceAllUsesWith(NewInst);
+        I->eraseFromParent();
+
+        OpBundles.clear();
+      }
+
       Instruction *I = BB->getFirstNonPHI();
       if (!I->isEHPad())
         continue;

llvm/lib/Transforms/Utils/Local.cpp

@@ -1180,9 +1180,7 @@ bool llvm::replaceDbgDeclareForAlloca(AllocaInst *AI, Value *NewAllocaAddress,
                            Deref, Offset);
 }
 
-/// changeToUnreachable - Insert an unreachable instruction before the specified
-/// instruction, making it and the rest of the code in the block dead.
-static void changeToUnreachable(Instruction *I, bool UseLLVMTrap) {
+void llvm::changeToUnreachable(Instruction *I, bool UseLLVMTrap) {
   BasicBlock *BB = I->getParent();
   // Loop over all of the successors, removing BB's entry from any PHI
   // nodes.