[Power9] Allow gpr callee saved spills in prologue to vector registers rather than stack

[syzaara]

Summary

Currently in llvm, CalleeSavedInfo can only assign a callee saved register to stack frame index to be spilled in the prologue. We would like to enable spilling gprs to vector registers. This patch adds the capability to spill to other registers aside from just the stack. It also adds the changes for power9 to spill gprs to volatile vector registers when they are available. This happens only for leaf functions when using the option -ppc-enable-pe-vector-spills.

[hfinkel]

This is pretty neat.

Is the CFI limitation temporary? I imagine that you can use MCCFIInstruction::createRegister to indicate that the value that used to be in the GPR is now in the vector register?

Will be also be able to use this same infrastructure in order to save callee-saved condition registers to otherwise available GPRs (instead of to the stack)?

[syzaara]

In D39386#910018, @hfinkel wrote:

This is pretty neat.

Is the CFI limitation temporary? I imagine that you can use MCCFIInstruction::createRegister to indicate that the value that used to be in the GPR is now in the vector register?

Will be also be able to use this same infrastructure in order to save callee-saved condition registers to otherwise available GPRs (instead of to the stack)?

The CFI limitation is temporary. I did notice the createRegister, but I don't think there is support in the debuggers/unwinders for moving from a vector register to a GPR. Once we know how to handle this, we can remove the CFI limitation. The same infrastructure can be used for other available GPRs and I think createRegister should be able to handle that.

[hfinkel]

In D39386#910734, @syzaara wrote:

In D39386#910018, @hfinkel wrote:

This is pretty neat.

Is the CFI limitation temporary? I imagine that you can use MCCFIInstruction::createRegister to indicate that the value that used to be in the GPR is now in the vector register?

Will be also be able to use this same infrastructure in order to save callee-saved condition registers to otherwise available GPRs (instead of to the stack)?

The CFI limitation is temporary. I did notice the createRegister, but I don't think there is support in the debuggers/unwinders for moving from a vector register to a GPR.

Okay. I'd hope that this would "just work" (I don't see any special cases in how DW_CFA_register, etc. is handled in GCC's unwind-dw2.c that would lead me to believe that it wouldn't). I suppose we'll find out.

Once we know how to handle this, we can remove the CFI limitation. The same infrastructure can be used for other available GPRs and I think createRegister should be able to handle that.

[syzaara]

Use createRegister to emit .cfi_register for debug info.

[syzaara]

@hfinkel Can you please finish reviewing this.

[thegameg]

include/llvm/CodeGen/MachineFrameInfo.h

36

How about

union {
  int FrameIdx;
  unsigned DstReg;
};

to make it more obvious that you can't have both?

270–272

