1 files changed, 85 insertions, 1 deletions

diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
index 7d48634f206..6b6f62e18ce 100644
--- a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
@@ -131,6 +131,11 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
   if (StackSize == 0 && !MFI.adjustsStack())
     return;
 
+  uint64_t FirstSPAdjustAmount = getFirstSPAdjustAmount(MF);
+  // Split the SP adjustment to reduce the offsets of callee saved spill.
+  if (FirstSPAdjustAmount)
+    StackSize = FirstSPAdjustAmount;
+
   // Allocate space on the stack if necessary.
   adjustReg(MBB, MBBI, DL, SPReg, SPReg, -StackSize, MachineInstr::FrameSetup);
 
@@ -170,7 +175,23 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
         nullptr, RI->getDwarfRegNum(FPReg, true), 0));
     BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
         .addCFIIndex(CFIIndex);
+  }
 
+  // Emit the second SP adjustment after saving callee saved registers.
+  if (FirstSPAdjustAmount) {
+    uint64_t SecondSPAdjustAmount = MFI.getStackSize() - FirstSPAdjustAmount;
+    assert(SecondSPAdjustAmount > 0 &&
+           "SecondSPAdjustAmount should be greater than zero");
+    adjustReg(MBB, MBBI, DL, SPReg, SPReg, -SecondSPAdjustAmount,
+              MachineInstr::FrameSetup);
+    // Emit ".cfi_def_cfa_offset StackSize"
+    unsigned CFIIndex = MF.addFrameInst(
+        MCCFIInstruction::createDefCfaOffset(nullptr, -MFI.getStackSize()));
+    BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+        .addCFIIndex(CFIIndex);
+  }
+
+  if (hasFP(MF)) {
     // Realign Stack
     const RISCVRegisterInfo *RI = STI.getRegisterInfo();
     if (RI->needsStackRealignment(MF)) {
@@ -224,6 +245,24 @@ void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
               MachineInstr::FrameDestroy);
   }
 
+  uint64_t FirstSPAdjustAmount = getFirstSPAdjustAmount(MF);
+  if (FirstSPAdjustAmount) {
+    uint64_t SecondSPAdjustAmount = MFI.getStackSize() - FirstSPAdjustAmount;
+    assert(SecondSPAdjustAmount > 0 &&
+           "SecondSPAdjustAmount should be greater than zero");
+
+    adjustReg(MBB, LastFrameDestroy, DL, SPReg, SPReg, SecondSPAdjustAmount,
+              MachineInstr::FrameDestroy);
+
+    // Emit ".cfi_def_cfa_offset FirstSPAdjustAmount"
+    unsigned CFIIndex =
+        MF.addFrameInst(
+             MCCFIInstruction::createDefCfaOffset(nullptr,
+                                                  -FirstSPAdjustAmount));
+    BuildMI(MBB, LastFrameDestroy, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+        .addCFIIndex(CFIIndex);
+  }
+
   if (hasFP(MF)) {
     // To find the instruction restoring FP from stack.
     for (auto &I = LastFrameDestroy; I != MBBI; ++I) {
@@ -256,6 +295,9 @@ void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
         .addCFIIndex(CFIIndex);
   }
 
+  if (FirstSPAdjustAmount)
+    StackSize = FirstSPAdjustAmount;
+
   // Deallocate stack
   adjustReg(MBB, MBBI, DL, SPReg, SPReg, StackSize, MachineInstr::FrameDestroy);
 
@@ -284,6 +326,8 @@ int RISCVFrameLowering::getFrameIndexReference(const MachineFunction &MF,
   int Offset = MFI.getObjectOffset(FI) - getOffsetOfLocalArea() +
                MFI.getOffsetAdjustment();
 
+  uint64_t FirstSPAdjustAmount = getFirstSPAdjustAmount(MF);
+
   if (CSI.size()) {
     MinCSFI = CSI[0].getFrameIdx();
     MaxCSFI = CSI[CSI.size() - 1].getFrameIdx();
@@ -291,7 +335,11 @@ int RISCVFrameLowering::getFrameIndexReference(const MachineFunction &MF,
 
   if (FI >= MinCSFI && FI <= MaxCSFI) {
     FrameReg = RISCV::X2;
-    Offset += MF.getFrameInfo().getStackSize();
+
+    if (FirstSPAdjustAmount)
+      Offset += FirstSPAdjustAmount;
+    else
+      Offset += MF.getFrameInfo().getStackSize();
   } else if (RI->needsStackRealignment(MF)) {
     assert(!MFI.hasVarSizedObjects() &&
            "Unexpected combination of stack realignment and varsized objects");
@@ -404,3 +452,39 @@ MachineBasicBlock::iterator RISCVFrameLowering::eliminateCallFramePseudoInstr(
 
   return MBB.erase(MI);
 }
+
+// We would like to split the SP adjustment to reduce prologue/epilogue
+// as following instructions. In this way, the offset of the callee saved
+// register could fit in a single store.
+//   add     sp,sp,-2032
+//   sw      ra,2028(sp)
+//   sw      s0,2024(sp)
+//   sw      s1,2020(sp)
+//   sw      s3,2012(sp)
+//   sw      s4,2008(sp)
+//   add     sp,sp,-64
+uint64_t
+RISCVFrameLowering::getFirstSPAdjustAmount(const MachineFunction &MF) const {
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  const std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
+  uint64_t StackSize = MFI.getStackSize();
+  uint64_t StackAlign = getStackAlignment();
+
+  // FIXME: Disable SplitSPAdjust if save-restore libcall enabled when the patch
+  //        landing. The callee saved registers will be pushed by the
+  //        save-restore libcalls, so we don't have to split the SP adjustment
+  //        in this case.
+  //
+  // Return the FirstSPAdjustAmount if the StackSize can not fit in signed
+  // 12-bit and there exists a callee saved register need to be pushed.
+  if (!isInt<12>(StackSize) && (CSI.size() > 0)) {
+    // FirstSPAdjustAmount is choosed as (2048 - StackAlign)
+    // because 2048 will cause sp = sp + 2048 in epilogue split into
+    // multi-instructions. The offset smaller than 2048 can fit in signle
+    // load/store instruction and we have to stick with the stack alignment.
+    // 2048 is 16-byte alignment. The stack alignment for RV32 and RV64 is 16,
+    // for RV32E is 4. So (2048 - StackAlign) will satisfy the stack alignment.
+    return 2048 - StackAlign;
+  }
+  return 0;
+}