[llvm] Rebased save csr in ra (PR #131845)

[Michael Maitland via llvm-commits]

================
@@ -786,6 +788,153 @@ void RISCVFrameLowering::allocateStack(MachineBasicBlock &MBB,
   }
 }
 
+struct CFIBuildInfo {
+  MachineBasicBlock *MBB;
+  MachineInstr *InsertAfterMI; // nullptr means insert at MBB.begin()
+  DebugLoc DL;
+  unsigned CFIIndex;
+};
+
+static void trackRegisterAndEmitCFIs(
+    MachineFunction &MF, MachineInstr &MI, MCRegister Reg, int DwarfEHRegNum,
+    const ReachingDefAnalysis &RDA, const TargetInstrInfo &TII,
+    const MachineFrameInfo &MFI, const RISCVRegisterInfo &TRI,
+    std::vector<CFIBuildInfo> &CFIBuildInfos,
+    std::unordered_set<MachineInstr *> &VisitedRestorePoints,
+    std::unordered_set<MachineInstr *> &VisitedDefs) {
+
+  if (VisitedRestorePoints.find(&MI) != VisitedRestorePoints.end()) {
+    return;
+  }
+  VisitedRestorePoints.insert(&MI);
+  SmallPtrSet<MachineInstr *, 2> Defs;
+  RDA.getGlobalReachingDefs(&MI, Reg, Defs);
+  MachineBasicBlock &EntryMBB = MF.front();
+  if (Defs.empty()) {
+    // it's a live-in register at the entry block.
+    // unsigned CFIIndex =
+    // MF.addFrameInst(MCCFIInstruction::createSameValue(nullptr,
+    // DwarfEHRegNum)); CFIBuildInfos.push_back({&EntryMBB, nullptr, DebugLoc(),
+    // CFIIndex});
+    return;
+  }
+
+  int FrameIndex = std::numeric_limits<int>::min();
+  for (MachineInstr *Def : Defs) {
+    if (VisitedDefs.find(Def) != VisitedDefs.end())
+      continue;
+    VisitedDefs.insert(Def);
+
+    MachineBasicBlock &MBB = *Def->getParent();
+    const DebugLoc &DL = Def->getDebugLoc();
+
+    if (Register StoredReg = TII.isStoreToStackSlot(*Def, FrameIndex)) {
+      assert(FrameIndex == Register::stackSlot2Index(Reg));
+
+      Register FrameReg;
+      StackOffset Offset =
+          MF.getSubtarget().getFrameLowering()->getFrameIndexReference(
+              MF, FrameIndex, FrameReg);
+      int64_t FixedOffset = Offset.getFixed();
+      // TODO:
+      assert(Offset.getScalable() == 0);
+
+      std::string CommentBuffer;
+      llvm::raw_string_ostream Comment(CommentBuffer);
+      int DwarfEHFrameReg = TRI.getDwarfRegNum(FrameReg, true);
+      Register LLVMReg = *TRI.getLLVMRegNum(DwarfEHRegNum, true);
+      Comment << printReg(LLVMReg, &TRI) << " = *(";
+      Comment << printReg(FrameReg, &TRI) << " + ";
+      Comment << FixedOffset << ")";
+      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createLLVMRegOffset(
+          nullptr, DwarfEHRegNum, DwarfEHFrameReg, FixedOffset, SMLoc(), Comment.str()));
+
+      CFIBuildInfos.push_back({&MBB, Def, DL, CFIIndex});
+      trackRegisterAndEmitCFIs(MF, *Def, StoredReg, DwarfEHRegNum, RDA, TII,
+                               MFI, TRI, CFIBuildInfos, VisitedRestorePoints,
+                               VisitedDefs);
+    } else if (Register LoadedReg = TII.isLoadFromStackSlot(*Def, FrameIndex)) {
+      assert(LoadedReg == Reg);
+
+      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createRegister(
+          nullptr, DwarfEHRegNum, TRI.getDwarfRegNum(LoadedReg, true)));
+      CFIBuildInfos.push_back({&MBB, Def, DL, CFIIndex});
+      trackRegisterAndEmitCFIs(MF, *Def, Register::index2StackSlot(FrameIndex),
+                               DwarfEHRegNum, RDA, TII, MFI, TRI, CFIBuildInfos,
+                               VisitedRestorePoints, VisitedDefs);
+    } else if (auto DstSrc = TII.isCopyInstr(*Def)) {
+      Register DstReg = DstSrc->Destination->getReg();
+      Register SrcReg = DstSrc->Source->getReg();
+      assert(DstReg == Reg);
+
+      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createRegister(
+          nullptr, DwarfEHRegNum, TRI.getDwarfRegNum(DstReg, true)));
+      CFIBuildInfos.push_back({&MBB, Def, DL, CFIIndex});
+      trackRegisterAndEmitCFIs(MF, *Def, SrcReg, DwarfEHRegNum, RDA, TII, MFI,
+                               TRI, CFIBuildInfos, VisitedRestorePoints,
+                               VisitedDefs);
+    } else {
+      llvm_unreachable("Unexpected instruction");
+    }
+  }
+  return;
+}
+
+int RISCVFrameLowering::getInitialCFAOffset(const MachineFunction &MF) const {
+  return 0;
+}
+
+Register
+RISCVFrameLowering::getInitialCFARegister(const MachineFunction &MF) const {
+  return RISCV::X2;
+}
+
+void RISCVFrameLowering::emitCFIsForCSRsHandledByRA(
+    MachineFunction &MF, ReachingDefAnalysis *RDA) const {
+  if (!STI.doCSRSavesInRA())
+    return;
+  const RISCVInstrInfo &TII = *STI.getInstrInfo();
+  const RISCVRegisterInfo &TRI = *STI.getRegisterInfo();
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  BitVector MustCalleeSavedRegs;
+  determineMustCalleeSaves(MF, MustCalleeSavedRegs);
+  const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs();
+  SmallVector<MCPhysReg, 4> EligibleRegs;
+  for (int i = 0; CSRegs[i]; ++i) {
+    unsigned Reg = CSRegs[i];
+    if (!MustCalleeSavedRegs.test(Reg))
+      EligibleRegs.push_back(CSRegs[i]);
+  }
+
+  SmallVector<MachineInstr *, 4> RestorePoints;
+  for (MachineBasicBlock &MBB : MF) {
+    if (MBB.isReturnBlock())
+      RestorePoints.push_back(&MBB.back());
+  }
+  // TODO: replace CFIBuildInfo with UnwindLocation from DebugInfo/DWARF/DWARFDebugFrame
+  std::vector<CFIBuildInfo> CFIBuildInfos;
+  for (MCPhysReg Reg : EligibleRegs) {
+    std::unordered_set<MachineInstr *> VisitedDefs;
+    for (MachineInstr *RestorePoint : RestorePoints) {
+      std::unordered_set<MachineInstr *> VisitedRestorePoints;
+      trackRegisterAndEmitCFIs(
+          MF, *RestorePoint, Reg, TRI.getDwarfRegNum(Reg, true), *RDA, TII, MFI,
+          TRI, CFIBuildInfos, VisitedRestorePoints, VisitedDefs);
+    }
+  }
+  for (CFIBuildInfo &Info : CFIBuildInfos) {
+    MachineBasicBlock::iterator InsertPos =
+        Info.InsertAfterMI ? ++(Info.InsertAfterMI->getIterator())
+                           : Info.MBB->begin();
+    BuildMI(*Info.MBB, InsertPos, Info.DL,
----------------
michaelmaitland wrote:

It looks like we are emitting CFI instruction here. Why did this not occur in `trackRegisterandEmitCFIs`?

https://github.com/llvm/llvm-project/pull/131845