[aarch64][win] Update Called Globals info when updating Call Site info (#122762)

Fixes the "use after poison" issue introduced by #121516 (see
<#121516 (comment)>).

The root cause of this issue is that #121516 introduced "Called Global"
information for call instructions modeling how "Call Site" info is
stored in the machine function, HOWEVER it didn't copy the
copy/move/erase operations for call site information.

The fix is to rename and update the existing copy/move/erase functions
so they also take care of Called Global info.

Author: dpaoliello

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -354,11 +354,6 @@ class LLVM_ABI MachineFunction {
   /// a table of valid targets for Windows EHCont Guard.
   std::vector<MCSymbol *> CatchretTargets;
 
-  /// Mapping of call instruction to the global value and target flags that it
-  /// calls, if applicable.
-  DenseMap<const MachineInstr *, std::pair<const GlobalValue *, unsigned>>
-      CalledGlobalsMap;
-
   /// \name Exception Handling
   /// \{
 
@@ -494,6 +489,11 @@ class LLVM_ABI MachineFunction {
     SmallVector<ArgRegPair, 1> ArgRegPairs;
   };
 
+  struct CalledGlobalInfo {
+    const GlobalValue *Callee;
+    unsigned TargetFlags;
+  };
+
 private:
   Delegate *TheDelegate = nullptr;
   GISelChangeObserver *Observer = nullptr;
@@ -506,6 +506,11 @@ class LLVM_ABI MachineFunction {
   /// instruction if debug entry value support is enabled.
   CallSiteInfoMap::iterator getCallSiteInfo(const MachineInstr *MI);
 
+  using CalledGlobalsMap = DenseMap<const MachineInstr *, CalledGlobalInfo>;
+  /// Mapping of call instruction to the global value and target flags that it
+  /// calls, if applicable.
+  CalledGlobalsMap CalledGlobalsInfo;
+
   // Callbacks for insertion and removal.
   void handleInsertion(MachineInstr &MI);
   void handleRemoval(MachineInstr &MI);
@@ -1189,22 +1194,20 @@ class LLVM_ABI MachineFunction {
 
   /// Tries to get the global and target flags for a call site, if the
   /// instruction is a call to a global.
-  std::pair<const GlobalValue *, unsigned>
-  tryGetCalledGlobal(const MachineInstr *MI) const {
-    return CalledGlobalsMap.lookup(MI);
+  CalledGlobalInfo tryGetCalledGlobal(const MachineInstr *MI) const {
+    return CalledGlobalsInfo.lookup(MI);
   }
 
   /// Notes the global and target flags for a call site.
-  void addCalledGlobal(const MachineInstr *MI,
-                       std::pair<const GlobalValue *, unsigned> Details) {
+  void addCalledGlobal(const MachineInstr *MI, CalledGlobalInfo Details) {
     assert(MI && "MI must not be null");
-    assert(Details.first && "Global must not be null");
-    CalledGlobalsMap.insert({MI, Details});
+    assert(Details.Callee && "Global must not be null");
+    CalledGlobalsInfo.insert({MI, Details});
   }
 
   /// Iterates over the full set of call sites and their associated globals.
   auto getCalledGlobals() const {
-    return llvm::make_range(CalledGlobalsMap.begin(), CalledGlobalsMap.end());
+    return llvm::make_range(CalledGlobalsInfo.begin(), CalledGlobalsInfo.end());
   }
 
   /// \name Exception Handling
@@ -1383,7 +1386,7 @@ class LLVM_ABI MachineFunction {
 
   /// Start tracking the arguments passed to the call \p CallI.
   void addCallSiteInfo(const MachineInstr *CallI, CallSiteInfo &&CallInfo) {
-    assert(CallI->isCandidateForCallSiteEntry());
+    assert(CallI->isCandidateForAdditionalCallInfo());
     bool Inserted =
         CallSitesInfo.try_emplace(CallI, std::move(CallInfo)).second;
     (void)Inserted;
@@ -1399,18 +1402,16 @@ class LLVM_ABI MachineFunction {
 
   /// 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);
+  void eraseAdditionalCallInfo(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);
+  void copyAdditionalCallInfo(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);
+  void moveAdditionalCallInfo(const MachineInstr *Old, const MachineInstr *New);
 
   unsigned getNewDebugInstrNum() {
     return ++DebugInstrNumberingCount;

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -957,13 +957,14 @@ class MachineInstr
     return hasProperty(MCID::Call, Type);
   }
 
-  /// Return true if this is a call instruction that may have an associated
-  /// call site entry in the debug info.
-  bool isCandidateForCallSiteEntry(QueryType Type = IgnoreBundle) const;
+  /// Return true if this is a call instruction that may have an additional
+  /// information associated with it.
+  bool isCandidateForAdditionalCallInfo(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;
+  /// updating additional call info (see \ref copyCallInfo, \ref moveCallInfo,
+  /// \ref eraseCallInfo).
+  bool shouldUpdateAdditionalCallInfo() const;
 
   /// Returns true if the specified instruction stops control flow
   /// from executing the instruction immediately following it.  Examples include

llvm/include/llvm/CodeGen/SelectionDAG.h

@@ -287,13 +287,14 @@ class SelectionDAG {
   SDDbgInfo *DbgInfo;
 
   using CallSiteInfo = MachineFunction::CallSiteInfo;
+  using CalledGlobalInfo = MachineFunction::CalledGlobalInfo;
 
   struct NodeExtraInfo {
     CallSiteInfo CSInfo;
     MDNode *HeapAllocSite = nullptr;
     MDNode *PCSections = nullptr;
     MDNode *MMRA = nullptr;
-    std::pair<const GlobalValue *, unsigned> CalledGlobal{};
+    CalledGlobalInfo CalledGlobal{};
     bool NoMerge = false;
   };
   /// Out-of-line extra information for SDNodes.
@@ -2380,8 +2381,7 @@ class SelectionDAG {
     SDEI[Node].CalledGlobal = {GV, OpFlags};
   }
   /// Return CalledGlobal associated with Node, or a nullopt if none exists.
-  std::optional<std::pair<const GlobalValue *, unsigned>>
-  getCalledGlobal(const SDNode *Node) {
+  std::optional<CalledGlobalInfo> getCalledGlobal(const SDNode *Node) {
     auto I = SDEI.find(Node);
     return I != SDEI.end()
                ? std::make_optional(std::move(I->second).CalledGlobal)

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp

@@ -918,7 +918,7 @@ void DwarfDebug::constructCallSiteEntryDIEs(const DISubprogram &SP,
 
       // Skip instructions which aren't calls. Both calls and tail-calling jump
       // instructions (e.g TAILJMPd64) are classified correctly here.
-      if (!MI.isCandidateForCallSiteEntry())
+      if (!MI.isCandidateForAdditionalCallInfo())
         continue;
 
       // Skip instructions marked as frame setup, as they are not interesting to
@@ -2019,7 +2019,7 @@ void DwarfDebug::beginInstruction(const MachineInstr *MI) {
 
   // When describing calls, we need a label for the call instruction.
   if (!NoDebug && SP->areAllCallsDescribed() &&
-      MI->isCandidateForCallSiteEntry(MachineInstr::AnyInBundle) &&
+      MI->isCandidateForAdditionalCallInfo(MachineInstr::AnyInBundle) &&
       (!MI->hasDelaySlot() || delaySlotSupported(*MI))) {
     const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
     bool IsTail = TII->isTailCall(*MI);

llvm/lib/CodeGen/BranchFolding.cpp

@@ -165,10 +165,10 @@ void BranchFolder::RemoveDeadBlock(MachineBasicBlock *MBB) {
   // Avoid matching if this pointer gets reused.
   TriedMerging.erase(MBB);
 
-  // Update call site info.
+  // Update call info.
   for (const MachineInstr &MI : *MBB)
-    if (MI.shouldUpdateCallSiteInfo())
-      MF->eraseCallSiteInfo(&MI);
+    if (MI.shouldUpdateAdditionalCallInfo())
+      MF->eraseAdditionalCallInfo(&MI);
 
   // Remove the block.
   MF->erase(MBB);

llvm/lib/CodeGen/IfConversion.cpp

@@ -1834,9 +1834,9 @@ 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->shouldUpdateCallSiteInfo())
-      MBB2.getParent()->eraseCallSiteInfo(&*DI2);
+    // info state if the instruction is call instruction.
+    if (DI2->shouldUpdateAdditionalCallInfo())
+      MBB2.getParent()->eraseAdditionalCallInfo(&*DI2);
 
     ++DI2;
     if (DI2 == MBB2.end())
@@ -1883,9 +1883,9 @@ bool IfConverter::IfConvertDiamondCommon(
     --DI1;
 
     // Since this instruction is going to be deleted, update call
-    // site info state if the instruction is call instruction.
-    if (DI1->shouldUpdateCallSiteInfo())
-      MBB1.getParent()->eraseCallSiteInfo(&*DI1);
+    // info state if the instruction is call instruction.
+    if (DI1->shouldUpdateAdditionalCallInfo())
+      MBB1.getParent()->eraseAdditionalCallInfo(&*DI1);
 
     // skip dbg_value instructions
     if (!DI1->isDebugInstr())
@@ -2169,9 +2169,9 @@ void IfConverter::CopyAndPredicateBlock(BBInfo &ToBBI, BBInfo &FromBBI,
       break;
 
     MachineInstr *MI = MF.CloneMachineInstr(&I);
-    // Make a copy of the call site info.
-    if (I.isCandidateForCallSiteEntry())
-      MF.copyCallSiteInfo(&I, MI);
+    // Make a copy of the call info.
+    if (I.isCandidateForAdditionalCallInfo())
+      MF.copyAdditionalCallInfo(&I, MI);
 
     ToBBI.BB->insert(ToBBI.BB->end(), MI);
     ToBBI.NonPredSize++;

llvm/lib/CodeGen/InlineSpiller.cpp

@@ -997,9 +997,9 @@ foldMemoryOperand(ArrayRef<std::pair<MachineInstr *, unsigned>> Ops,
       HSpiller.rmFromMergeableSpills(*MI, FI))
     --NumSpills;
   LIS.ReplaceMachineInstrInMaps(*MI, *FoldMI);
-  // Update the call site info.
-  if (MI->isCandidateForCallSiteEntry())
-    MI->getMF()->moveCallSiteInfo(MI, FoldMI);
+  // Update the call info.
+  if (MI->isCandidateForAdditionalCallInfo())
+    MI->getMF()->moveAdditionalCallInfo(MI, FoldMI);
 
   // If we've folded a store into an instruction labelled with debug-info,
   // record a substitution from the old operand to the memory operand. Handle

llvm/lib/CodeGen/LiveRangeEdit.cpp

@@ -253,9 +253,9 @@ bool LiveRangeEdit::foldAsLoad(LiveInterval *LI,
     return false;
   LLVM_DEBUG(dbgs() << "                folded: " << *FoldMI);
   LIS.ReplaceMachineInstrInMaps(*UseMI, *FoldMI);
-  // Update the call site info.
-  if (UseMI->shouldUpdateCallSiteInfo())
-    UseMI->getMF()->moveCallSiteInfo(UseMI, FoldMI);
+  // Update the call info.
+  if (UseMI->shouldUpdateAdditionalCallInfo())
+    UseMI->getMF()->moveAdditionalCallInfo(UseMI, FoldMI);
   UseMI->eraseFromParent();
   DefMI->addRegisterDead(LI->reg(), nullptr);
   Dead.push_back(DefMI);

llvm/lib/CodeGen/MIRPrinter.cpp

@@ -605,17 +605,14 @@ void MIRPrinter::convertCalledGlobals(yaml::MachineFunction &YMF,
                                       const MachineFunction &MF,
                                       MachineModuleSlotTracker &MST) {
   for (const auto &[CallInst, CG] : MF.getCalledGlobals()) {
-    // If the call instruction was dropped, then we don't need to print it.
-    auto BB = CallInst->getParent();
-    if (BB) {
-      yaml::MachineInstrLoc CallSite;
-      CallSite.BlockNum = CallInst->getParent()->getNumber();
-      CallSite.Offset = std::distance(CallInst->getParent()->instr_begin(),
-                                      CallInst->getIterator());
-
-      yaml::CalledGlobal YamlCG{CallSite, CG.first->getName().str(), CG.second};
-      YMF.CalledGlobals.push_back(YamlCG);
-    }
+    yaml::MachineInstrLoc CallSite;
+    CallSite.BlockNum = CallInst->getParent()->getNumber();
+    CallSite.Offset = std::distance(CallInst->getParent()->instr_begin(),
+                                    CallInst->getIterator());
+
+    yaml::CalledGlobal YamlCG{CallSite, CG.Callee->getName().str(),
+                              CG.TargetFlags};
+    YMF.CalledGlobals.push_back(YamlCG);
   }
 
   // Sort by position of call instructions.