This should also be updated to warn users that not all CSRs are saved in the frame (and that something like CSInfo.size() * CSRSize is not totally accurate anymore, even though I'm not sure if any part of the code relies on this to allocate / estimate stack size).

lib/Target/PowerPC/PPCFrameLowering.cpp

1984

I may be missing something, but why is this needed? Won't it be def'd anyway?

lib/Target/PowerPC/PPCFrameLowering.h

102

grps -> gprs

test/CodeGen/PowerPC/prolog_vec_spills.ll

20

You can also write MIR tests with stuff like:

%R14 = IMPLICIT_DEF
%R15 = IMPLICIT_DEF
%R16 = IMPLICIT_DEF

to have more clear tests (with -run-pass prologepilog).

[syzaara]

lib/Target/PowerPC/PPCFrameLowering.cpp

1984

Without this, we will get machine verify errors in the epilogue that we are using an undefined physical register. This can be fixed by adding the register as live into the block. However, adding it to only the epilogue isn't enough since other blocks may clobber the register in between the prologue and epilogue.

[thegameg]

lib/Target/PowerPC/PPCFrameLowering.cpp

1984

Ah, right, I see. There is a function called updateLiveness in PrologEpilogInserter, which propagates the liveness of CSRs. You might want to look into that.

[syzaara]

Add handling in updateLiveness to add spilled destination register as live into all MachineBasicBlocks of the MachineFunction.

[thegameg]

lib/CodeGen/PrologEpilogInserter.cpp

475

I agree that in the general case all the basic blocks need to have this register as live in, but if the prologue/epilogue of the function is shrink-wrapped, only the blocks between the prologue and the epilogue need to have this as live-in.
I think at this point in the function the Visited set contains all the blocks *outside* the prologue/epilogue region. Correct me if I'm wrong but I think you're missing a:

if (Visited.count(&MBB))
  continue;

[syzaara]

@hfinkel Hi Hal, are you okay with approving this now?

[hfinkel]

lib/CodeGen/PrologEpilogInserter.cpp

475

Please add a comment here explaining that the visited set contains the blocks outside the prologue/epilogue region. That's likely not obvious.

[nemanjai]

lib/Target/PowerPC/PPCFrameLowering.cpp

42

I am really not a fan of defining this here. This should be available from the .td files. Can we not get it from there?

1976

No magic numbers please. If you are using the static array you defined above, you can insert a marker at the end and loop until you get to it (perhaps PPC::NoRegister)

1987

I'm just curious about this entire section. I imagine this is just a copy of the generic implementation. If that is the case, can we just return false from this function and modify the generic implementation to just ignore any that are isSpilledToReg() (assuming the target has already done that)?
Perhaps if we do that, the generic code can also assert that the DestReg is set for any such CSI objects.

[nemanjai]

@syzaara Do you plan to update this to address the comments?

[nemanjai]

Thank you for updating the code to get the CSR list from RegInfo. It looks much cleaner than having a static list of registers. I think the code can be cleaned up a bit further still, but overall this looks great.

llvm/include/llvm/CodeGen/MachineFrameInfo.h

64 ↗

(On Diff #148042)

I think setFrameIdx() and setDstReg() should probably update SpilledToReg.

265 ↗

(On Diff #148042)

... this data *is* used...

llvm/lib/Target/PowerPC/PPCFrameLowering.cpp

1963 ↗

(On Diff #148042)

A comment here such as:
// Build a BitVector of VSRs that can be used for spilling GPRs.

1971 ↗

(On Diff #148042)

Why the 3 register classes? Don't we only want to keep VF/F8 registers in this BitVector?

1976 ↗

(On Diff #148042)

Can you also clear any of the registers that satisfy MF.getRegInfo().isPhysRegUsed(VolatileVFReg) so you:

• Provide early exit if there are no unused caller-saved VSRs
• Simplify the spilling code below (so you don't need to loop and check this there)

1980 ↗

(On Diff #148042)

If you do the above, you can sink this into the loop so we exit as soon as we have used up all the usable VSRs.

1990 ↗

(On Diff #148042)

If you do the above, this whole thing just becomes a call to find_first().

[nemanjai]

LGTM.
Perhaps @thegameg or @MatzeB want to look at it as well (particularly the target independent parts). And perhaps @hfinkel and @echristo are interested in having another look. Otherwise, I suppose this is ready to go.

include/llvm/CodeGen/MachineFrameInfo.h

54

You initialize this in the only constructor. I don't think there's any point in initializing it here.

lib/CodeGen/PrologEpilogInserter.cpp

474

... outside of the region delimited by the prologue/epilogue.
I think that's what this comment is meant to say, isn't it?

[thegameg]

Please add a MIR test for this, it would make the test much clearer.

Other than that and the comment below, it LGTM!

lib/Target/PowerPC/PPCFrameLowering.cpp

1972

for (unsigned Reg : BVAllocatable.set_bits()) should be equivalent here.

[thegameg]

Thank you! A few comments on the MIR test though. Thanks for putting up the test.

llvm/test/CodeGen/MIR/PowerPC/prolog_vec_spills.mir

3 ↗

(On Diff #162899)

You should be able to remove the LLVM IR and keep the MIR only.

Usually the LLVM IR is kept around if there are any references to IR values from MI.

49 ↗

(On Diff #162899)

Is this needed?
I think you can either use BLR8 implicit undef $lr8, implicit undef $rm and get rid of the $r3 and $x3 defs, or:

$lr8 = IMPLICIT_DEF
$rm = IMPLICIT_DEF
BLR8 implicit $lr8, implicit $rm

59 ↗

(On Diff #162899)

Same as the previous function.

71 ↗

(On Diff #162899)

$cr0 = IMPLICIT_DEF

83 ↗

(On Diff #162899)

I would use IMPLICIT_DEFs here as well to express the intent of the test better.

99 ↗

(On Diff #162899)

CHECK-LABEL: name:

102 ↗

(On Diff #162899)

If these are expected to be generated next to each other, I would use CHECK-NEXT: here.

108 ↗

(On Diff #162899)

CHECK-LABEL: name:

[thegameg]

llvm/test/CodeGen/MIR/PowerPC/prolog_vec_spills.mir

18 ↗

(On Diff #163344)

This line is probably not needed.

55 ↗

(On Diff #163344)

Same.

[MatzeB]

LGTM too; some nitpicks below.

llvm/include/llvm/CodeGen/MachineFrameInfo.h

37 ↗

(On Diff #163344)

If DstReg must be a physical register then you could use MCPhysReg instead of unsigned.

54 ↗

(On Diff #163344)

You could use bool SpilledToReg = false; here instead of setting it in the constructor.

llvm/lib/CodeGen/PrologEpilogInserter.cpp

78–79 ↗

(On Diff #163344)

I assume adding the statistics is independent of the main change. It's fine for review like this, but would be nice if you could separate it into a separate commit when committing.

[syzaara]

llvm/include/llvm/CodeGen/MachineFrameInfo.h

37 ↗

(On Diff #163344)

I'm preferring to keep it as unsigned so that it is consistent with unsigned Reg (the other register variable in the class).

If you'd like me to change it, I can go ahead with that. Would you also want me to change the function definitions for unsigned getDstReg() and void setDstReg(unsigned SpillReg).

And what about the users of this function, like:

unsigned SpilledReg = CSI[I].getDstReg();
unsigned CFIRegister = MF.addFrameInst(MCCFIInstruction::createRegister(
    nullptr, MRI->getDwarfRegNum(Reg, true),
    MRI->getDwarfRegNum(SpilledReg, true)));

SpilledReg is passed to getDwarfRegNum as unsigned.

54 ↗

(On Diff #163344)

Also preferring to keep this the same way so its consistent with the rest of the variables in the class being set in the constructor.

[syzaara]

llvm/include/llvm/CodeGen/MachineFrameInfo.h

37 ↗

(On Diff #163344)

@MatzeB can you please provide some suggestions on the question above?

[• wuzish]

llvm/trunk/lib/Target/PowerPC/PPCFrameLowering.cpp

1975 ↗

(On Diff #173347)

I am wondering why it's not enabled by default? @syzaara

include/llvm/CodeGen/MachineFrameInfo.h

 class raw_ostream;
 class MachineFunction;
 class MachineBasicBlock;
 class BitVector;
 class AllocaInst;
 
 /// The CalleeSavedInfo class tracks the information need to locate where a
 /// callee saved register is in the current frame.
+/// Callee saved reg can also be saved to a different register rather than
+/// on the stack by setting DstReg instead of FrameIdx.
 class CalleeSavedInfo {
   unsigned Reg;
+  union {
     int FrameIdx;
+    unsigned DstReg;
+  };
   /// Flag indicating whether the register is actually restored in the epilog.
   /// In most cases, if a register is saved, it is also restored. There are
   /// some situations, though, when this is not the case. For example, the
   /// LR register on ARM is usually saved, but on exit from the function its
   /// saved value may be loaded directly into PC. Since liveness tracking of
   /// physical registers treats callee-saved registers are live outside of
   /// the function, LR would be treated as live-on-exit, even though in these
   /// scenarios it is not. This flag is added to indicate that the saved
   /// register described by this object is not restored in the epilog.
   /// The long-term solution is to model the liveness of callee-saved registers
   /// by implicit uses on the return instructions, however, the required
   /// changes in the ARM backend would be quite extensive.
   bool Restored;
+  /// Flag indicating whether the register is spilled to stack or another
+  /// register.
+  bool SpilledToReg = false;
 
 public:
   explicit CalleeSavedInfo(unsigned R, int FI = 0)
-  : Reg(R), FrameIdx(FI), Restored(true) {}
+  : Reg(R), FrameIdx(FI), Restored(true), SpilledToReg(false) {}
 
   // Accessors.
   unsigned getReg()                        const { return Reg; }
   int getFrameIdx()                        const { return FrameIdx; }
-  void setFrameIdx(int FI)                       { FrameIdx = FI; }
+  unsigned getDstReg()                     const { return DstReg; }
+  void setFrameIdx(int FI) {
+    FrameIdx = FI;
+    SpilledToReg = false;
+  }
+  void setDstReg(unsigned SpillReg) {
+    DstReg = SpillReg;
+    SpilledToReg = true;
+  }
   bool isRestored()                        const { return Restored; }
   void setRestored(bool R)                       { Restored = R; }
+  bool isSpilledToReg()                    const { return SpilledToReg; }
 };
 
 /// The MachineFrameInfo class represents an abstract stack frame until
 /// prolog/epilog code is inserted.  This class is key to allowing stack frame
 /// representation optimizations, such as frame pointer elimination.  It also
 /// allows more mundane (but still important) optimizations, such as reordering
 /// of abstract objects on the stack frame.
 ///
class MachineFrameInfo {
 
   /// This contains the size of the largest call frame if the target uses frame
   /// setup/destroy pseudo instructions (as defined in the TargetFrameInfo
   /// class).  This information is important for frame pointer elimination.
   /// It is only valid during and after prolog/epilog code insertion.
   unsigned MaxCallFrameSize = ~0u;
 
   /// The prolog/epilog code inserter fills in this vector with each
-  /// callee saved register saved in the frame.  Beyond its use by the prolog/
-  /// epilog code inserter, this data used for debug info and exception
-  /// handling.
+  /// callee saved register saved in either the frame or a different
+  /// register.  Beyond its use by the prolog/ epilog code inserter,
+  /// this data is used for debug info and exception handling.
   std::vector<CalleeSavedInfo> CSInfo;
 
   /// Has CSInfo been set yet?
   bool CSIValid = false;
 
   /// References to frame indices which are mapped
   /// into the local frame allocation block. <FrameIdx, LocalOffset>
   SmallVector<std::pair<int, int64_t>, 32> LocalFrameObjects;

lib/CodeGen/PrologEpilogInserter.cpp

 #include <vector>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "prologepilog"
 
 using MBBVector = SmallVector<MachineBasicBlock *, 4>;
 
+STATISTIC(NumLeafFuncWithSpills, "Number of leaf functions with CSRs");
+STATISTIC(NumFuncSeen, "Number of functions seen in PEI");
+
+
 namespace {
 
 class PEI : public MachineFunctionPass {
 public:
   static char ID;
 
   PEI() : MachineFunctionPass(ID) {
     initializePEIPass(*PassRegistry::getPassRegistry());

 }
 
 /// StackObjSet - A set of stack object indexes
 using StackObjSet = SmallSetVector<int, 8>;
 
 /// runOnMachineFunction - Insert prolog/epilog code and replace abstract
 /// frame indexes with appropriate references.
 bool PEI::runOnMachineFunction(MachineFunction &MF) {
+  NumFuncSeen++;
   const Function &F = MF.getFunction();
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
 
   RS = TRI->requiresRegisterScavenging(MF) ? new RegScavenger() : nullptr;
   FrameIndexVirtualScavenging = TRI->requiresFrameIndexScavenging(MF);
   FrameIndexEliminationScavenging = (RS && !FrameIndexVirtualScavenging) ||
     TRI->requiresFrameIndexReplacementScavenging(MF);
if (!TFI->assignCalleeSavedSpillSlots(F, RegInfo, CSI)) {
 
     unsigned NumFixedSpillSlots;
     const TargetFrameLowering::SpillSlot *FixedSpillSlots =
         TFI->getCalleeSavedSpillSlots(NumFixedSpillSlots);
 
     // Now that we know which registers need to be saved and restored, allocate
     // stack slots for them.
     for (auto &CS : CSI) {
+      // If the target has spilled this register to another register, we don't
+      // need to allocate a stack slot.
+      if (CS.isSpilledToReg())
+        continue;
+
       unsigned Reg = CS.getReg();
       const TargetRegisterClass *RC = RegInfo->getMinimalPhysRegClass(Reg);
 
       int FrameIdx;
       if (RegInfo->hasReservedSpillSlot(F, Reg, FrameIdx)) {
         CS.setFrameIdx(FrameIdx);
         continue;
       }
static void updateLiveness(MachineFunction &MF) {
   for (unsigned i = 0, e = CSI.size(); i != e; ++i) {
     for (MachineBasicBlock *MBB : Visited) {
       MCPhysReg Reg = CSI[i].getReg();
       // Add the callee-saved register as live-in.
       // It's killed at the spill.
       if (!MRI.isReserved(Reg) && !MBB->isLiveIn(Reg))
         MBB->addLiveIn(Reg);
     }
+    // If callee-saved register is spilled to another register rather than
+    // spilling to stack, the destination register has to be marked as live for
+    // each MBB between the prologue and epilogue so that it is not clobbered
+    // before it is reloaded in the epilogue. The Visited set contains all
+    // blocks outside the prologue and epilogue.
+    if (CSI[i].isSpilledToReg()) {
+      for (MachineBasicBlock &MBB : MF) {
+        if (Visited.count(&MBB))
+          continue;
+        MCPhysReg DstReg = CSI[i].getDstReg();
+        if (!MBB.isLiveIn(DstReg))
+          MBB.addLiveIn(DstReg);
+      }
     }
   }
 
+}
+
 /// Insert restore code for the callee-saved registers used in the function.
 static void insertCSRSaves(MachineBasicBlock &SaveBlock,
                            ArrayRef<CalleeSavedInfo> CSI) {
   MachineFunction &MF = *SaveBlock.getParent();
   const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
   const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
 
void PEI::spillCalleeSavedRegs(MachineFunction &MF) {
   assignCalleeSavedSpillSlots(MF, SavedRegs, MinCSFrameIndex, MaxCSFrameIndex);
 
   // Add the code to save and restore the callee saved registers.
   if (!F.hasFnAttribute(Attribute::Naked)) {
     MFI.setCalleeSavedInfoValid(true);
 
     std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
     if (!CSI.empty()) {
+      if (!MFI.hasCalls())
+        NumLeafFuncWithSpills++;
+
       for (MachineBasicBlock *SaveBlock : SaveBlocks) {
         insertCSRSaves(*SaveBlock, CSI);
         // Update the live-in information of all the blocks up to the save
         // point.
         updateLiveness(MF);
       }
       for (MachineBasicBlock *RestoreBlock : RestoreBlocks)
         insertCSRRestores(*RestoreBlock, CSI);

lib/Target/PowerPC/PPCFrameLowering.h

public:
   void processFunctionBeforeFrameFinalized(MachineFunction &MF,
                                      RegScavenger *RS = nullptr) const override;
   void addScavengingSpillSlot(MachineFunction &MF, RegScavenger *RS) const;
 
   bool spillCalleeSavedRegisters(MachineBasicBlock &MBB,
                                  MachineBasicBlock::iterator MI,
                                  const std::vector<CalleeSavedInfo> &CSI,
                                  const TargetRegisterInfo *TRI) const override;
+  /// This function will assign callee saved gprs to volatile vector registers
+  /// for prologue spills when applicable. It returns false if there are any
+  /// registers which were not spilled to volatile vector registers.
+  bool
+  assignCalleeSavedSpillSlots(MachineFunction &MF,
+                              const TargetRegisterInfo *TRI,
+                              std::vector<CalleeSavedInfo> &CSI) const override;
 
   MachineBasicBlock::iterator
   eliminateCallFramePseudoInstr(MachineFunction &MF, MachineBasicBlock &MBB,
                                 MachineBasicBlock::iterator I) const override;
 
   bool restoreCalleeSavedRegisters(MachineBasicBlock &MBB,
                                   MachineBasicBlock::iterator MI,
                                   std::vector<CalleeSavedInfo> &CSI,

lib/Target/PowerPC/PPCFrameLowering.cpp

 //===----------------------------------------------------------------------===//
 
 #include "PPCFrameLowering.h"
 #include "PPCInstrBuilder.h"
 #include "PPCInstrInfo.h"
 #include "PPCMachineFunctionInfo.h"
 #include "PPCSubtarget.h"
 #include "PPCTargetMachine.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/Function.h"
 #include "llvm/Target/TargetOptions.h"
 
 using namespace llvm;
 
+#define DEBUG_TYPE "framelowering"
+STATISTIC(NumNoNeedForFrame, "Number of functions without frames");
+STATISTIC(NumPESpillVSR, "Number of spills to vector in prologue");
+STATISTIC(NumPEReloadVSR, "Number of reloads from vector in epilogue");
+
+static cl::opt<bool>
+EnablePEVectorSpills("ppc-enable-pe-vector-spills",
+                     cl::desc("Enable spills in prologue to vector registers."),
+                     cl::init(true), cl::Hidden);
+
 /// VRRegNo - Map from a numbered VR register to its enum value.
 ///
 static const MCPhysReg VRRegNo[] = {
  PPC::V0 , PPC::V1 , PPC::V2 , PPC::V3 , PPC::V4 , PPC::V5 , PPC::V6 , PPC::V7 ,
  PPC::V8 , PPC::V9 , PPC::V10, PPC::V11, PPC::V12, PPC::V13, PPC::V14, PPC::V15,
  PPC::V16, PPC::V17, PPC::V18, PPC::V19, PPC::V20, PPC::V21, PPC::V22, PPC::V23,
  PPC::V24, PPC::V25, PPC::V26, PPC::V27, PPC::V28, PPC::V29, PPC::V30, PPC::V31
 };
 
bool CanUseRedZone = !MFI.hasVarSizedObjects() && // No dynamic alloca.
                        !RegInfo->hasBasePointer(MF); // No special alignment.
 
   // Note: for PPC32 SVR4ABI (Non-DarwinABI), we can still generate stackless
   // code if all local vars are reg-allocated.
   bool FitsInRedZone = FrameSize <= Subtarget.getRedZoneSize();
 
   // Check whether we can skip adjusting the stack pointer (by using red zone)
   if (!DisableRedZone && CanUseRedZone && FitsInRedZone) {
+    NumNoNeedForFrame++;
     // No need for frame
     if (UpdateMF)
       MFI.setStackSize(0);
     return 0;
   }
 
   // Get the maximum call frame size of all the calls.
   unsigned maxCallFrameSize = MFI.getMaxCallFrameSize();
for (unsigned I = 0, E = CSI.size(); I != E; ++I) {
         unsigned CRReg = isELFv2ABI? Reg : (unsigned) PPC::CR2;
         unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createOffset(
             nullptr, MRI->getDwarfRegNum(CRReg, true), 8));
         BuildMI(MBB, MBBI, dl, TII.get(TargetOpcode::CFI_INSTRUCTION))
             .addCFIIndex(CFIIndex);
         continue;
       }
 
+      if (CSI[I].isSpilledToReg()) {
+        unsigned SpilledReg = CSI[I].getDstReg();
+        unsigned CFIRegister = MF.addFrameInst(MCCFIInstruction::createRegister(
+            nullptr, MRI->getDwarfRegNum(Reg, true),
+            MRI->getDwarfRegNum(SpilledReg, true)));
+        BuildMI(MBB, MBBI, dl, TII.get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIRegister);
+      } else {
         int Offset = MFI.getObjectOffset(CSI[I].getFrameIdx());
         unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createOffset(
             nullptr, MRI->getDwarfRegNum(Reg, true), Offset));
         BuildMI(MBB, MBBI, dl, TII.get(TargetOpcode::CFI_INSTRUCTION))
             .addCFIIndex(CFIIndex);
       }
     }
   }
+}
 
 void PPCFrameLowering::emitEpilogue(MachineFunction &MF,
                                     MachineBasicBlock &MBB) const {
   MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
   DebugLoc dl;
 
   if (MBBI != MBB.end())
     dl = MBBI->getDebugLoc();
void PPCFrameLowering::processFunctionBeforeFrameFinalized(MachineFunction &MF,
   }
 
   // General register save area starts right below the Floating-point
   // register save area.
   if (HasGPSaveArea || HasG8SaveArea) {
     // Move general register save area spill slots down, taking into account
     // the size of the Floating-point register save area.
     for (unsigned i = 0, e = GPRegs.size(); i != e; ++i) {
+      if (!GPRegs[i].isSpilledToReg()) {
         int FI = GPRegs[i].getFrameIdx();
-
         MFI.setObjectOffset(FI, LowerBound + MFI.getObjectOffset(FI));
       }
+    }
 
     // Move general register save area spill slots down, taking into account
     // the size of the Floating-point register save area.
     for (unsigned i = 0, e = G8Regs.size(); i != e; ++i) {
+      if (!G8Regs[i].isSpilledToReg()) {
         int FI = G8Regs[i].getFrameIdx();
-
         MFI.setObjectOffset(FI, LowerBound + MFI.getObjectOffset(FI));
       }
+    }
 
     unsigned MinReg =
       std::min<unsigned>(TRI->getEncodingValue(MinGPR),
                          TRI->getEncodingValue(MinG8R));
 
     if (Subtarget.isPPC64()) {
       LowerBound -= (31 - MinReg + 1) * 8;
     } else {
if (MFI.hasVarSizedObjects() || spillsCR(MF) || spillsVRSAVE(MF) ||
 
     // These kinds of spills might need two registers.
     if (spillsCR(MF) || spillsVRSAVE(MF) || HasAlVars)
       RS->addScavengingFrameIndex(MFI.CreateStackObject(Size, Align, false));
 
   }
 }
 
+// This function checks if a callee saved gpr can be spilled to a volatile
+// vector register. This occurs for leaf functions when the option
+// ppc-enable-pe-vector-spills is enabled. If there are any remaining registers
+// which were not spilled to vectors, return false so the target independent
+// code can handle them by assigning a FrameIdx to a stack slot.
+bool PPCFrameLowering::assignCalleeSavedSpillSlots(
+    MachineFunction &MF, const TargetRegisterInfo *TRI,
+    std::vector<CalleeSavedInfo> &CSI) const {
+
+  if (CSI.empty())
+    return true; // Early exit if no callee saved registers are modified!
+
+  // Early exit if cannot spill gprs to volatile vector registers.
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+  if (!EnablePEVectorSpills || MFI.hasCalls() || !Subtarget.hasP9Vector())
+    return false;
+
+  // Build a BitVector of VSRs that can be used for spilling GPRs.
+  BitVector BVAllocatable = TRI->getAllocatableSet(MF);
+  BitVector BVCalleeSaved(TRI->getNumRegs());
+  const PPCRegisterInfo *RegInfo = Subtarget.getRegisterInfo();
+  const MCPhysReg *CSRegs = RegInfo->getCalleeSavedRegs(&MF);
+  for (unsigned i = 0; CSRegs[i]; ++i)
+    BVCalleeSaved.set(CSRegs[i]);
+
+  for (unsigned Reg = BVAllocatable.find_first(); Reg < BVAllocatable.size();
+       Reg = BVAllocatable.find_next(Reg)) {
+    // Set to 0 if the register is not a volatile VF/F8 register, or if it is
+    // used in the function.
+    if (BVCalleeSaved[Reg] ||
+        (!PPC::F8RCRegClass.contains(Reg) &&
+         !PPC::VFRCRegClass.contains(Reg)) ||
+        (MF.getRegInfo().isPhysRegUsed(Reg)))
+      BVAllocatable.reset(Reg);
+  }
+
+  bool AllSpilledToReg = true;
+  for (auto &CS : CSI) {
+    if (BVAllocatable.none())
+      return false;
+
+    unsigned Reg = CS.getReg();
+    if (!PPC::G8RCRegClass.contains(Reg) && !PPC::GPRCRegClass.contains(Reg)) {
+      AllSpilledToReg = false;
+      continue;
+    }
+
+    unsigned VolatileVFReg = BVAllocatable.find_first();
+    if (VolatileVFReg < BVAllocatable.size()) {
+      CS.setDstReg(VolatileVFReg);
+      BVAllocatable.reset(VolatileVFReg);
+    } else {
+      AllSpilledToReg = false;
+    }
+  }
+  return AllSpilledToReg;
+}
+
+
 bool
 PPCFrameLowering::spillCalleeSavedRegisters(MachineBasicBlock &MBB,
                                      MachineBasicBlock::iterator MI,
                                      const std::vector<CalleeSavedInfo> &CSI,
                                      const TargetRegisterInfo *TRI) const {
 
   // Currently, this function only handles SVR4 32- and 64-bit ABIs.
   // Return false otherwise to maintain pre-existing behavior.
if (IsCRField) {
 
         MBB.insert(MI, CRMIB);
         MBB.insert(MI, addFrameReference(BuildMI(*MF, DL, TII.get(PPC::STW))
                                          .addReg(PPC::R12,
                                                  getKillRegState(true)),
                                          CSI[i].getFrameIdx()));
       }
     } else {
+      if (CSI[i].isSpilledToReg()) {
+        NumPESpillVSR++;
+        BuildMI(MBB, MI, DL, TII.get(PPC::MTVSRD), CSI[i].getDstReg())
+          .addReg(Reg, getKillRegState(true));
+      } else {
         const TargetRegisterClass *RC = TRI->getMinimalPhysRegClass(Reg);
         // Use !IsLiveIn for the kill flag.
         // We do not want to kill registers that are live in this function
         // before their use because they will become undefined registers.
         TII.storeRegToStackSlot(MBB, MI, Reg, !IsLiveIn,
                                 CSI[i].getFrameIdx(), RC, TRI);
       }
     }
+  }
   return true;
 }
 
 static void
 restoreCRs(bool isPPC64, bool is31,
            bool CR2Spilled, bool CR3Spilled, bool CR4Spilled,
            MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
            const std::vector<CalleeSavedInfo> &CSI, unsigned CSIIndex) {
if (Reg == PPC::CR2) {
           && !(PPC::CR2 <= Reg && Reg <= PPC::CR4)) {
         bool is31 = needsFP(*MF);
         restoreCRs(Subtarget.isPPC64(), is31,
                    CR2Spilled, CR3Spilled, CR4Spilled,
                    MBB, I, CSI, CSIIndex);
         CR2Spilled = CR3Spilled = CR4Spilled = false;
       }
 
+      if (CSI[i].isSpilledToReg()) {
+        DebugLoc DL;
+        NumPEReloadVSR++;
+        BuildMI(MBB, I, DL, TII.get(PPC::MFVSRD), Reg)
+            .addReg(CSI[i].getDstReg(), getKillRegState(true));
+      } else {
        // Default behavior for non-CR saves.
         const TargetRegisterClass *RC = TRI->getMinimalPhysRegClass(Reg);
-      TII.loadRegFromStackSlot(MBB, I, Reg, CSI[i].getFrameIdx(),
-                               RC, TRI);
+        TII.loadRegFromStackSlot(MBB, I, Reg, CSI[i].getFrameIdx(), RC, TRI);
         assert(I != MBB.begin() &&
                "loadRegFromStackSlot didn't insert any code!");
       }
+    }
 
     // Insert in reverse order.
     if (AtStart)
       I = MBB.begin();
     else {
       I = BeforeI;
       ++I;
     }

test/CodeGen/PowerPC/prolog_vec_spills.ll

+; RUN: llc -verify-machineinstrs -mtriple=powerpc64le-unknown-linux-gnu -ppc-enable-pe-vector-spills -mcpu=pwr9 < %s | FileCheck %s -check-prefix=CHECK
+
+define signext i32 @test1(i32 signext %a, i32 signext %b) #0 {
+entry:
+; CHECK-LABEL: test1
+; CHECK: mtvsrd [[REG1:[0-9]+]], 14
+; CHECK: mtvsrd [[REG2:[0-9]+]], 15
+; CHECK: mtvsrd [[REG3:[0-9]+]], 16
+; CHECK: mffprd   16, [[REG3]]
+; CHECK: mffprd   15, [[REG2]]
+; CHECK: mffprd   14, [[REG1]]
+  %a.addr = alloca i32, align 4
+  %b.addr = alloca i32, align 4
+  %dst = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  store i32 %b, i32* %b.addr, align 4
+  call void asm sideeffect "", "~{v20},~{f14},~{f0}"()
+  %0 = load i32, i32* %a.addr, align 4
+  %1 = load i32, i32* %b.addr, align 4
+  %2 = call i32 asm "add $0, $1, $2", "=r,r,r,~{r14},~{r15},~{r16}"(i32 %0, i32 %1)
+  store i32 %2, i32* %dst, align 4
+  %3 = load i32, i32* %dst, align 4
+  ret i32 %3
+}
+
+; Test with multiple basic blocks.
+define signext i32 @test2(i32 signext %a, i32 signext %b) {
+entry:
+; CHECK-LABEL: test2
+; CHECK: mtvsrd [[REG1:[0-9]+]], 14
+; CHECK: mtvsrd [[REG2:[0-9]+]], 15
+; CHECK: mtvsrd [[REG3:[0-9]+]], 16
+; CHECK: mffprd   16, [[REG3]]
+; CHECK: mffprd   15, [[REG2]]
+; CHECK: mffprd   14, [[REG1]]
+  %0 = tail call i32 asm "add $0, $1, $2", "=r,r,r,~{r14},~{r15},~{r16}"(i32 %a, i32 %b)
+  %cmp = icmp sgt i32 %b, %a
+  %cmp112 = icmp sgt i32 %b, 0
+  %or.cond = and i1 %cmp, %cmp112
+  br i1 %or.cond, label %for.body.lr.ph, label %if.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  %add = sub i32 %b, %a
+  %1 = add i32 %0, %b
+  %2 = add i32 %b, -1
+  %3 = mul i32 %add, %2
+  %4 = add i32 %1, %3
+  %5 = sub i32 %4, %a
+  br label %if.end
+
+if.end:                                           ; preds = %for.body.lr.ph, %entry
+  %dst.1 = phi i32 [ %0, %entry ], [ %5, %for.body.lr.ph ]
+  ret i32 %dst.1
+}
+
+attributes #0 = { nounwind }