[RISCV] Handle variable sized objects with the stack need to be realigned

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

lib/Target/RISCV/RISCVFrameLowering.cpp

@@ -32,6 +32,13 @@ bool RISCVFrameLowering::hasFP(const MachineFunction &MF) const {
          MFI.isFrameAddressTaken();
 }
 
+bool RISCVFrameLowering::hasBP(const MachineFunction &MF) const {
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  const TargetRegisterInfo *TRI = STI.getRegisterInfo();
+
+  return MFI.hasVarSizedObjects() && TRI->needsStackRealignment(MF);
+}
+
 // Determines the size of the frame and maximum call frame size.
 void RISCVFrameLowering::determineFrameLayout(MachineFunction &MF) const {
   MachineFrameInfo &MFI = MF.getFrameInfo();
@@ -108,14 +115,9 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
   const RISCVInstrInfo *TII = STI.getInstrInfo();
   MachineBasicBlock::iterator MBBI = MBB.begin();
 
-  if (RI->needsStackRealignment(MF) && MFI.hasVarSizedObjects()) {
-    report_fatal_error(
-        "RISC-V backend can't currently handle functions that need stack "
-        "realignment and have variable sized objects");
-  }
-
   Register FPReg = getFPReg(STI);
   Register SPReg = getSPReg(STI);
+  Register BPReg = RISCVABI::getBPReg();
 
   // Debug location must be unknown since the first debug location is used
   // to determine the end of the prologue.
@@ -229,6 +231,15 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
             .addReg(VR)
             .addImm(ShiftAmount);
       }
+      // FP will be used to restore the frame in the epilogue, so we need
+      // another base register BP to record SP after re-alignment. SP will
+      // track the current stack after allocating variable sized objects.
+      if (hasBP(MF)) {
+        // move BP, SP
+        BuildMI(MBB, MBBI, DL, TII->get(RISCV::ADDI), BPReg)
+            .addReg(SPReg)
+            .addImm(0);
+      }
     }
   }
 }
