[LoongArch] Split SP adjustment

[wangleiat]

Summary

This patch split the SP adjustment to reduce the instructions in
prologue and epilogue. In this way, the offset of the callee saved
register could fit in a single store.

Similar to D68011(RISCV).

Depends on D136074

[SixWeining]

LGTM.

llvm/lib/Target/LoongArch/LoongArchFrameLowering.cpp

303

Nit: Inline one-time-used variable.

llvm/test/CodeGen/LoongArch/stack-realignment.ll

557

With the addition of this patch, 7 instructions are reduced for this function.

[Allen]

llvm/test/CodeGen/LoongArch/split-sp-adjust.ll

6

Sorry for the naive question: do you split the offset of SP as the constraint of the load/store, which should be less than 2048 ?

[wangleiat]

Address @SixWeining's comments.

[wangleiat]

llvm/lib/Target/LoongArch/LoongArchFrameLowering.cpp

303

Nit: Inline one-time-used variable.

Ok, I will fix it.

llvm/test/CodeGen/LoongArch/split-sp-adjust.ll

6

Thank you very much for your interest in LoongArch.
As you guessed, if we don't split the adjustment of SP, the load/store for callee-saved registers will use multiple instructions when offset >= 2048.

[wangleiat]

Rebase and added some spaces.

llvm/lib/Target/LoongArch/LoongArchFrameLowering.h

eliminateCallFramePseudoInstr(MachineFunction &MF, MachineBasicBlock &MBB,
                                 MachineBasicBlock::iterator MI) const override;
 
   StackOffset getFrameIndexReference(const MachineFunction &MF, int FI,
                                      Register &FrameReg) const override;
 
   bool hasFP(const MachineFunction &MF) const override;
   bool hasBP(const MachineFunction &MF) const;
 
+  uint64_t getFirstSPAdjustAmount(const MachineFunction &MF) const;
+
 private:
   void determineFrameLayout(MachineFunction &MF) const;
   void adjustReg(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                  const DebugLoc &DL, Register DestReg, Register SrcReg,
                  int64_t Val, MachineInstr::MIFlag Flag) const;
 };
 } // end namespace llvm
 #endif // LLVM_LIB_TARGET_LOONGARCH_LOONGARCHFRAMELOWERING_H

llvm/lib/Target/LoongArch/LoongArchFrameLowering.cpp

void LoongArchFrameLowering::emitPrologue(MachineFunction &MF,
   // to determine the end of the prologue.
   DebugLoc DL;
 
   // Determine the correct frame layout
   determineFrameLayout(MF);
 
   // First, compute final stack size.
   uint64_t StackSize = MFI.getStackSize();
+  uint64_t RealStackSize = StackSize;
 
   // Early exit if there is no need to allocate space in the stack.
   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;
+
   // Adjust stack.
   adjustReg(MBB, MBBI, DL, SPReg, SPReg, -StackSize, MachineInstr::FrameSetup);
   // Emit ".cfi_def_cfa_offset StackSize".
   unsigned CFIIndex =
       MF.addFrameInst(MCCFIInstruction::cfiDefCfaOffset(nullptr, StackSize));
   BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
       .addCFIIndex(CFIIndex)
       .setMIFlag(MachineInstr::FrameSetup);
if (hasFP(MF)) {
 
     // Emit ".cfi_def_cfa $fp, LoongArchFI->getVarArgsSaveSize()"
     unsigned CFIIndex = MF.addFrameInst(
         MCCFIInstruction::cfiDefCfa(nullptr, RI->getDwarfRegNum(FPReg, true),
                                     LoongArchFI->getVarArgsSaveSize()));
     BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
         .addCFIIndex(CFIIndex)
         .setMIFlag(MachineInstr::FrameSetup);
+  }
+
+  // Emit the second SP adjustment after saving callee saved registers.
+  if (FirstSPAdjustAmount) {
+    uint64_t SecondSPAdjustAmount = RealStackSize - FirstSPAdjustAmount;
+    assert(SecondSPAdjustAmount > 0 &&
+           "SecondSPAdjustAmount should be greater than zero");
+    adjustReg(MBB, MBBI, DL, SPReg, SPReg, -SecondSPAdjustAmount,
+              MachineInstr::FrameSetup);
+
+    if (!hasFP(MF)) {
+      // If we are using a frame-pointer, and thus emitted ".cfi_def_cfa fp, 0",
+      // don't emit an sp-based .cfi_def_cfa_offset
+      // Emit ".cfi_def_cfa_offset RealStackSize"
+      unsigned CFIIndex = MF.addFrameInst(
+          MCCFIInstruction::cfiDefCfaOffset(nullptr, RealStackSize));
+      BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex)
+          .setMIFlag(MachineInstr::FrameSetup);
+    }
+  }
 
