[RISCV] fix RISCVPushPopOptimizer pass

llvm/lib/Target/RISCV/RISCVFrameLowering.cpp

@@ -761,18 +761,26 @@ void RISCVFrameLowering::deallocateStack(MachineFunction &MF,
                                          const DebugLoc &DL, uint64_t StackSize,
                                          int64_t CFAOffset) const {
   const RISCVRegisterInfo *RI = STI.getRegisterInfo();
+  const RISCVInstrInfo *TII = STI.getInstrInfo();
 
   Register SPReg = getSPReg(STI);
 
   RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg, StackOffset::getFixed(StackSize),
                 MachineInstr::FrameDestroy, getStackAlign());
+
+  unsigned CFIIndex =
+      MF.addFrameInst(MCCFIInstruction::cfiDefCfaOffset(nullptr, CFAOffset));
+  BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+      .addCFIIndex(CFIIndex)
+      .setMIFlag(MachineInstr::FrameDestroy);
 }
 
 void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
                                       MachineBasicBlock &MBB) const {
   const RISCVRegisterInfo *RI = STI.getRegisterInfo();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   auto *RVFI = MF.getInfo<RISCVMachineFunctionInfo>();
+  const RISCVInstrInfo *TII = STI.getInstrInfo();
   Register FPReg = getFPReg(STI);
   Register SPReg = getSPReg(STI);
 
@@ -819,13 +827,22 @@ void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
 
   bool RestoreFP = RI->hasStackRealignment(MF) || MFI.hasVarSizedObjects() ||
                    !hasReservedCallFrame(MF);
-
   if (RVVStackSize) {
     // If RestoreFP the stack pointer will be restored using the frame pointer
     // value.
-    if (!RestoreFP)
+    if (!RestoreFP) {
       adjustStackForRVV(MF, MBB, LastFrameDestroy, DL, RVVStackSize,
                         MachineInstr::FrameDestroy);
+
+      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::cfiDefCfa(
+          nullptr, RI->getDwarfRegNum(SPReg, true), RealStackSize));
+      BuildMI(MBB, LastFrameDestroy, DL,
+              TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex)
+          .setMIFlag(MachineInstr::FrameDestroy);
+    }
+
+    emitCalleeSavedRVVEpilogCFI(MBB, LastFrameDestroy);
   }
 
   if (FirstSPAdjustAmount) {
@@ -836,10 +853,18 @@ void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
 
     // If RestoreFP the stack pointer will be restored using the frame pointer
     // value.
-    if (!RestoreFP)
+    if (!RestoreFP) {
       RI->adjustReg(MBB, LastFrameDestroy, DL, SPReg, SPReg,
                     StackOffset::getFixed(SecondSPAdjustAmount),
                     MachineInstr::FrameDestroy, getStackAlign());
+
+      unsigned CFIIndex = MF.addFrameInst(
+          MCCFIInstruction::cfiDefCfaOffset(nullptr, FirstSPAdjustAmount));
+      BuildMI(MBB, LastFrameDestroy, DL,
+              TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex)
+          .setMIFlag(MachineInstr::FrameDestroy);
+    }
   }
 
   // Restore the stack pointer using the value of the frame pointer. Only
@@ -858,6 +883,25 @@ void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
     RI->adjustReg(MBB, LastFrameDestroy, DL, SPReg, FPReg,
                   StackOffset::getFixed(-FPOffset), MachineInstr::FrameDestroy,
                   getStackAlign());
+
+    unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::cfiDefCfa(
+        nullptr, RI->getDwarfRegNum(SPReg, true), RealStackSize));
+    BuildMI(MBB, LastFrameDestroy, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+        .addCFIIndex(CFIIndex)
+        .setMIFlag(MachineInstr::FrameDestroy);
+  }
+
+  if (getLibCallID(MF, CSI) + 1) {
+    // tail __riscv_restore_[0-12] instruction is considered as a terminator,
+    // therefor it is unnecessary to place any CFI instructions after it. Just
+    // deallocate stack if needed and return.
+    if (StackSize != 0)
+      deallocateStack(MF, MBB, MBBI, DL, StackSize,
+                      RVFI->getLibCallStackSize());
+
+    // Emit epilogue for shadow call stack.
+    emitSCSEpilogue(MF, MBB, MBBI, DL);
+    return;
   }
 
   bool ApplyPop = RVFI->isPushable(MF) && MBBI != MBB.end() &&
@@ -876,6 +920,23 @@ void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
                       /*stack_adj of cm.pop instr*/ RealStackSize - StackSize);
 
     ++MBBI;
