[MIR] Add call-site-info support for stack args

[milica-lazarevic]

Summary

Attaching information about arguments forwarding via stack to the MachineFunction, for every MachineInstruction that is candidate for CallSiteInfo.

llvm/include/llvm/CodeGen/MachineFunction.h

struct ArgRegPair {
     ArgRegPair(Register R, unsigned Arg) : Reg(R), ArgNo(Arg) {
       assert(Arg < (1 << 16) && "Arg out of range");
     }
   };
   /// Vector of call argument and its forwarding register.
   using CallSiteInfo = SmallVector<ArgRegPair, 1>;
   using CallSiteInfoImpl = SmallVectorImpl<ArgRegPair>;
 
+  struct ArgStackOffsetPair {
+    uint64_t Offset;
+    uint16_t ArgNo;
+    ArgStackOffsetPair(unsigned O, unsigned Arg) : Offset(O), ArgNo(Arg) {
+      assert(Arg < (1 << 16) && "Arg out of range");
+    }
+  };
+  /// Vector of call argument and its forwarding stack offset.
+  using CallSiteInfoStackFwd = SmallVector<ArgStackOffsetPair, 1>;
+  using CallSiteInfoStackFwdImpl = SmallVectorImpl<ArgStackOffsetPair>;
+
 private:
   Delegate *TheDelegate = nullptr;
   GISelChangeObserver *Observer = nullptr;
 
   using CallSiteInfoMap = DenseMap<const MachineInstr *, CallSiteInfo>;
   /// Map a call instruction to call site arguments forwarding info.
   CallSiteInfoMap CallSitesInfo;
 
+  using CallSiteInfoStackFwdMap =
+      DenseMap<const MachineInstr *, CallSiteInfoStackFwd>;
+
+  /// Map a call instruction to call site arguments stack forwarding info.
+  CallSiteInfoStackFwdMap CallSitesInfoStackFwd;
+
   /// A helper function that returns call site info for a give call
   /// instruction if debug entry value support is enabled.
   CallSiteInfoMap::iterator getCallSiteInfo(const MachineInstr *MI);
 
   // Callbacks for insertion and removal.
   void handleInsertion(MachineInstr &MI);
   void handleRemoval(MachineInstr &MI);
   friend struct ilist_traits<MachineInstr>;
void addCallArgsForwardingRegs(const MachineInstr *CallI,
     (void)Inserted;
     assert(Inserted && "Call site info not unique");
   }
 
   const CallSiteInfoMap &getCallSitesInfo() const {
     return CallSitesInfo;
   }
 
+  /// Start tracking the arguments passed on the stack to the call \p CallI.
+  void addCallArgsForwardingStackOffsets(const MachineInstr *CallI,
+                                         CallSiteInfoStackFwdImpl &&CallInfo) {
+    assert(CallI->isCandidateForCallSiteEntry());
+    bool Inserted =
+        CallSitesInfoStackFwd.try_emplace(CallI, std::move(CallInfo)).second;
+    (void)Inserted;
+    assert(Inserted && "Call site info not unique");
+  }
+
+  const CallSiteInfoStackFwdMap &getCallSitesInfoStackFwd() const {
+    return CallSitesInfoStackFwd;
+  }
+
   /// Following functions update call site info. They should be called before
   /// removing, replacing or copying call instruction.
 
   /// 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

llvm/include/llvm/CodeGen/SelectionDAG.h

class SelectionDAG {
   BumpPtrAllocator Allocator;
 
   /// Tracks dbg_value and dbg_label information through SDISel.
   SDDbgInfo *DbgInfo;
 
   using CallSiteInfo = MachineFunction::CallSiteInfo;
   using CallSiteInfoImpl = MachineFunction::CallSiteInfoImpl;
 
+  using CallSiteInfoStackFwd = MachineFunction::CallSiteInfoStackFwd;
+  using CallSiteInfoStackFwdImpl = MachineFunction::CallSiteInfoStackFwdImpl;
+
+  DenseMap<const SDNode *, CallSiteInfoStackFwd> SDCallSiteInfoStackFwd;
+
   struct CallSiteDbgInfo {
     CallSiteInfo CSInfo;
     MDNode *HeapAllocSite = nullptr;
     bool NoMerge = false;
   };
 
   DenseMap<const SDNode *, CallSiteDbgInfo> SDCallSiteDbgInfo;
 
public:
 
   CallSiteInfo getSDCallSiteInfo(const SDNode *CallNode) {
     auto I = SDCallSiteDbgInfo.find(CallNode);
     if (I != SDCallSiteDbgInfo.end())
       return std::move(I->second).CSInfo;
     return CallSiteInfo();
   }
 
+  void addCallSiteInfoStackFwd(const SDNode *CallNode,
+                               CallSiteInfoStackFwdImpl &&CallInfo) {
+    SDCallSiteInfoStackFwd[CallNode] = std::move(CallInfo);
+  }
+
+  CallSiteInfoStackFwd getSDCallSiteInfoStackFwd(const SDNode *CallNode) {
+    auto I = SDCallSiteInfoStackFwd.find(CallNode);
+    if (I != SDCallSiteInfoStackFwd.end())
+      return std::move(I->second);
+    return CallSiteInfoStackFwd();
+  }
+
   void addHeapAllocSite(const SDNode *Node, MDNode *MD) {
     SDCallSiteDbgInfo[Node].HeapAllocSite = MD;
   }
 
   /// Return the HeapAllocSite type associated with the SDNode, if it exists.
   MDNode *getHeapAllocSite(const SDNode *Node) {
     auto It = SDCallSiteDbgInfo.find(Node);
     if (It == SDCallSiteDbgInfo.end())

llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp

if (Before == BB->end()) {
       // beginning of the block.
       MI = &Emitter.getBlock()->instr_front();
     } else {
       // Return first instruction after the pre-existing instructions.
       MI = &*std::next(Before);
     }
 
     if (MI->isCandidateForCallSiteEntry() &&
-        DAG->getTarget().Options.EmitCallSiteInfo)
+        DAG->getTarget().Options.EmitCallSiteInfo) {
       MF.addCallArgsForwardingRegs(MI, DAG->getSDCallSiteInfo(Node));
+      MF.addCallArgsForwardingStackOffsets(
+          MI, DAG->getSDCallSiteInfoStackFwd(Node));
+    }
 
     if (DAG->getNoMergeSiteInfo(Node)) {
       MI->setFlag(MachineInstr::MIFlag::NoMerge);
     }
 
     return MI;
   };