[llvm] [NFC][RISCV] Refactor allocation of the stack space (PR #116625)

[Raphael Moreira Zinsly via llvm-commits]

https://github.com/rzinsly created https://github.com/llvm/llvm-project/pull/116625

Separates the stack allocations from prologue in preparation for the stack clash protection support.

>From 4678b1c85da95b6f4ea08b35b3d567e516e1277b Mon Sep 17 00:00:00 2001
From: Raphael Moreira Zinsly <rzinsly at ventanamicro.com>
Date: Mon, 18 Nov 2024 10:57:10 -0300
Subject: [PATCH] [NFC][RISCV] Refactor allocation of the stack space

Separates the stack allocations from prologue in preparation for the
stack clash protection support.
---
 llvm/lib/Target/RISCV/RISCVFrameLowering.cpp | 43 +++++++++++---------
 llvm/lib/Target/RISCV/RISCVFrameLowering.h   |  3 ++
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
index 1d91d46cb30ee0..8ade1943a3433b 100644
--- a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
@@ -610,6 +610,25 @@ static MCCFIInstruction createDefCFAOffset(const TargetRegisterInfo &TRI,
                                         Comment.str());
 }
 
+void RISCVFrameLowering::allocateStack(MachineBasicBlock &MBB,
+                                       MachineBasicBlock::iterator MBBI,
+                                       StackOffset Offset, bool EmitCFI,
+                                       unsigned CFIIndex) const {
+  DebugLoc DL;
+  const RISCVRegisterInfo *RI = STI.getRegisterInfo();
+  const RISCVInstrInfo *TII = STI.getInstrInfo();
+
+  RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg, Offset, MachineInstr::FrameSetup,
+                getStackAlign());
+
+  if (EmitCFI) {
+    // Emit ".cfi_def_cfa_offset StackSize"
+    BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+        .addCFIIndex(CFIIndex)
+        .setMIFlag(MachineInstr::FrameSetup);
+  }
+}
+
 void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
                                       MachineBasicBlock &MBB) const {
   MachineFrameInfo &MFI = MF.getFrameInfo();
@@ -726,16 +745,9 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
 
   if (StackSize != 0) {
     // Allocate space on the stack if necessary.
-    RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg,
-                  StackOffset::getFixed(-StackSize), MachineInstr::FrameSetup,
-                  getStackAlign());
-
-    // 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);
+    allocateStack(MBB, MBBI, StackOffset::getFixed(-StackSize), true, CFIIndex);
   }
 
   // The frame pointer is callee-saved, and code has been generated for us to
@@ -776,20 +788,13 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
         getStackSizeWithRVVPadding(MF) - FirstSPAdjustAmount;
     assert(SecondSPAdjustAmount > 0 &&
            "SecondSPAdjustAmount should be greater than zero");
-    RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg,
-                  StackOffset::getFixed(-SecondSPAdjustAmount),
-                  MachineInstr::FrameSetup, getStackAlign());
 
     // 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
-    if (!hasFP(MF)) {
-      // Emit ".cfi_def_cfa_offset StackSize"
-      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::cfiDefCfaOffset(
-          nullptr, getStackSizeWithRVVPadding(MF)));
-      BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
-          .addCFIIndex(CFIIndex)
-          .setMIFlag(MachineInstr::FrameSetup);
-    }
+    unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::cfiDefCfaOffset(
+        nullptr, getStackSizeWithRVVPadding(MF)));
+    allocateStack(MBB, MBBI, StackOffset::getFixed(-SecondSPAdjustAmount),
+                  !hasFP(MF), CFIIndex);
   }
 
   if (RVVStackSize) {
diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.h b/llvm/lib/Target/RISCV/RISCVFrameLowering.h
index c106b7b6754653..9aff4dc9e4089a 100644
--- a/llvm/lib/Target/RISCV/RISCVFrameLowering.h
+++ b/llvm/lib/Target/RISCV/RISCVFrameLowering.h
@@ -78,6 +78,9 @@ class RISCVFrameLowering : public TargetFrameLowering {
     return StackId != TargetStackID::ScalableVector;
   }
 
+  void allocateStack(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+                     StackOffset Offset, bool EmitCFI, unsigned CFIIndex) const;
+
 protected:
   const RISCVSubtarget &STI;