[CallSiteInfo] Handle bundles when updating call site info

This will address the issue: P8198 and P8199 (from D73534).

The methods was not handle bundles properly.

Differential Revision: https://reviews.llvm.org/D74904

(cherry picked from commit 016d91c)

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -998,21 +998,19 @@ class MachineFunction {
   /// Following functions update call site info. They should be called before
   /// removing, replacing or copying call instruction.
 
-  /// Move the call site info from \p Old to \New call site info. This function
-  /// is used when we are replacing one call instruction with another one to
-  /// the same callee.
-  void moveCallSiteInfo(const MachineInstr *Old,
-                        const MachineInstr *New);
-
   /// Erase the call site info for \p MI. It is used to remove a call
   /// instruction from the instruction stream.
   void eraseCallSiteInfo(const MachineInstr *MI);
-
   /// Copy the call site info from \p Old to \ New. Its usage is when we are
   /// making a copy of the instruction that will be inserted at different point
   /// of the instruction stream.
   void copyCallSiteInfo(const MachineInstr *Old,
                         const MachineInstr *New);
+  /// Move the call site info from \p Old to \New call site info. This function
+  /// is used when we are replacing one call instruction with another one to
+  /// the same callee.
+  void moveCallSiteInfo(const MachineInstr *Old,
+                        const MachineInstr *New);
 };
 
 //===--------------------------------------------------------------------===//

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -685,7 +685,11 @@ class MachineInstr
 
   /// Return true if this is a call instruction that may have an associated
   /// call site entry in the debug info.
-  bool isCandidateForCallSiteEntry() const;
+  bool isCandidateForCallSiteEntry(QueryType Type = IgnoreBundle) const;
+  /// Return true if copying, moving, or erasing this instruction requires
+  /// updating Call Site Info (see \ref copyCallSiteInfo, \ref moveCallSiteInfo,
+  /// \ref eraseCallSiteInfo).
+  bool shouldUpdateCallSiteInfo() const;
 
   /// Returns true if the specified instruction stops control flow
   /// from executing the instruction immediately following it.  Examples include

llvm/lib/CodeGen/BranchFolding.cpp

