[LoongArch] Handle register spill in BranchRelaxation pass

When the range of the unconditional branch is overflow, the indirect
branch way is used. The case when there is no scavenged register for
indirect branch needs to spill register to stack.

Reviewed By: SixWeining, wangleiat

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

llvm/lib/Target/LoongArch/LoongArchFrameLowering.cpp

@@ -118,12 +118,32 @@ void LoongArchFrameLowering::determineFrameLayout(MachineFunction &MF) const {
   MFI.setStackSize(FrameSize);
 }
 
+static uint64_t estimateFunctionSizeInBytes(const LoongArchInstrInfo *TII,
+                                            const MachineFunction &MF) {
+  uint64_t FuncSize = 0;
+  for (auto &MBB : MF)
+    for (auto &MI : MBB)
+      FuncSize += TII->getInstSizeInBytes(MI);
+  return FuncSize;
+}
+
 void LoongArchFrameLowering::processFunctionBeforeFrameFinalized(
     MachineFunction &MF, RegScavenger *RS) const {
   const LoongArchRegisterInfo *RI = STI.getRegisterInfo();
   const TargetRegisterClass &RC = LoongArch::GPRRegClass;
+  const LoongArchInstrInfo *TII = STI.getInstrInfo();
+  LoongArchMachineFunctionInfo *LAFI =
+      MF.getInfo<LoongArchMachineFunctionInfo>();
   MachineFrameInfo &MFI = MF.getFrameInfo();
 
+  // Far branches beyond 27-bit offset require a spill slot for scratch register.
+  if (!isInt<27>(estimateFunctionSizeInBytes(TII, MF))) {
+    int FI = MFI.CreateStackObject(RI->getSpillSize(RC), RI->getSpillAlign(RC),
+                                   false);
+    RS->addScavengingFrameIndex(FI);
+    if (LAFI->getBranchRelaxationSpillFrameIndex() == -1)
+      LAFI->setBranchRelaxationSpillFrameIndex(FI);
+  }
   // estimateStackSize has been observed to under-estimate the final stack
   // size, so give ourselves wiggle-room by checking for stack size
   // representable an 11-bit signed field rather than 12-bits.

llvm/lib/Target/LoongArch/LoongArchInstrInfo.cpp

@@ -371,6 +371,9 @@ void LoongArchInstrInfo::insertIndirectBranch(MachineBasicBlock &MBB,
 
   MachineFunction *MF = MBB.getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();
+  const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
+  LoongArchMachineFunctionInfo *LAFI =
+      MF->getInfo<LoongArchMachineFunctionInfo>();
 
   if (!isInt<32>(BrOffset))
     report_fatal_error(
@@ -379,26 +382,45 @@ void LoongArchInstrInfo::insertIndirectBranch(MachineBasicBlock &MBB,
   Register ScratchReg = MRI.createVirtualRegister(&LoongArch::GPRRegClass);
   auto II = MBB.end();
 
-  MachineInstr &MI =
+  MachineInstr &PCALAU12I =
       *BuildMI(MBB, II, DL, get(LoongArch::PCALAU12I), ScratchReg)
            .addMBB(&DestBB, LoongArchII::MO_PCREL_HI);
-  BuildMI(MBB, II, DL,
-          get(STI.is64Bit() ? LoongArch::ADDI_D : LoongArch::ADDI_W),
-          ScratchReg)
-      .addReg(ScratchReg)
-      .addMBB(&DestBB, LoongArchII::MO_PCREL_LO);
+  MachineInstr &ADDI =
+      *BuildMI(MBB, II, DL,
+               get(STI.is64Bit() ? LoongArch::ADDI_D : LoongArch::ADDI_W),
+               ScratchReg)
+           .addReg(ScratchReg)
+           .addMBB(&DestBB, LoongArchII::MO_PCREL_LO);
   BuildMI(MBB, II, DL, get(LoongArch::PseudoBRIND))
       .addReg(ScratchReg, RegState::Kill)
       .addImm(0);
 
   RS->enterBasicBlockEnd(MBB);
-  Register Scav = RS->scavengeRegisterBackwards(LoongArch::GPRRegClass,
-                                                MI.getIterator(), false, 0);
-  // TODO: When there is no scavenged register, it needs to specify a register.
-  assert(Scav != LoongArch::NoRegister && "No register is scavenged!");
+  Register Scav = RS->scavengeRegisterBackwards(
+      LoongArch::GPRRegClass, PCALAU12I.getIterator(), /*RestoreAfter=*/false,
+      /*SPAdj=*/0, /*AllowSpill=*/false);
+  if (Scav != LoongArch::NoRegister)
+    RS->setRegUsed(Scav);
+  else {
+    // When there is no scavenged register, it needs to specify a register.
+    // Specify t8 register because it won't be used too often.
+    Scav = LoongArch::R20;
+    int FrameIndex = LAFI->getBranchRelaxationSpillFrameIndex();
+    if (FrameIndex == -1)
+      report_fatal_error("The function size is incorrectly estimated.");
+    storeRegToStackSlot(MBB, PCALAU12I, Scav, /*IsKill=*/true, FrameIndex,
+                        &LoongArch::GPRRegClass, TRI);
+    TRI->eliminateFrameIndex(std::prev(PCALAU12I.getIterator()),
+                             /*SpAdj=*/0, /*FIOperandNum=*/1);
+    PCALAU12I.getOperand(1).setMBB(&RestoreBB);
+    ADDI.getOperand(2).setMBB(&RestoreBB);
+    loadRegFromStackSlot(RestoreBB, RestoreBB.end(), Scav, FrameIndex,
+                         &LoongArch::GPRRegClass, TRI);
+    TRI->eliminateFrameIndex(RestoreBB.back(),
+                             /*SpAdj=*/0, /*FIOperandNum=*/1);
+  }
   MRI.replaceRegWith(ScratchReg, Scav);
   MRI.clearVirtRegs();
-  RS->setRegUsed(Scav);
 }
 
 static unsigned getOppositeBranchOpc(unsigned Opc) {

llvm/lib/Target/LoongArch/LoongArchMachineFunctionInfo.h

@@ -32,6 +32,10 @@ class LoongArchMachineFunctionInfo : public MachineFunctionInfo {
   /// Size of stack frame to save callee saved registers
   unsigned CalleeSavedStackSize = 0;
 
+  /// FrameIndex of the spill slot when there is no scavenged register in
+  /// insertIndirectBranch.
+  int BranchRelaxationSpillFrameIndex = -1;
+
 public:
   LoongArchMachineFunctionInfo(const MachineFunction &MF) {}
 
@@ -50,6 +54,13 @@ class LoongArchMachineFunctionInfo : public MachineFunctionInfo {
 
   unsigned getCalleeSavedStackSize() const { return CalleeSavedStackSize; }
   void setCalleeSavedStackSize(unsigned Size) { CalleeSavedStackSize = Size; }
+
+  int getBranchRelaxationSpillFrameIndex() {
+    return BranchRelaxationSpillFrameIndex;
+  }
+  void setBranchRelaxationSpillFrameIndex(int Index) {
+    BranchRelaxationSpillFrameIndex = Index;
+  }
 };
 
 } // end namespace llvm