+    // Update CFA offset. After CM_POP SP should be equal to CFA, so CFA offset
+    // should be a zero.
+    unsigned CFIIndex =
+        MF.addFrameInst(MCCFIInstruction::cfiDefCfaOffset(nullptr, 0));
+    BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+        .addCFIIndex(CFIIndex)
+        .setMIFlag(MachineInstr::FrameDestroy);
+  }
+
+  // Recover callee-saved registers.
+  for (const auto &Entry : CSI) {
+    Register Reg = Entry.getReg();
+    unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createRestore(
+        nullptr, RI->getDwarfRegNum(Reg, true)));
+    BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+        .addCFIIndex(CFIIndex)
+        .setMIFlag(MachineInstr::FrameDestroy);
   }
 
   // Deallocate stack if we didn't already do it during cm.pop handling and
@@ -1614,6 +1675,28 @@ void RISCVFrameLowering::emitCalleeSavedRVVPrologCFI(
   }
 }
 
+void RISCVFrameLowering::emitCalleeSavedRVVEpilogCFI(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MI) const {
+  MachineFunction *MF = MBB.getParent();
+  const MachineFrameInfo &MFI = MF->getFrameInfo();
+  const RISCVRegisterInfo *RI = STI.getRegisterInfo();
+  const TargetInstrInfo &TII = *STI.getInstrInfo();
+  DebugLoc DL = MBB.findDebugLoc(MI);
+
+  const auto &RVVCSI = getRVVCalleeSavedInfo(*MF, MFI.getCalleeSavedInfo());
+  for (auto &CS : RVVCSI) {
+    int FI = CS.getFrameIdx();
+    if (FI >= 0 && MFI.getStackID(FI) == TargetStackID::ScalableVector) {
+      Register Reg = CS.getReg();
+      unsigned CFIIndex = MF->addFrameInst(MCCFIInstruction::createRestore(
+          nullptr, RI->getDwarfRegNum(Reg, true)));
+      BuildMI(MBB, MI, DL, TII.get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex)
+          .setMIFlag(MachineInstr::FrameDestroy);
+    }
+  }
+}
+
 bool RISCVFrameLowering::restoreCalleeSavedRegisters(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
     MutableArrayRef<CalleeSavedInfo> CSI, const TargetRegisterInfo *TRI) const {

llvm/lib/Target/RISCV/RISCVFrameLowering.h

@@ -91,6 +91,8 @@ class RISCVFrameLowering : public TargetFrameLowering {
   void emitCalleeSavedRVVPrologCFI(MachineBasicBlock &MBB,
                                    MachineBasicBlock::iterator MI,
                                    bool HasFP) const;
+  void emitCalleeSavedRVVEpilogCFI(MachineBasicBlock &MBB,
+                                   MachineBasicBlock::iterator MI) const;
   void deallocateStack(MachineFunction &MF, MachineBasicBlock &MBB,
                        MachineBasicBlock::iterator MBBI, const DebugLoc &DL,
                        uint64_t StackSize, int64_t CFAOffset) const;

llvm/lib/Target/RISCV/RISCVPushPopOptimizer.cpp

@@ -45,10 +45,25 @@ char RISCVPushPopOpt::ID = 0;
 INITIALIZE_PASS(RISCVPushPopOpt, "riscv-push-pop-opt", RISCV_PUSH_POP_OPT_NAME,
                 false, false)
 
+template <typename IterT>
+static IterT nextNoDebugNoCFIInst(const IterT It, const IterT End) {
+  return std::find_if_not(std::next(It), End, [](const auto &Inst) {
+    return Inst.isDebugInstr() || Inst.isCFIInstruction() ||
+           Inst.isPseudoProbe();
+  });
+}
+
+static void eraseCFIInst(const MachineBasicBlock::iterator Begin,
+                         const MachineBasicBlock::iterator End) {
+  for (auto &Inst : llvm::make_early_inc_range(llvm::make_range(Begin, End)))
+    if (Inst.isCFIInstruction())
+      Inst.eraseFromParent();
+}
+
 // Check if POP instruction was inserted into the MBB and return iterator to it.
 static MachineBasicBlock::iterator containsPop(MachineBasicBlock &MBB) {
   for (MachineBasicBlock::iterator MBBI = MBB.begin(); MBBI != MBB.end();
-       MBBI = next_nodbg(MBBI, MBB.end()))
+       MBBI = nextNoDebugNoCFIInst(MBBI, MBB.end()))
     if (MBBI->getOpcode() == RISCV::CM_POP)
       return MBBI;
 
@@ -76,6 +91,9 @@ bool RISCVPushPopOpt::usePopRet(MachineBasicBlock::iterator &MBBI,
   for (unsigned i = FirstNonDeclaredOp; i < MBBI->getNumOperands(); ++i)
     PopRetBuilder.add(MBBI->getOperand(i));
 
+  // Remove CFI instructions, they are not needed for cm.popret and cm.popretz
+  eraseCFIInst(MBBI, NextI);
+
   MBBI->eraseFromParent();
   NextI->eraseFromParent();
   return true;
@@ -92,8 +110,8 @@ bool RISCVPushPopOpt::adjustRetVal(MachineBasicBlock::iterator &MBBI) {
   // Since POP instruction is in Epilogue no normal instructions will follow
   // after it. Therefore search only previous ones to find the return value.
   for (MachineBasicBlock::reverse_iterator I =
-           next_nodbg(MBBI.getReverse(), RE);
-       I != RE; I = next_nodbg(I, RE)) {
+           nextNoDebugNoCFIInst(MBBI.getReverse(), RE);
+       I != RE; I = nextNoDebugNoCFIInst(I, RE)) {
     MachineInstr &MI = *I;
     if (auto OperandPair = TII->isCopyInstrImpl(MI)) {
       Register DestReg = OperandPair->Destination->getReg();
@@ -138,7 +156,7 @@ bool RISCVPushPopOpt::runOnMachineFunction(MachineFunction &Fn) {
   bool Modified = false;
   for (auto &MBB : Fn) {
     MachineBasicBlock::iterator MBBI = containsPop(MBB);
-    MachineBasicBlock::iterator NextI = next_nodbg(MBBI, MBB.end());
+    MachineBasicBlock::iterator NextI = nextNoDebugNoCFIInst(MBBI, MBB.end());
     if (MBBI != MBB.end() && NextI != MBB.end() &&
         NextI->getOpcode() == RISCV::PseudoRET)
       Modified |= usePopRet(MBBI, NextI, adjustRetVal(MBBI));

llvm/lib/Target/RISCV/RISCVTargetMachine.cpp

@@ -171,6 +171,8 @@ RISCVTargetMachine::RISCVTargetMachine(const Target &T, const Triple &TT,
 
   if (TT.isOSFuchsia() && !TT.isArch64Bit())
     report_fatal_error("Fuchsia is only supported for 64-bit");
+
+  setCFIFixup(true);
 }
 
 const RISCVSubtarget *

llvm/test/CodeGen/RISCV/GlobalISel/stacksave-stackrestore.ll

@@ -25,10 +25,15 @@ define void @test_scoped_alloca(i64 %n) {
 ; RV32-NEXT:    call use_addr
 ; RV32-NEXT:    mv sp, s1
 ; RV32-NEXT:    addi sp, s0, -16
+; RV32-NEXT:    .cfi_def_cfa sp, 16
 ; RV32-NEXT:    lw ra, 12(sp) # 4-byte Folded Reload
 ; RV32-NEXT:    lw s0, 8(sp) # 4-byte Folded Reload
 ; RV32-NEXT:    lw s1, 4(sp) # 4-byte Folded Reload
+; RV32-NEXT:    .cfi_restore ra
+; RV32-NEXT:    .cfi_restore s0
+; RV32-NEXT:    .cfi_restore s1
 ; RV32-NEXT:    addi sp, sp, 16
+; RV32-NEXT:    .cfi_def_cfa_offset 0
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: test_scoped_alloca:
@@ -51,10 +56,15 @@ define void @test_scoped_alloca(i64 %n) {
 ; RV64-NEXT:    call use_addr
 ; RV64-NEXT:    mv sp, s1
 ; RV64-NEXT:    addi sp, s0, -32
+; RV64-NEXT:    .cfi_def_cfa sp, 32
 ; RV64-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
 ; RV64-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
 ; RV64-NEXT:    ld s1, 8(sp) # 8-byte Folded Reload
+; RV64-NEXT:    .cfi_restore ra
+; RV64-NEXT:    .cfi_restore s0
+; RV64-NEXT:    .cfi_restore s1
 ; RV64-NEXT:    addi sp, sp, 32
+; RV64-NEXT:    .cfi_def_cfa_offset 0
 ; RV64-NEXT:    ret
   %sp = call ptr @llvm.stacksave.p0()
   %addr = alloca i8, i64 %n