@@ -308,11 +319,13 @@ int RISCVFrameLowering::getFrameIndexReference(const MachineFunction &MF,
       Offset += FirstSPAdjustAmount;
     else
       Offset += MF.getFrameInfo().getStackSize();
-  } else if (RI->needsStackRealignment(MF)) {
+  } else if (RI->needsStackRealignment(MF) && !MFI.isFixedObjectIndex(FI)) {
-    assert(!MFI.hasVarSizedObjects() &&
-           "Unexpected combination of stack realignment and varsized objects");
     // If the stack was realigned, the frame pointer is set in order to allow
-    // SP to be restored, but we still access stack objects using SP.
+    // SP to be restored, so we need another base register to record the stack
+    // after realignment.
+    if (hasBP(MF))
+      FrameReg = RISCVABI::getBPReg();
+    else
       FrameReg = RISCV::X2;
     Offset += MF.getFrameInfo().getStackSize();
   } else {
@@ -335,6 +348,9 @@ void RISCVFrameLowering::determineCalleeSaves(MachineFunction &MF,
     SavedRegs.set(RISCV::X1);
     SavedRegs.set(RISCV::X8);
   }
+  // Mark BP as used if function has dedicated base pointer.
+  if (hasBP(MF))
+    SavedRegs.set(RISCVABI::getBPReg());
 
   // If interrupt is enabled and there are calls in the handler,
   // unconditionally save all Caller-saved registers and

lib/Target/RISCV/RISCVFrameLowering.h

@@ -40,6 +40,8 @@ public:
 
   bool hasFP(const MachineFunction &MF) const override;
 
+  bool hasBP(const MachineFunction &MF) const;
+
   bool hasReservedCallFrame(const MachineFunction &MF) const override;
   MachineBasicBlock::iterator
   eliminateCallFramePseudoInstr(MachineFunction &MF, MachineBasicBlock &MBB,

lib/Target/RISCV/RISCVRegisterInfo.cpp

@@ -66,7 +66,7 @@ RISCVRegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
 }
 
 BitVector RISCVRegisterInfo::getReservedRegs(const MachineFunction &MF) const {
-  const TargetFrameLowering *TFI = getFrameLowering(MF);
+  const RISCVFrameLowering *TFI = getFrameLowering(MF);
   BitVector Reserved(getNumRegs());
 
   // Mark any registers requested to be reserved as such
@@ -82,6 +82,10 @@ BitVector RISCVRegisterInfo::getReservedRegs(const MachineFunction &MF) const {
   markSuperRegs(Reserved, RISCV::X4); // tp
   if (TFI->hasFP(MF))
     markSuperRegs(Reserved, RISCV::X8); // fp
+  // Reserve the base register if we need to realign the stack and allocate
+  // variable-sized objects at runtime.
+  if (TFI->hasBP(MF))
+    markSuperRegs(Reserved, RISCVABI::getBPReg()); // bp
   assert(checkAllSuperRegsMarked(Reserved));
   return Reserved;
 }

lib/Target/RISCV/Utils/RISCVBaseInfo.cpp

@@ -66,6 +66,12 @@ ABI computeTargetABI(const Triple &TT, FeatureBitset FeatureBits,
     return ABI_LP64;
   return ABI_ILP32;
 }
+
+// To avoid the BP value clobbered by a function call, we need to choose a
+// callee saved register to save the value. RV32E only has X8 and X9 as callee
+// saved registers and X8 will be used as fp. So we choose X9 as bp.
+Register getBPReg() { return RISCV::X9; }
+
 } // namespace RISCVABI
 
 namespace RISCVFeatures {

lib/Target/RISCV/Utils/RISCVBaseInfo.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_LIB_TARGET_RISCV_MCTARGETDESC_RISCVBASEINFO_H
 #define LLVM_LIB_TARGET_RISCV_MCTARGETDESC_RISCVBASEINFO_H
 
+#include "RISCVRegisterInfo.h"
 #include "MCTargetDesc/RISCVMCTargetDesc.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
@@ -195,6 +196,9 @@ enum ABI {
 ABI computeTargetABI(const Triple &TT, FeatureBitset FeatureBits,
                      StringRef ABIName);
 
+// Returns the register used to hold the stack pointer after realignment.
+Register getBPReg();
+
 } // namespace RISCVABI
 
 namespace RISCVFeatures {

test/CodeGen/RISCV/stack-realignment-unsupported.ll

@@ -1,13 +0,0 @@
-; RUN: not llc -mtriple=riscv32 < %s 2>&1 | FileCheck %s
-; RUN: not llc -mtriple=riscv64 < %s 2>&1 | FileCheck %s
-
-; CHECK: LLVM ERROR: RISC-V backend can't currently handle functions that need stack realignment and have variable sized objects
-
-declare void @callee(i8*, i32*)
-
-define void @caller(i32 %n) nounwind {
-  %1 = alloca i8, i32 %n
-  %2 = alloca i32, align 64
-  call void @callee(i8* %1, i32 *%2)
-  ret void
-}

test/CodeGen/RISCV/stack-realignment-with-variable-sized-objects.ll

@@ -0,0 +1,72 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 -verify-machineinstrs < %s \
+; RUN:   | FileCheck %s -check-prefix=RV32I
+; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s \
+; RUN:   | FileCheck %s -check-prefix=RV64I
+
+declare void @callee(i8*, i32*)
+
+define void @caller(i32 %n) {
+; RV32I-LABEL: caller:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    addi sp, sp, -128
+; RV32I-NEXT:    .cfi_def_cfa_offset 128
+; RV32I-NEXT:    sw ra, 124(sp)
+; RV32I-NEXT:    sw s0, 120(sp)
+; RV32I-NEXT:    sw s1, 116(sp)
+; RV32I-NEXT:    .cfi_offset ra, -4
+; RV32I-NEXT:    .cfi_offset s0, -8
+; RV32I-NEXT:    .cfi_offset s1, -12
+; RV32I-NEXT:    addi s0, sp, 128
+; RV32I-NEXT:    .cfi_def_cfa s0, 0
+; RV32I-NEXT:    andi sp, sp, -64
+; RV32I-NEXT:    mv s1, sp
+; RV32I-NEXT:    addi a0, a0, 15
+; RV32I-NEXT:    andi a0, a0, -16
+; RV32I-NEXT:    sub a0, sp, a0
+; RV32I-NEXT:    mv sp, a0
+; RV32I-NEXT:    addi a1, s1, 64
+; RV32I-NEXT:    call callee
+; RV32I-NEXT:    addi sp, s0, -128
+; RV32I-NEXT:    lw s1, 116(sp)
+; RV32I-NEXT:    lw s0, 120(sp)
+; RV32I-NEXT:    lw ra, 124(sp)
+; RV32I-NEXT:    addi sp, sp, 128
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: caller:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    addi sp, sp, -128
+; RV64I-NEXT:    .cfi_def_cfa_offset 128
+; RV64I-NEXT:    sd ra, 120(sp)
+; RV64I-NEXT:    sd s0, 112(sp)
+; RV64I-NEXT:    sd s1, 104(sp)
+; RV64I-NEXT:    .cfi_offset ra, -8
+; RV64I-NEXT:    .cfi_offset s0, -16
+; RV64I-NEXT:    .cfi_offset s1, -24
+; RV64I-NEXT:    addi s0, sp, 128
+; RV64I-NEXT:    .cfi_def_cfa s0, 0
+; RV64I-NEXT:    andi sp, sp, -64
+; RV64I-NEXT:    mv s1, sp
+; RV64I-NEXT:    slli a0, a0, 32
+; RV64I-NEXT:    srli a0, a0, 32
+; RV64I-NEXT:    addi a0, a0, 15
+; RV64I-NEXT:    addi a1, zero, 1
+; RV64I-NEXT:    slli a1, a1, 33
+; RV64I-NEXT:    addi a1, a1, -16
+; RV64I-NEXT:    and a0, a0, a1
+; RV64I-NEXT:    sub a0, sp, a0
+; RV64I-NEXT:    mv sp, a0
+; RV64I-NEXT:    addi a1, s1, 64
+; RV64I-NEXT:    call callee
+; RV64I-NEXT:    addi sp, s0, -128
+; RV64I-NEXT:    ld s1, 104(sp)
+; RV64I-NEXT:    ld s0, 112(sp)
+; RV64I-NEXT:    ld ra, 120(sp)
+; RV64I-NEXT:    addi sp, sp, 128
+; RV64I-NEXT:    ret
+  %1 = alloca i8, i32 %n
+  %2 = alloca i32, align 64
+  call void @callee(i8* %1, i32 *%2)
+  ret void
+}