@@ -170,7 +170,7 @@ void BranchFolder::RemoveDeadBlock(MachineBasicBlock *MBB) {
 
   // Update call site info.
   std::for_each(MBB->begin(), MBB->end(), [MF](const MachineInstr &MI) {
-    if (MI.isCandidateForCallSiteEntry())
+    if (MI.shouldUpdateCallSiteInfo())
       MF->eraseCallSiteInfo(&MI);
   });
   // Remove the block.

llvm/lib/CodeGen/IfConversion.cpp

@@ -1851,7 +1851,7 @@ bool IfConverter::IfConvertDiamondCommon(
   while (NumDups1 != 0) {
     // Since this instruction is going to be deleted, update call
     // site info state if the instruction is call instruction.
-    if (DI2->isCandidateForCallSiteEntry())
+    if (DI2->shouldUpdateCallSiteInfo())
       MBB2.getParent()->eraseCallSiteInfo(&*DI2);
 
     ++DI2;
@@ -1900,7 +1900,7 @@ bool IfConverter::IfConvertDiamondCommon(
 
     // Since this instruction is going to be deleted, update call
     // site info state if the instruction is call instruction.
-    if (DI1->isCandidateForCallSiteEntry())
+    if (DI1->shouldUpdateCallSiteInfo())
       MBB1.getParent()->eraseCallSiteInfo(&*DI1);
 
     // skip dbg_value instructions

llvm/lib/CodeGen/LiveRangeEdit.cpp

@@ -232,7 +232,7 @@ bool LiveRangeEdit::foldAsLoad(LiveInterval *LI,
   LLVM_DEBUG(dbgs() << "                folded: " << *FoldMI);
   LIS.ReplaceMachineInstrInMaps(*UseMI, *FoldMI);
   // Update the call site info.
-  if (UseMI->isCandidateForCallSiteEntry())
+  if (UseMI->shouldUpdateCallSiteInfo())
     UseMI->getMF()->moveCallSiteInfo(UseMI, FoldMI);
   UseMI->eraseFromParent();
   DefMI->addRegisterDead(LI->reg, nullptr);

llvm/lib/CodeGen/MachineFunction.cpp

@@ -856,48 +856,68 @@ MachineFunction::getCallSiteInfo(const MachineInstr *MI) {
   return CallSitesInfo.find(MI);
 }
 
-void MachineFunction::moveCallSiteInfo(const MachineInstr *Old,
-                                       const MachineInstr *New) {
-  assert(Old->isCandidateForCallSiteEntry() &&
-         "Call site info refers only to call (MI) candidates");
-
-  if (!New->isCandidateForCallSiteEntry())
-    return eraseCallSiteInfo(Old);
+/// Return the call machine instruction or find a call within bundle.
+static const MachineInstr *getCallInstr(const MachineInstr *MI) {
+  if (!MI->isBundle())
+    return MI;
 
-  CallSiteInfoMap::iterator CSIt = getCallSiteInfo(Old);
-  if (CSIt == CallSitesInfo.end())
-    return;
+  for (auto &BMI : make_range(getBundleStart(MI->getIterator()),
+                              getBundleEnd(MI->getIterator())))
+    if (BMI.isCandidateForCallSiteEntry())
+      return &BMI;
 
-  CallSiteInfo CSInfo = std::move(CSIt->second);
-  CallSitesInfo.erase(CSIt);
-  CallSitesInfo[New] = CSInfo;
+  llvm_unreachable("Unexpected bundle without a call site candidate");
 }
 
 void MachineFunction::eraseCallSiteInfo(const MachineInstr *MI) {
-  assert(MI->isCandidateForCallSiteEntry() &&
-         "Call site info refers only to call (MI) candidates");
-  CallSiteInfoMap::iterator CSIt = getCallSiteInfo(MI);
+  assert(MI->shouldUpdateCallSiteInfo() &&
+         "Call site info refers only to call (MI) candidates or "
+         "candidates inside bundles");
+
+  const MachineInstr *CallMI = getCallInstr(MI);
+  CallSiteInfoMap::iterator CSIt = getCallSiteInfo(CallMI);
   if (CSIt == CallSitesInfo.end())
     return;
   CallSitesInfo.erase(CSIt);
 }
 
 void MachineFunction::copyCallSiteInfo(const MachineInstr *Old,
                                        const MachineInstr *New) {
-  assert(Old->isCandidateForCallSiteEntry() &&
-         "Call site info refers only to call (MI) candidates");
+  assert(Old->shouldUpdateCallSiteInfo() &&
+         "Call site info refers only to call (MI) candidates or "
+         "candidates inside bundles");
 
   if (!New->isCandidateForCallSiteEntry())
     return eraseCallSiteInfo(Old);
 
-  CallSiteInfoMap::iterator CSIt = getCallSiteInfo(Old);
+  const MachineInstr *OldCallMI = getCallInstr(Old);
+  CallSiteInfoMap::iterator CSIt = getCallSiteInfo(OldCallMI);
   if (CSIt == CallSitesInfo.end())
     return;
 
   CallSiteInfo CSInfo = CSIt->second;
   CallSitesInfo[New] = CSInfo;
 }
 
+void MachineFunction::moveCallSiteInfo(const MachineInstr *Old,
+                                       const MachineInstr *New) {
+  assert(Old->shouldUpdateCallSiteInfo() &&
+         "Call site info refers only to call (MI) candidates or "
+         "candidates inside bundles");
+
+  if (!New->isCandidateForCallSiteEntry())
+    return eraseCallSiteInfo(Old);
+
+  const MachineInstr *OldCallMI = getCallInstr(Old);
+  CallSiteInfoMap::iterator CSIt = getCallSiteInfo(OldCallMI);
+  if (CSIt == CallSitesInfo.end())
+    return;
+
+  CallSiteInfo CSInfo = std::move(CSIt->second);
+  CallSitesInfo.erase(CSIt);
+  CallSitesInfo[New] = CSInfo;
+}
+
 /// \}
 
 //===----------------------------------------------------------------------===//

llvm/lib/CodeGen/MachineInstr.cpp

@@ -696,8 +696,8 @@ void MachineInstr::eraseFromBundle() {
   getParent()->erase_instr(this);
 }
 
-bool MachineInstr::isCandidateForCallSiteEntry() const {
-  if (!isCall(MachineInstr::IgnoreBundle))
+bool MachineInstr::isCandidateForCallSiteEntry(QueryType Type) const {
+  if (!isCall(Type))
     return false;
   switch (getOpcode()) {
   case TargetOpcode::PATCHABLE_EVENT_CALL:
@@ -710,6 +710,12 @@ bool MachineInstr::isCandidateForCallSiteEntry() const {
   return true;
 }
 
+bool MachineInstr::shouldUpdateCallSiteInfo() const {
+  if (isBundle())
+    return isCandidateForCallSiteEntry(MachineInstr::AnyInBundle);
+  return isCandidateForCallSiteEntry();
+}
+
 unsigned MachineInstr::getNumExplicitOperands() const {
   unsigned NumOperands = MCID->getNumOperands();
   if (!MCID->isVariadic())

llvm/lib/CodeGen/MachineLICM.cpp

@@ -1369,7 +1369,7 @@ MachineInstr *MachineLICMBase::ExtractHoistableLoad(MachineInstr *MI) {
   // Otherwise we successfully unfolded a load that we can hoist.
 
   // Update the call site info.
-  if (MI->isCandidateForCallSiteEntry())
+  if (MI->shouldUpdateCallSiteInfo())
     MF.eraseCallSiteInfo(MI);
 
   MI->eraseFromParent();

llvm/lib/CodeGen/MachineOutliner.cpp

@@ -1262,7 +1262,7 @@ bool MachineOutliner::outline(Module &M,
                   MOP.getReg(), true, /* isDef = true */
                   true /* isImp = true */));
           }
-          if (MI.isCandidateForCallSiteEntry())
+          if (MI.shouldUpdateCallSiteInfo())
             MI.getMF()->eraseCallSiteInfo(&MI);
         };
         // Copy over the defs in the outlined range.

llvm/lib/CodeGen/PeepholeOptimizer.cpp

@@ -1777,7 +1777,7 @@ bool PeepholeOptimizer::runOnMachineFunction(MachineFunction &MF) {
               LocalMIs.erase(DefMI);
               LocalMIs.insert(FoldMI);
               // Update the call site info.
-              if (MI->isCandidateForCallSiteEntry())
+              if (MI->shouldUpdateCallSiteInfo())
                 MI->getMF()->moveCallSiteInfo(MI, FoldMI);
               MI->eraseFromParent();
               DefMI->eraseFromParent();

llvm/lib/CodeGen/TailDuplicator.cpp

@@ -1005,7 +1005,7 @@ void TailDuplicator::removeDeadBlock(
   MachineFunction *MF = MBB->getParent();
   // Update the call site info.
   std::for_each(MBB->begin(), MBB->end(), [MF](const MachineInstr &MI) {
-    if (MI.isCandidateForCallSiteEntry())
+    if (MI.shouldUpdateCallSiteInfo())
       MF->eraseCallSiteInfo(&MI);
   });
 

llvm/lib/CodeGen/TargetInstrInfo.cpp

@@ -142,7 +142,7 @@ TargetInstrInfo::ReplaceTailWithBranchTo(MachineBasicBlock::iterator Tail,
   // from the end of MBB.
   while (Tail != MBB->end()) {
     auto MI = Tail++;
-    if (MI->isCandidateForCallSiteEntry())
+    if (MI->shouldUpdateCallSiteInfo())
       MBB->getParent()->eraseCallSiteInfo(&*MI);
     MBB->erase(MI);
   }

llvm/lib/CodeGen/UnreachableBlockElim.cpp

@@ -151,7 +151,7 @@ bool UnreachableMachineBlockElim::runOnMachineFunction(MachineFunction &F) {
   for (unsigned i = 0, e = DeadBlocks.size(); i != e; ++i) {
     // Remove any call site information for calls in the block.
     for (auto &I : DeadBlocks[i]->instrs())
-      if (I.isCandidateForCallSiteEntry())
+      if (I.shouldUpdateCallSiteInfo())
         DeadBlocks[i]->getParent()->eraseCallSiteInfo(&I);
 
     DeadBlocks[i]->eraseFromParent();

llvm/lib/CodeGen/XRayInstrumentation.cpp

@@ -111,7 +111,7 @@ void XRayInstrumentation::replaceRetWithPatchableRet(
         for (auto &MO : T.operands())
           MIB.add(MO);
         Terminators.push_back(&T);
-        if (T.isCandidateForCallSiteEntry())
+        if (T.shouldUpdateCallSiteInfo())
           MF.eraseCallSiteInfo(&T);
       }
     }

llvm/lib/Target/AArch64/AArch64CleanupLocalDynamicTLSPass.cpp

@@ -106,7 +106,7 @@ struct LDTLSCleanup : public MachineFunctionPass {
                              .addReg(TLSBaseAddrReg);
 
     // Update the call site info.
-    if (I.isCandidateForCallSiteEntry())
+    if (I.shouldUpdateCallSiteInfo())
       I.getMF()->eraseCallSiteInfo(&I);
 
     // Erase the TLS_base_addr instruction.