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

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

5 files changed, 44 insertions, 12 deletions

diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
index 4ba4606612b..61516960b39 100644
--- a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
+++ b/llvm/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,12 +319,14 @@ int RISCVFrameLowering::getFrameIndexReference(const MachineFunction &MF,
       Offset += FirstSPAdjustAmount;
     else
       Offset += MF.getFrameInfo().getStackSize();
-  } else if (RI->needsStackRealignment(MF)) {
-    assert(!MFI.hasVarSizedObjects() &&
-           "Unexpected combination of stack realignment and varsized objects");
+  } else if (RI->needsStackRealignment(MF) && !MFI.isFixedObjectIndex(FI)) {
     // 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.
-    FrameReg = RISCV::X2;
+    // 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 {
     FrameReg = RI->getFrameRegister(MF);
@@ -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
diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.h b/llvm/lib/Target/RISCV/RISCVFrameLowering.h
index f4a5949773d..3a16cf93cf1 100644
--- a/llvm/lib/Target/RISCV/RISCVFrameLowering.h
+++ b/llvm/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,
diff --git a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
index d4c65a0b22b..1d41994ef1e 100644
--- a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
+++ b/llvm/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;
 }
diff --git a/llvm/lib/Target/RISCV/Utils/RISCVBaseInfo.cpp b/llvm/lib/Target/RISCV/Utils/RISCVBaseInfo.cpp
index bc5395768ca..9c54c3749a1 100644
--- a/llvm/lib/Target/RISCV/Utils/RISCVBaseInfo.cpp
+++ b/llvm/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 {
diff --git a/llvm/lib/Target/RISCV/Utils/RISCVBaseInfo.h b/llvm/lib/Target/RISCV/Utils/RISCVBaseInfo.h
index 30e475e80a0..738f635ada9 100644
--- a/llvm/lib/Target/RISCV/Utils/RISCVBaseInfo.h
+++ b/llvm/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 {