+  if (hasFP(MF)) {
     // Realign stack.
     if (RI->hasStackRealignment(MF)) {
       Align MaxAlignment = MFI.getMaxAlign();
       unsigned ShiftAmount = Log2(MaxAlignment);
       Register VR =
           MF.getRegInfo().createVirtualRegister(&LoongArch::GPRRegClass);
       BuildMI(MBB, MBBI, DL,
               TII->get(STI.is64Bit() ? LoongArch::SRLI_D : LoongArch::SRLI_W),
void LoongArchFrameLowering::emitEpilogue(MachineFunction &MF,
   // Restore the stack pointer.
   if (RI->hasStackRealignment(MF) || MFI.hasVarSizedObjects()) {
     assert(hasFP(MF) && "frame pointer should not have been eliminated");
     adjustReg(MBB, LastFrameDestroy, DL, SPReg, LoongArch::R22,
               -StackSize + LoongArchFI->getVarArgsSaveSize(),
               MachineInstr::FrameDestroy);
   }
 
+  uint64_t FirstSPAdjustAmount = getFirstSPAdjustAmount(MF);
+  if (FirstSPAdjustAmount) {
+    uint64_t SecondSPAdjustAmount = StackSize - FirstSPAdjustAmount;
+    assert(SecondSPAdjustAmount > 0 &&
+           "SecondSPAdjustAmount should be greater than zero");
+
+    adjustReg(MBB, LastFrameDestroy, DL, SPReg, SPReg, SecondSPAdjustAmount,
+              MachineInstr::FrameDestroy);
+    StackSize = FirstSPAdjustAmount;
+  }
+
   // Deallocate stack
   adjustReg(MBB, MBBI, DL, SPReg, SPReg, StackSize, MachineInstr::FrameDestroy);
 }
 
+// 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.
+// e.g.
+//   addi.d  $sp, $sp, -2032
+//   st.d    $ra, $sp, 2024
+//   st.d    $fp, $sp, 2016
+//   addi.d  $sp, $sp, -16
+uint64_t LoongArchFrameLowering::getFirstSPAdjustAmount(
+    const MachineFunction &MF) const {
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  const std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
+
+  // Return the FirstSPAdjustAmount if the StackSize can not fit in a signed
+  // 12-bit and there exists a callee-saved register needing to be pushed.
+  if (!isInt<12>(MFI.getStackSize()) && (CSI.size() > 0)) {
+    // FirstSPAdjustAmount is chosen as (2048 - StackAlign) because 2048 will
+    // cause sp = sp + 2048 in the epilogue to be split into multiple
+    // instructions. Offsets smaller than 2048 can fit in a single load/store
+    // instruction, and we have to stick with the stack alignment.
+    // So (2048 - StackAlign) will satisfy the stack alignment.
+    return 2048 - getStackAlign().value();
+  }
+  return 0;
+}
+
 void LoongArchFrameLowering::determineCalleeSaves(MachineFunction &MF,
                                                   BitVector &SavedRegs,
                                                   RegScavenger *RS) const {
   TargetFrameLowering::determineCalleeSaves(MF, SavedRegs, RS);
   // Unconditionally spill RA and FP only if the function uses a frame
   // pointer.
   if (hasFP(MF)) {
     SavedRegs.set(LoongArch::R1);

 }
 
 StackOffset LoongArchFrameLowering::getFrameIndexReference(
     const MachineFunction &MF, int FI, Register &FrameReg) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const TargetRegisterInfo *RI = MF.getSubtarget().getRegisterInfo();
   auto *LoongArchFI = MF.getInfo<LoongArchMachineFunctionInfo>();
   uint64_t StackSize = MFI.getStackSize();
+  uint64_t FirstSPAdjustAmount = getFirstSPAdjustAmount(MF);
 
   // Callee-saved registers should be referenced relative to the stack
   // pointer (positive offset), otherwise use the frame pointer (negative
   // offset).
   const auto &CSI = MFI.getCalleeSavedInfo();
   int MinCSFI = 0;
   int MaxCSFI = -1;
   StackOffset Offset =
       StackOffset::getFixed(MFI.getObjectOffset(FI) - getOffsetOfLocalArea() +
                             MFI.getOffsetAdjustment());
 
   if (CSI.size()) {
     MinCSFI = CSI[0].getFrameIdx();
     MaxCSFI = CSI[CSI.size() - 1].getFrameIdx();
   }
 
   if (FI >= MinCSFI && FI <= MaxCSFI) {
     FrameReg = LoongArch::R3;
+    if (FirstSPAdjustAmount)
+      Offset += StackOffset::getFixed(FirstSPAdjustAmount);
+    else
       Offset += StackOffset::getFixed(StackSize);
   } else if (RI->hasStackRealignment(MF) && !MFI.isFixedObjectIndex(FI)) {
     // If the stack was realigned, the frame pointer is set in order to allow
     // SP to be restored, so we need another base register to record the stack
     // after realignment.
     FrameReg = hasBP(MF) ? LoongArchABI::getBPReg() : LoongArch::R3;
     Offset += StackOffset::getFixed(StackSize);
   } else {
     FrameReg = RI->getFrameRegister(MF);
     if (hasFP(MF))
       Offset += StackOffset::getFixed(LoongArchFI->getVarArgsSaveSize());
     else
       Offset += StackOffset::getFixed(StackSize);
   }
 
   return Offset;
 }

llvm/test/CodeGen/LoongArch/split-sp-adjust.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc --mtriple=loongarch64 --verify-machineinstrs < %s \
+; RUN:   | FileCheck %s
+
+;; The stack size is 2048 and the SP adjustment will be split.
+define i32 @SplitSP() nounwind {
+; CHECK-LABEL: SplitSP:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi.d $sp, $sp, -2032
+; CHECK-NEXT:    st.d $ra, $sp, 2024 # 8-byte Folded Spill
+; CHECK-NEXT:    addi.d $sp, $sp, -16
+; CHECK-NEXT:    addi.d $a0, $sp, 12
+; CHECK-NEXT:    bl %plt(foo)
+; CHECK-NEXT:    move $a0, $zero
+; CHECK-NEXT:    addi.d $sp, $sp, 16
+; CHECK-NEXT:    ld.d $ra, $sp, 2024 # 8-byte Folded Reload
+; CHECK-NEXT:    addi.d $sp, $sp, 2032
+; CHECK-NEXT:    ret
+entry:
+  %xx = alloca [2028 x i8], align 1
+  %0 = getelementptr inbounds [2028 x i8], ptr %xx, i32 0, i32 0
+  %call = call i32 @foo(ptr nonnull %0)
+  ret i32 0
+}
+
+;; The stack size is 2032 and the SP adjustment will not be split.
+define i32 @NoSplitSP() nounwind {
+; CHECK-LABEL: NoSplitSP:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi.d $sp, $sp, -2032
+; CHECK-NEXT:    st.d $ra, $sp, 2024 # 8-byte Folded Spill
+; CHECK-NEXT:    addi.d $a0, $sp, 0
+; CHECK-NEXT:    bl %plt(foo)
+; CHECK-NEXT:    move $a0, $zero
+; CHECK-NEXT:    ld.d $ra, $sp, 2024 # 8-byte Folded Reload
+; CHECK-NEXT:    addi.d $sp, $sp, 2032
+; CHECK-NEXT:    ret
+entry:
+  %xx = alloca [2024 x i8], align 1
+  %0 = getelementptr inbounds [2024 x i8], ptr %xx, i32 0, i32 0
+  %call = call i32 @foo(ptr nonnull %0)
+  ret i32 0
+}
+
+declare i32 @foo(ptr)

llvm/test/CodeGen/LoongArch/stack-realignment.ll

; LA64-NEXT:    ret
   %1 = alloca i8, align 1024
   call void @callee(i8* %1)
   ret void
 }
 
 define void @caller2048() {
 ; LA32-LABEL: caller2048:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    addi.w $sp, $sp, -2048
-; LA32-NEXT:    .cfi_def_cfa_offset 2048
-; LA32-NEXT:    st.w $ra, $sp, 2044 # 4-byte Folded Spill
-; LA32-NEXT:    st.w $fp, $sp, 2040 # 4-byte Folded Spill
+; LA32-NEXT:    addi.w $sp, $sp, -2032
+; LA32-NEXT:    .cfi_def_cfa_offset 2032
+; LA32-NEXT:    st.w $ra, $sp, 2028 # 4-byte Folded Spill
+; LA32-NEXT:    st.w $fp, $sp, 2024 # 4-byte Folded Spill
 ; LA32-NEXT:    .cfi_offset 1, -4
 ; LA32-NEXT:    .cfi_offset 22, -8
 ; LA32-NEXT:    addi.w $fp, $sp, 2032
-; LA32-NEXT:    addi.w $fp, $fp, 16
 ; LA32-NEXT:    .cfi_def_cfa 22, 0
+; LA32-NEXT:    addi.w $sp, $sp, -16
 ; LA32-NEXT:    srli.w $a0, $sp, 11
 ; LA32-NEXT:    slli.w $sp, $a0, 11
 ; LA32-NEXT:    addi.w $a0, $sp, 0
 ; LA32-NEXT:    bl %plt(callee)
 ; LA32-NEXT:    addi.w $sp, $fp, -2048
-; LA32-NEXT:    ld.w $fp, $sp, 2040 # 4-byte Folded Reload
-; LA32-NEXT:    ld.w $ra, $sp, 2044 # 4-byte Folded Reload
-; LA32-NEXT:    addi.w $sp, $sp, 2032
 ; LA32-NEXT:    addi.w $sp, $sp, 16
+; LA32-NEXT:    ld.w $fp, $sp, 2024 # 4-byte Folded Reload
+; LA32-NEXT:    ld.w $ra, $sp, 2028 # 4-byte Folded Reload
+; LA32-NEXT:    addi.w $sp, $sp, 2032
 ; LA32-NEXT:    ret
 ;
 ; LA64-LABEL: caller2048:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    addi.d $sp, $sp, -2048
-; LA64-NEXT:    .cfi_def_cfa_offset 2048
-; LA64-NEXT:    st.d $ra, $sp, 2040 # 8-byte Folded Spill
-; LA64-NEXT:    st.d $fp, $sp, 2032 # 8-byte Folded Spill
+; LA64-NEXT:    addi.d $sp, $sp, -2032
+; LA64-NEXT:    .cfi_def_cfa_offset 2032
+; LA64-NEXT:    st.d $ra, $sp, 2024 # 8-byte Folded Spill
+; LA64-NEXT:    st.d $fp, $sp, 2016 # 8-byte Folded Spill
 ; LA64-NEXT:    .cfi_offset 1, -8
 ; LA64-NEXT:    .cfi_offset 22, -16
 ; LA64-NEXT:    addi.d $fp, $sp, 2032
-; LA64-NEXT:    addi.d $fp, $fp, 16
 ; LA64-NEXT:    .cfi_def_cfa 22, 0
+; LA64-NEXT:    addi.d $sp, $sp, -16
 ; LA64-NEXT:    srli.d $a0, $sp, 11
 ; LA64-NEXT:    slli.d $sp, $a0, 11
 ; LA64-NEXT:    addi.d $a0, $sp, 0
 ; LA64-NEXT:    bl %plt(callee)
 ; LA64-NEXT:    addi.d $sp, $fp, -2048
-; LA64-NEXT:    ld.d $fp, $sp, 2032 # 8-byte Folded Reload
-; LA64-NEXT:    ld.d $ra, $sp, 2040 # 8-byte Folded Reload
-; LA64-NEXT:    addi.d $sp, $sp, 2032
 ; LA64-NEXT:    addi.d $sp, $sp, 16
+; LA64-NEXT:    ld.d $fp, $sp, 2016 # 8-byte Folded Reload
+; LA64-NEXT:    ld.d $ra, $sp, 2024 # 8-byte Folded Reload
+; LA64-NEXT:    addi.d $sp, $sp, 2032
 ; LA64-NEXT:    ret
   %1 = alloca i8, align 2048
   call void @callee(i8* %1)
   ret void
 }
 
 define void @caller_no_realign2048() "no-realign-stack" {
 ; LA32-LABEL: caller_no_realign2048:
; LA64-NEXT:    ret
   %1 = alloca i8, align 2048
   call void @callee(i8* %1)
   ret void
 }
 
 define void @caller4096() {
 ; LA32-LABEL: caller4096:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    lu12i.w $a0, 1
-; LA32-NEXT:    sub.w $sp, $sp, $a0
-; LA32-NEXT:    .cfi_def_cfa_offset 4096
-; LA32-NEXT:    ori $a0, $zero, 4092
-; LA32-NEXT:    add.w $a0, $sp, $a0
-; LA32-NEXT:    st.w $ra, $a0, 0 # 4-byte Folded Spill
-; LA32-NEXT:    ori $a0, $zero, 4088
-; LA32-NEXT:    add.w $a0, $sp, $a0
-; LA32-NEXT:    st.w $fp, $a0, 0 # 4-byte Folded Spill
+; LA32-NEXT:    addi.w $sp, $sp, -2032
+; LA32-NEXT:    .cfi_def_cfa_offset 2032
+; LA32-NEXT:    st.w $ra, $sp, 2028 # 4-byte Folded Spill
+; LA32-NEXT:    st.w $fp, $sp, 2024 # 4-byte Folded Spill
 ; LA32-NEXT:    .cfi_offset 1, -4
 ; LA32-NEXT:    .cfi_offset 22, -8
-; LA32-NEXT:    lu12i.w $a0, 1
-; LA32-NEXT:    add.w $fp, $sp, $a0
+; LA32-NEXT:    addi.w $fp, $sp, 2032
 ; LA32-NEXT:    .cfi_def_cfa 22, 0
+; LA32-NEXT:    addi.w $sp, $sp, -2048
+; LA32-NEXT:    addi.w $sp, $sp, -16
 ; LA32-NEXT:    srli.w $a0, $sp, 12
 ; LA32-NEXT:    slli.w $sp, $a0, 12
 ; LA32-NEXT:    addi.w $a0, $sp, 0
 ; LA32-NEXT:    bl %plt(callee)
 ; LA32-NEXT:    lu12i.w $a0, 1
 ; LA32-NEXT:    sub.w $sp, $fp, $a0
-; LA32-NEXT:    ori $a0, $zero, 4088
-; LA32-NEXT:    add.w $a0, $sp, $a0
-; LA32-NEXT:    ld.w $fp, $a0, 0 # 4-byte Folded Reload
-; LA32-NEXT:    ori $a0, $zero, 4092
-; LA32-NEXT:    add.w $a0, $sp, $a0
-; LA32-NEXT:    ld.w $ra, $a0, 0 # 4-byte Folded Reload
-; LA32-NEXT:    lu12i.w $a0, 1
-; LA32-NEXT:    add.w $sp, $sp, $a0
+; LA32-NEXT:    addi.w $sp, $sp, 2032
+; LA32-NEXT:    addi.w $sp, $sp, 32
+; LA32-NEXT:    ld.w $fp, $sp, 2024 # 4-byte Folded Reload
+; LA32-NEXT:    ld.w $ra, $sp, 2028 # 4-byte Folded Reload
+; LA32-NEXT:    addi.w $sp, $sp, 2032
 ; LA32-NEXT:    ret
 ;
 ; LA64-LABEL: caller4096:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    lu12i.w $a0, 1
-; LA64-NEXT:    sub.d $sp, $sp, $a0
-; LA64-NEXT:    .cfi_def_cfa_offset 4096
-; LA64-NEXT:    ori $a0, $zero, 4088
-; LA64-NEXT:    add.d $a0, $sp, $a0
-; LA64-NEXT:    st.d $ra, $a0, 0 # 8-byte Folded Spill
-; LA64-NEXT:    ori $a0, $zero, 4080
-; LA64-NEXT:    add.d $a0, $sp, $a0
-; LA64-NEXT:    st.d $fp, $a0, 0 # 8-byte Folded Spill
+; LA64-NEXT:    addi.d $sp, $sp, -2032
+; LA64-NEXT:    .cfi_def_cfa_offset 2032
+; LA64-NEXT:    st.d $ra, $sp, 2024 # 8-byte Folded Spill
+; LA64-NEXT:    st.d $fp, $sp, 2016 # 8-byte Folded Spill
 ; LA64-NEXT:    .cfi_offset 1, -8
 ; LA64-NEXT:    .cfi_offset 22, -16
-; LA64-NEXT:    lu12i.w $a0, 1
-; LA64-NEXT:    add.d $fp, $sp, $a0
+; LA64-NEXT:    addi.d $fp, $sp, 2032
 ; LA64-NEXT:    .cfi_def_cfa 22, 0
+; LA64-NEXT:    addi.d $sp, $sp, -2048
+; LA64-NEXT:    addi.d $sp, $sp, -16
 ; LA64-NEXT:    srli.d $a0, $sp, 12
 ; LA64-NEXT:    slli.d $sp, $a0, 12
 ; LA64-NEXT:    addi.d $a0, $sp, 0
 ; LA64-NEXT:    bl %plt(callee)
 ; LA64-NEXT:    lu12i.w $a0, 1
 ; LA64-NEXT:    sub.d $sp, $fp, $a0
-; LA64-NEXT:    ori $a0, $zero, 4080
-; LA64-NEXT:    add.d $a0, $sp, $a0
-; LA64-NEXT:    ld.d $fp, $a0, 0 # 8-byte Folded Reload
-; LA64-NEXT:    ori $a0, $zero, 4088
-; LA64-NEXT:    add.d $a0, $sp, $a0
-; LA64-NEXT:    ld.d $ra, $a0, 0 # 8-byte Folded Reload
-; LA64-NEXT:    lu12i.w $a0, 1
-; LA64-NEXT:    add.d $sp, $sp, $a0
+; LA64-NEXT:    addi.d $sp, $sp, 2032
+; LA64-NEXT:    addi.d $sp, $sp, 32
+; LA64-NEXT:    ld.d $fp, $sp, 2016 # 8-byte Folded Reload
+; LA64-NEXT:    ld.d $ra, $sp, 2024 # 8-byte Folded Reload
+; LA64-NEXT:    addi.d $sp, $sp, 2032
 ; LA64-NEXT:    ret
   %1 = alloca i8, align 4096
   call void @callee(i8* %1)
   ret void
 }
 
 define void @caller_no_realign4096() "no-realign-stack" {
 ; LA32-LABEL: caller_no_realign4096: