[AMDGPU][SIFrameLowering] Unify PEI SGPR spill saves and restores

[cdevadas]

Summary

There is a lot of customization and eventually code duplication in the
frame lowering that handles special SGPR spills like the one needed for
the Frame Pointer. Incorporating any additional SGPR spill currently
makes it difficult during PEI. This patch introduces a new spill builder
to efficiently handle such spill requirements. Various spill methods are
special handled using a separate class.

[sebastian-ne]

Looks ok to me in general.

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

31

This function doesn’t do anything with CSRs, although the name says so?

99–100

Comment should say <FI> for the renamed variable.

271

Why does this code split the copies into SubRegs and copyFromScratchSGPR doesn’t?

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

37–38 ↗

(On Diff #454702)

Why not make this an enum class RegSaveKind?

[cdevadas]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

31

You're right. Neither scratch nor CSR should be in the name. I will change it.

99–100

Thanks. Will change it.

271

For the purpose of emitting CFI directives for each copy. There is no CFI emitted in the epilogue.
There is more coming in the downstream version of this patch.

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

37–38 ↗

(On Diff #454702)

Will do.

[cdevadas]

Addressed review comments.

[sebastian-ne]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

271

Hm, does that mean CFI cannot represent register pairs?
If so, can we emit a single copy instruction for the super register and multiple CFI instructions for the subregisters?

[cdevadas]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

271

Sure, we will get the benefit for 64-bit copies as we have the 'S_MOV_B64' instruction. But for the CFI emission, we need composite expressions for register pairs. @scott.linder to comment about its impact.
For now, I will change the code to emit a single COPY instruction for the register pair while retaining separate CFI instructions for equivalent SubRegs.

[scott.linder]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

271

I agree with using the most efficient representation of the copy (sounds like a 64-bit MOV is best) and generating whatever CFI we need to describe it.

In the case of the 64-bit program address space PC and the wave64 EXEC this will be a pair of SGPRs, described via a composite expression in the CFI, like how CD mentions. For every other case we can generate independent CFI instructions to separately describe the SubRegs which correspond to DWARF registers.

For example, we don't have a DWARF register number for the pair SGPR40_SGPR41, whereas we do for SGPR40 and SGPR41 separately, so we can still emit something along the lines of:

$sgpr40_sgpr41 = S_MOV_B64 $sgpr50_sgpr51
CFI_INSTRUCTION register sgpr50, sgpr40
CFI_INSTRUCTION register sgpr51, sgpr41

[cdevadas]

Restricted the code to emit a single COPY in case of copy to scratch SGPR. The related CFIs still have split components per sub reg.

[sebastian-ne]

Thanks, looks good to me.

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

53–54

No else after return

[cdevadas]

Removed else after the return statement.

[arsenm]

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

30–37 ↗

(On Diff #462456)

Don't see why this needs to pollute SIRegisterInfo. Belongs in MFI or maybe FrameLowering

[cdevadas]

Moved CustomSGPRSaveInfo class into MFI

[cdevadas]

Used PrologEpilog instead of the prefix "Custom" to refer SGPR spills during PEI. Also, fixed the review title and summary.

[arsenm]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

195

Really we ought to be stepping back from the end of the entry block and always using reverse liveness (but I guess you're just moving this function, so that's a separate change)

[cdevadas]

Rebase.

[cdevadas]

code rebase

llvm/lib/Target/AMDGPU/SIFrameLowering.h

void emitEpilogue(MachineFunction &MF,
                     MachineBasicBlock &MBB) const override;
   StackOffset getFrameIndexReference(const MachineFunction &MF, int FI,
                                      Register &FrameReg) const override;
 
   void determineCalleeSaves(MachineFunction &MF, BitVector &SavedRegs,
                             RegScavenger *RS = nullptr) const override;
   void determineCalleeSavesSGPR(MachineFunction &MF, BitVector &SavedRegs,
                                 RegScavenger *RS = nullptr) const;
+  void determineCustomSGPRSaves(MachineFunction &MF,
+                                BitVector &SavedRegs) const;
   bool
   assignCalleeSavedSpillSlots(MachineFunction &MF,
                               const TargetRegisterInfo *TRI,
                               std::vector<CalleeSavedInfo> &CSI) const override;
 
   bool allocateScavengingFrameIndexesNearIncomingSP(
     const MachineFunction &MF) const override;
 

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

 #define DEBUG_TYPE "frame-info"
 
 static cl::opt<bool> EnableSpillVGPRToAGPR(
   "amdgpu-spill-vgpr-to-agpr",
   cl::desc("Enable spilling VGPRs to AGPRs"),
   cl::ReallyHidden,
   cl::init(true));
 
+// Find a register matching \p RC from \p LiveRegs which is unused and available
+// throughout the function. On failure, returns AMDGPU::NoRegister.
+static MCRegister findUnusedRegister(MachineRegisterInfo &MRI,
+                                     const LivePhysRegs &LiveRegs,
+                                     const TargetRegisterClass &RC) {
+  for (MCRegister Reg : RC) {
+    if (!MRI.isPhysRegUsed(Reg) && LiveRegs.available(MRI, Reg))
+      return Reg;
+  }
+  return MCRegister();
+}
+
 // Find a scratch register that we can use in the prologue. We avoid using
 // callee-save registers since they may appear to be free when this is called
 // from canUseAsPrologue (during shrink wrapping), but then no longer be free
 // when this is called from emitPrologue.
 static MCRegister findScratchNonCalleeSaveRegister(MachineRegisterInfo &MRI,
                                                    LivePhysRegs &LiveRegs,
                                                    const TargetRegisterClass &RC,
                                                    bool Unused = false) {
   // Mark callee saved registers as used so we will not choose them.
   const MCPhysReg *CSRegs = MRI.getCalleeSavedRegs();
   for (unsigned i = 0; CSRegs[i]; ++i)
     LiveRegs.addReg(CSRegs[i]);
 
-  if (Unused) {
-    // We are looking for a register that can be used throughout the entire
-    // function, so any use is unacceptable.
-    for (MCRegister Reg : RC) {
-      if (!MRI.isPhysRegUsed(Reg) && LiveRegs.available(MRI, Reg))
-        return Reg;
-    }
-  } else {
-    for (MCRegister Reg : RC) {
-      if (LiveRegs.available(MRI, Reg))
-        return Reg;
-    }
-  }
+  if (Unused)
+    return findUnusedRegister(MRI, LiveRegs, RC);
+
 
   return MCRegister();
 }
 
-static void getVGPRSpillLaneOrTempRegister(MachineFunction &MF,
-                                           LivePhysRegs &LiveRegs,
-                                           Register &TempSGPR,
-                                           Optional<int> &FrameIndex,
-                                           bool IsFP) {
+static void getVGPRSpillLaneOrTempRegister(
+    MachineFunction &MF, LivePhysRegs &LiveRegs, Register SGPR,
+    const TargetRegisterClass &RC = AMDGPU::SReg_32_XM0_XEXECRegClass) {
   SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   MachineFrameInfo &FrameInfo = MF.getFrameInfo();
 
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
+  unsigned Size = TRI->getSpillSize(RC);
+  Align Alignment = TRI->getSpillAlign(RC);
 
-  // We need to save and restore the current FP/BP.
+  // We need to save and restore the given SGPR.
 
-  // 1: Try to save the FP/BP in an unused SGPR.
-  TempSGPR = findScratchNonCalleeSaveRegister(
-      MF.getRegInfo(), LiveRegs, AMDGPU::SReg_32_XM0_XEXECRegClass, true);
+  // 1: Try to save the given register into an unused scratch SGPR. The LiveRegs
+  // should have all the callee saved registers marked as used.
+  Register ScratchSGPR = findUnusedRegister(MF.getRegInfo(), LiveRegs, RC);
 
-  if (!TempSGPR) {
-    int NewFI = FrameInfo.CreateStackObject(4, Align(4), true, nullptr,
-                                            TargetStackID::SGPRSpill);
+  if (!ScratchSGPR) {
+    int FI = FrameInfo.CreateStackObject(Size, Alignment, true, nullptr,
 
     if (TRI->spillSGPRToVGPR() &&
-        MFI->allocateSGPRSpillToVGPR(MF, NewFI, /* IsPEI */ true)) {
-      // 2: There's no free lane to spill, and no free register to save FP/BP,
-      // so we're forced to spill another VGPR to use for the spill.
-      FrameIndex = NewFI;
-
-      LLVM_DEBUG(
-          auto Spill = MFI->getSGPRToVGPRCustomSpills(NewFI).front();
-          dbgs() << (IsFP ? "FP" : "BP") << " requires fallback spill to "
-                 << printReg(Spill.VGPR, TRI) << ':' << Spill.Lane << '\n';);
+        MFI->allocateSGPRSpillToVGPR(MF, FI, /* IsPEI */ true)) {
+      // 2: There's no free lane to spill, and no free register to save the
+      // SGPR, so we're forced to take another VGPR to use for the spill.
+      MFI->addToCustomSGPRSpills(
+          SGPR, CustomSGPRSaveInfo(SGPRSaveKind::SPILL_TO_VGPR_LANE, FI));
+
+      LLVM_DEBUG(auto Spill = MFI->getSGPRToVGPRCustomSpills(FI).front();
+                 dbgs() << printReg(SGPR, TRI) << " requires fallback spill to "
+                        << printReg(Spill.VGPR, TRI) << ':' << Spill.Lane
+                        << '\n';);
     } else {
-      // Remove dead <NewFI> index
-      MF.getFrameInfo().RemoveStackObject(NewFI);
-      // 3: If all else fails, spill the FP/BP to memory.
-      FrameIndex = FrameInfo.CreateSpillStackObject(4, Align(4));
-      LLVM_DEBUG(dbgs() << "Reserved FI " << FrameIndex << " for spilling "
-                        << (IsFP ? "FP" : "BP") << '\n');
+      // Remove dead <FI> index
+      MF.getFrameInfo().RemoveStackObject(FI);
+      // 3: If all else fails, spill the register to memory.
+      FI = FrameInfo.CreateSpillStackObject(Size, Alignment);
+      MFI->addToCustomSGPRSpills(
+          SGPR, CustomSGPRSaveInfo(SGPRSaveKind::SPILL_TO_MEM, FI));
+      LLVM_DEBUG(dbgs() << "Reserved FI " << FI << " for spilling "
+                        << printReg(SGPR, TRI) << '\n');
     }
   } else {
-    LLVM_DEBUG(dbgs() << "Saving " << (IsFP ? "FP" : "BP") << " with copy to "
-                      << printReg(TempSGPR, TRI) << '\n');
+    MFI->addToCustomSGPRSpills(
+        SGPR,
+        CustomSGPRSaveInfo(SGPRSaveKind::COPY_TO_SCRATCH_SGPR, ScratchSGPR));
+    LiveRegs.addReg(ScratchSGPR);
+    LLVM_DEBUG(dbgs() << "Saving " << printReg(SGPR, TRI) << " with copy to "
+                      << printReg(ScratchSGPR, TRI) << '\n');
   }
 }
 
 // We need to specially emit stack operations here because a different frame
 // register is used than in the rest of the function, as getFrameRegister would
 // use.
 static void buildPrologSpill(const GCNSubtarget &ST, const SIRegisterInfo &TRI,
                              const SIMachineFunctionInfo &FuncInfo,
                              LivePhysRegs &LiveRegs, MachineFunction &MF,
                              MachineBasicBlock &MBB,
                              MachineBasicBlock::iterator I, const DebugLoc &DL,
-                             Register SpillReg, int FI) {
+                             Register SpillReg, int FI, int64_t DwordOff = 0) {
   unsigned Opc = ST.enableFlatScratch() ? AMDGPU::SCRATCH_STORE_DWORD_SADDR
                                         : AMDGPU::BUFFER_STORE_DWORD_OFFSET;
 
   MachineFrameInfo &FrameInfo = MF.getFrameInfo();
   MachinePointerInfo PtrInfo = MachinePointerInfo::getFixedStack(MF, FI);
   MachineMemOperand *MMO = MF.getMachineMemOperand(
       PtrInfo, MachineMemOperand::MOStore, FrameInfo.getObjectSize(FI),
       FrameInfo.getObjectAlign(FI));
   LiveRegs.addReg(SpillReg);
   bool IsKill = !MBB.isLiveIn(SpillReg);
   TRI.buildSpillLoadStore(MBB, I, DL, Opc, FI, SpillReg, IsKill,
-                          FuncInfo.getStackPtrOffsetReg(), 0, MMO, nullptr,
-                          &LiveRegs);
+                          FuncInfo.getStackPtrOffsetReg(), DwordOff, MMO,
+                          nullptr, &LiveRegs);
   if (IsKill)
     LiveRegs.removeReg(SpillReg);
 }
 
-static void buildEpilogRestore(const GCNSubtarget &ST,
-                               const SIRegisterInfo &TRI,
-                               const SIMachineFunctionInfo &FuncInfo,
-                               LivePhysRegs &LiveRegs, MachineFunction &MF,
-                               MachineBasicBlock &MBB,
-                               MachineBasicBlock::iterator I,
-                               const DebugLoc &DL, Register SpillReg, int FI) {
+static void buildEpilogRestore(
+    const GCNSubtarget &ST, const SIRegisterInfo &TRI,
+    const SIMachineFunctionInfo &FuncInfo, LivePhysRegs &LiveRegs,
+    MachineFunction &MF, MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
+    const DebugLoc &DL, Register SpillReg, int FI, int64_t DwordOff = 0) {
   unsigned Opc = ST.enableFlatScratch() ? AMDGPU::SCRATCH_LOAD_DWORD_SADDR
                                         : AMDGPU::BUFFER_LOAD_DWORD_OFFSET;
 
   MachineFrameInfo &FrameInfo = MF.getFrameInfo();
   MachinePointerInfo PtrInfo = MachinePointerInfo::getFixedStack(MF, FI);
   MachineMemOperand *MMO = MF.getMachineMemOperand(
       PtrInfo, MachineMemOperand::MOLoad, FrameInfo.getObjectSize(FI),
       FrameInfo.getObjectAlign(FI));
   TRI.buildSpillLoadStore(MBB, I, DL, Opc, FI, SpillReg, false,
-                          FuncInfo.getStackPtrOffsetReg(), 0, MMO, nullptr,
-                          &LiveRegs);
+                          FuncInfo.getStackPtrOffsetReg(), DwordOff, MMO,
+                          nullptr, &LiveRegs);
 }
 
 static void buildGitPtr(MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
                         const DebugLoc &DL, const SIInstrInfo *TII,
                         Register TargetReg) {
   MachineFunction *MF = MBB.getParent();
   const SIMachineFunctionInfo *MFI = MF->getInfo<SIMachineFunctionInfo>();
   const SIRegisterInfo *TRI = &TII->getRegisterInfo();
static void buildGitPtr(MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
   }
   Register GitPtrLo = MFI->getGITPtrLoReg(*MF);
   MF->getRegInfo().addLiveIn(GitPtrLo);
   MBB.addLiveIn(GitPtrLo);
   BuildMI(MBB, I, DL, SMovB32, TargetLo)
     .addReg(GitPtrLo);
 }
 
+static void initLiveRegs(LivePhysRegs &LiveRegs, const SIRegisterInfo &TRI,
+                         const SIMachineFunctionInfo *FuncInfo,
+                         MachineFunction &MF, MachineBasicBlock &MBB,
+                         MachineBasicBlock::iterator MBBI, bool IsProlog) {
+  if (LiveRegs.empty()) {
+    LiveRegs.init(TRI);
+    if (IsProlog) {
+      LiveRegs.addLiveIns(MBB);
+    } else {
+      // In epilog.
+      LiveRegs.addLiveOuts(MBB);
+      LiveRegs.stepBackward(*MBBI);
+    }
+  }
+}
+
+namespace llvm {
+
+// SpillBuilder for saving/restoring custom SGPR spills.
+// Custom spills are those special SGPR spills delayed until the current
+// function's frame is finalized. The spills for FP, BP, etc. come under this
+// category. For a given register, the builder uses the CustomSGPRSaveInfo to
+// decide the spill method.
+class CustomSGPRSpillBuilder {
+  MachineBasicBlock::iterator MI;
+  MachineBasicBlock &MBB;
+  MachineFunction &MF;
+  const GCNSubtarget &ST;
+  MachineFrameInfo &MFI;
+  SIMachineFunctionInfo *FuncInfo;
+  const SIInstrInfo *TII;
+  const SIRegisterInfo &TRI;
+  Register SuperReg;
+  const CustomSGPRSaveInfo SI;
+  LivePhysRegs &LiveRegs;
+  const DebugLoc &DL;
+  ArrayRef<int16_t> SplitParts;
+  unsigned NumSubRegs;
+  unsigned EltSize = 4;
+
+  void saveToMemory(const int FI) const {
+    MachineRegisterInfo &MRI = MF.getRegInfo();
+    assert(!MFI.isDeadObjectIndex(FI));
+
+    initLiveRegs(LiveRegs, TRI, FuncInfo, MF, MBB, MI, /*IsProlog*/ true);
+
+    MCPhysReg TmpVGPR = findScratchNonCalleeSaveRegister(
+        MRI, LiveRegs, AMDGPU::VGPR_32RegClass);
+    if (!TmpVGPR)
+      report_fatal_error("failed to find free scratch register");
+
+    for (unsigned I = 0, DwordOff = 0; I < NumSubRegs; ++I) {
+      Register SubReg = NumSubRegs == 1
+                            ? SuperReg
+                            : Register(TRI.getSubReg(SuperReg, SplitParts[I]));
+      BuildMI(MBB, MI, DL, TII->get(AMDGPU::V_MOV_B32_e32), TmpVGPR)
+          .addReg(SubReg);
+
+      buildPrologSpill(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MI, DL, TmpVGPR,
+                       FI, DwordOff);
+      DwordOff += 4;
+    }
+  }
+
+  void saveToVGPRLane(const int FI) const {
+    assert(!MFI.isDeadObjectIndex(FI));
+
+    assert(MFI.getStackID(FI) == TargetStackID::SGPRSpill);
+    ArrayRef<SIRegisterInfo::SpilledReg> Spill =
+        FuncInfo->getSGPRToVGPRCustomSpills(FI);
+    assert(Spill.size() == NumSubRegs);
+
+    for (unsigned I = 0; I < NumSubRegs; ++I) {
+      Register SubReg = NumSubRegs == 1
+                            ? SuperReg
+                            : Register(TRI.getSubReg(SuperReg, SplitParts[I]));
+      BuildMI(MBB, MI, DL, TII->get(AMDGPU::V_WRITELANE_B32), Spill[I].VGPR)
+          .addReg(SubReg)
+          .addImm(Spill[I].Lane)
+          .addReg(Spill[I].VGPR, RegState::Undef);
+    }
+  }
+
+  void copyToScratchSGPR(Register DstReg) const {
+    BuildMI(MBB, MI, DL, TII->get(AMDGPU::COPY), DstReg)
+        .addReg(SuperReg)
+        .setMIFlag(MachineInstr::FrameSetup);
+  }
+
+  void restoreFromMemory(const int FI) {
+    MachineRegisterInfo &MRI = MF.getRegInfo();
+
+    initLiveRegs(LiveRegs, TRI, FuncInfo, MF, MBB, MI, /*IsProlog*/ false);
+    MCPhysReg TmpVGPR = findScratchNonCalleeSaveRegister(
+        MRI, LiveRegs, AMDGPU::VGPR_32RegClass);
+    if (!TmpVGPR)
+      report_fatal_error("failed to find free scratch register");
+
+    for (unsigned I = 0, DwordOff = 0; I < NumSubRegs; ++I) {
+      Register SubReg = NumSubRegs == 1
+                            ? SuperReg
+                            : Register(TRI.getSubReg(SuperReg, SplitParts[I]));
+
+      buildEpilogRestore(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MI, DL, TmpVGPR,
+                         FI, DwordOff);
+      BuildMI(MBB, MI, DL, TII->get(AMDGPU::V_READFIRSTLANE_B32), SubReg)
+          .addReg(TmpVGPR, RegState::Kill);
+      DwordOff += 4;
+    }
+  }
+
+  void restoreFromVGPRLane(const int FI) {
+    assert(MFI.getStackID(FI) == TargetStackID::SGPRSpill);
+    ArrayRef<SIRegisterInfo::SpilledReg> Spill =
+        FuncInfo->getSGPRToVGPRCustomSpills(FI);
+    assert(Spill.size() == NumSubRegs);
+
+    for (unsigned I = 0; I < NumSubRegs; ++I) {
+      Register SubReg = NumSubRegs == 1
+                            ? SuperReg
+                            : Register(TRI.getSubReg(SuperReg, SplitParts[I]));
+      BuildMI(MBB, MI, DL, TII->get(AMDGPU::V_READLANE_B32), SubReg)
+          .addReg(Spill[I].VGPR)
+          .addImm(Spill[I].Lane);
+    }
+  }
+
+  void copyFromScratchSGPR(Register SrcReg) const {
+    BuildMI(MBB, MI, DL, TII->get(AMDGPU::COPY), SuperReg)
+        .addReg(SrcReg)
+        .setMIFlag(MachineInstr::FrameDestroy);
+  }
+
+public:
+  CustomSGPRSpillBuilder(Register Reg, const CustomSGPRSaveInfo SI,
+                         MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
+                         const DebugLoc &DL, const SIInstrInfo *TII,
+                         const SIRegisterInfo &TRI, LivePhysRegs &LiveRegs)
+      : MI(MI), MBB(MBB), MF(*MBB.getParent()),
+        ST(MF.getSubtarget<GCNSubtarget>()), MFI(MF.getFrameInfo()),
+        FuncInfo(MF.getInfo<SIMachineFunctionInfo>()), TII(TII), TRI(TRI),
+        SuperReg(Reg), SI(SI), LiveRegs(LiveRegs), DL(DL) {
+    const TargetRegisterClass *RC = TRI.getPhysRegClass(SuperReg);
+    SplitParts = TRI.getRegSplitParts(RC, EltSize);
+    NumSubRegs = SplitParts.empty() ? 1 : SplitParts.size();
+
+    assert(SuperReg != AMDGPU::M0 && "m0 should never spill");
+  }
+
+  void save() {
+    switch (SI.getKind()) {
+    case SGPRSaveKind::SPILL_TO_MEM:
+      return saveToMemory(SI.getIndex());
+    case SGPRSaveKind::SPILL_TO_VGPR_LANE:
+      return saveToVGPRLane(SI.getIndex());
+    case SGPRSaveKind::COPY_TO_SCRATCH_SGPR:
+      return copyToScratchSGPR(SI.getReg());
+    }
+  }
+
+  void restore() {
+    switch (SI.getKind()) {
+    case SGPRSaveKind::SPILL_TO_MEM:
+      return restoreFromMemory(SI.getIndex());
+    case SGPRSaveKind::SPILL_TO_VGPR_LANE:
+      return restoreFromVGPRLane(SI.getIndex());
+    case SGPRSaveKind::COPY_TO_SCRATCH_SGPR:
+      return copyFromScratchSGPR(SI.getReg());
+    }
+  }
+};
+
+} // namespace llvm
+
 // Emit flat scratch setup code, assuming `MFI->hasFlatScratchInit()`
 void SIFrameLowering::emitEntryFunctionFlatScratchInit(
     MachineFunction &MF, MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
     const DebugLoc &DL, Register ScratchWaveOffsetReg) const {
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIInstrInfo *TII = ST.getInstrInfo();
   const SIRegisterInfo *TRI = &TII->getRegisterInfo();
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
case TargetStackID::SGPRSpill:
     return true;
   case TargetStackID::ScalableVector:
   case TargetStackID::WasmLocal:
     return false;
   }
   llvm_unreachable("Invalid TargetStackID::Value");
 }
 
-static void initLiveRegs(LivePhysRegs &LiveRegs, const SIRegisterInfo &TRI,
-                         const SIMachineFunctionInfo *FuncInfo,
-                         MachineFunction &MF, MachineBasicBlock &MBB,
-                         MachineBasicBlock::iterator MBBI, bool IsProlog) {
-  if (LiveRegs.empty()) {
-    LiveRegs.init(TRI);
-    if (IsProlog) {
-      LiveRegs.addLiveIns(MBB);
-    } else {
-      // In epilog.
-      LiveRegs.addLiveOuts(MBB);
-      LiveRegs.stepBackward(*MBBI);
-    }
-  }
-}
-
 // Activate all lanes, returns saved exec.
 static Register buildScratchExecCopy(LivePhysRegs &LiveRegs,
                                      MachineFunction &MF,
                                      MachineBasicBlock &MBB,
                                      MachineBasicBlock::iterator MBBI,
                                      const DebugLoc &DL, bool IsProlog) {
   Register ScratchExecCopy;
   MachineRegisterInfo &MRI = MF.getRegInfo();
const unsigned OrSaveExec =
       ST.isWave32() ? AMDGPU::S_OR_SAVEEXEC_B32 : AMDGPU::S_OR_SAVEEXEC_B64;
   auto SaveExec = BuildMI(MBB, MBBI, DL, TII->get(OrSaveExec), ScratchExecCopy)
     .addImm(-1);
   SaveExec->getOperand(3).setIsDead(); // Mark SCC as dead.
 
   return ScratchExecCopy;
 }
 
-// A StackID of SGPRSpill implies that this is a spill from SGPR to VGPR.
-// Otherwise we are spilling to memory.
-static bool spilledToMemory(const MachineFunction &MF, int SaveIndex) {
-  const MachineFrameInfo &MFI = MF.getFrameInfo();
-  return MFI.getStackID(SaveIndex) != TargetStackID::SGPRSpill;
-}
-
 void SIFrameLowering::emitPrologue(MachineFunction &MF,
                                    MachineBasicBlock &MBB) const {
   SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
   if (FuncInfo->isEntryFunction()) {
     emitEntryFunctionPrologue(MF, MBB);
     return;
   }
 
   MachineFrameInfo &MFI = MF.getFrameInfo();
-  MachineRegisterInfo &MRI = MF.getRegInfo();
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIInstrInfo *TII = ST.getInstrInfo();
   const SIRegisterInfo &TRI = TII->getRegisterInfo();
 
   Register StackPtrReg = FuncInfo->getStackPtrOffsetReg();
   Register FramePtrReg = FuncInfo->getFrameOffsetReg();
   Register BasePtrReg =
       TRI.hasBasePointer(MF) ? TRI.getBaseRegister() : Register();
   LivePhysRegs LiveRegs;
 
   MachineBasicBlock::iterator MBBI = MBB.begin();
   // DebugLoc must be unknown since the first instruction with DebugLoc is used
   // to determine the end of the prologue.
   DebugLoc DL;
 
   bool HasFP = false;
   bool HasBP = false;
   uint32_t NumBytes = MFI.getStackSize();
   uint32_t RoundedSize = NumBytes;
   // To avoid clobbering VGPRs in lanes that weren't active on function entry,
   // turn on all lanes before doing the spill to memory.
   Register ScratchExecCopy;
 
-  Optional<int> FPSaveIndex = FuncInfo->FramePointerSaveIndex;
-  Optional<int> BPSaveIndex = FuncInfo->BasePointerSaveIndex;
-
   // Spill Whole-Wave Mode VGPRs.
   for (const auto &Reg : FuncInfo->getWWMSpills()) {
     Register VGPR = Reg.first;
     int FI = Reg.second;
     if (!ScratchExecCopy)
       ScratchExecCopy = buildScratchExecCopy(LiveRegs, MF, MBB, MBBI, DL,
                                              /*IsProlog*/ true);
 
     buildPrologSpill(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MBBI, DL, VGPR, FI);
   }
 
   if (ScratchExecCopy) {
     // FIXME: Split block and make terminator.
     unsigned ExecMov = ST.isWave32() ? AMDGPU::S_MOV_B32 : AMDGPU::S_MOV_B64;
     MCRegister Exec = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
     BuildMI(MBB, MBBI, DL, TII->get(ExecMov), Exec)
         .addReg(ScratchExecCopy, RegState::Kill);
     LiveRegs.addReg(ScratchExecCopy);
   }
 
-  auto SaveSGPRToMemory = [&](Register Reg, const int FI) {
-    assert(!MFI.isDeadObjectIndex(FI));
-
-    initLiveRegs(LiveRegs, TRI, FuncInfo, MF, MBB, MBBI, /*IsProlog*/ true);
-
-    MCPhysReg TmpVGPR = findScratchNonCalleeSaveRegister(
-        MRI, LiveRegs, AMDGPU::VGPR_32RegClass);
-    if (!TmpVGPR)
-      report_fatal_error("failed to find free scratch register");
-
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::V_MOV_B32_e32), TmpVGPR)
-        .addReg(Reg);
-
-    buildPrologSpill(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MBBI, DL, TmpVGPR,
-                     FI);
-  };
-
-  auto SaveSGPRToVGPRLane = [&](Register Reg, const int FI) {
-    assert(!MFI.isDeadObjectIndex(FI));
-
-    assert(MFI.getStackID(FI) == TargetStackID::SGPRSpill);
-    ArrayRef<SIRegisterInfo::SpilledReg> Spill =
-        FuncInfo->getSGPRToVGPRCustomSpills(FI);
-    assert(Spill.size() == 1);
-
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::V_WRITELANE_B32), Spill[0].VGPR)
-        .addReg(Reg)
-        .addImm(Spill[0].Lane)
-        .addReg(Spill[0].VGPR, RegState::Undef);
-  };
-
-  if (FPSaveIndex) {
-    if (spilledToMemory(MF, *FPSaveIndex))
-      SaveSGPRToMemory(FramePtrReg, *FPSaveIndex);
-    else
-      SaveSGPRToVGPRLane(FramePtrReg, *FPSaveIndex);
-  }
-
-  // Emit the copy if we need an FP, and are using a free SGPR to save it.
-  if (FuncInfo->SGPRForFPSaveRestoreCopy) {
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::COPY),
-            FuncInfo->SGPRForFPSaveRestoreCopy)
-        .addReg(FramePtrReg)
-        .setMIFlag(MachineInstr::FrameSetup);
+  for (const auto &Spill : FuncInfo->getCustomSGPRSpills()) {
+    CustomSGPRSpillBuilder CSB(Spill.first, Spill.second, MBB, MBBI, DL, TII,
+                               TRI, LiveRegs);
+    CSB.save();
   }
 
-  if (BPSaveIndex) {
-    if (spilledToMemory(MF, *BPSaveIndex))
-      SaveSGPRToMemory(BasePtrReg, *BPSaveIndex);
-    else
-      SaveSGPRToVGPRLane(BasePtrReg, *BPSaveIndex);
-  }
-
-  // Emit the copy if we need a BP, and are using a free SGPR to save it.
-  if (FuncInfo->SGPRForBPSaveRestoreCopy) {
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::COPY),
-            FuncInfo->SGPRForBPSaveRestoreCopy)
-        .addReg(BasePtrReg)
-        .setMIFlag(MachineInstr::FrameSetup);
-  }
-
-  // If a copy has been emitted for FP and/or BP, Make the SGPRs
-  // used in the copy instructions live throughout the function.
-  SmallVector<MCPhysReg, 2> TempSGPRs;
-  if (FuncInfo->SGPRForFPSaveRestoreCopy)
-    TempSGPRs.push_back(FuncInfo->SGPRForFPSaveRestoreCopy);
-
-  if (FuncInfo->SGPRForBPSaveRestoreCopy)
-    TempSGPRs.push_back(FuncInfo->SGPRForBPSaveRestoreCopy);
-
-  if (!TempSGPRs.empty()) {
+  // If a copy to scratch SGPR has been chosen for any of the custom SGPR
+  // spills, make such scratch registers live throughout the function.
+  SmallVector<Register, 1> ScratchSGPRs;
+  FuncInfo->getAllScratchSGPRCopyDstRegs(ScratchSGPRs);
+  if (!ScratchSGPRs.empty()) {
     for (MachineBasicBlock &MBB : MF) {
-      for (MCPhysReg Reg : TempSGPRs)
+      for (MCPhysReg Reg : ScratchSGPRs)
         MBB.addLiveIn(Reg);
 
       MBB.sortUniqueLiveIns();
     }
     if (!LiveRegs.empty()) {
-      LiveRegs.addReg(FuncInfo->SGPRForFPSaveRestoreCopy);
-      LiveRegs.addReg(FuncInfo->SGPRForBPSaveRestoreCopy);
+      for (MCPhysReg Reg : ScratchSGPRs)
+        LiveRegs.addReg(Reg);
     }
   }
 
   if (TRI.hasStackRealignment(MF)) {
     HasFP = true;
     const unsigned Alignment = MFI.getMaxAlign().value();
 
     RoundedSize += Alignment;
void SIFrameLowering::emitPrologue(MachineFunction &MF,
   if (HasFP && RoundedSize != 0) {
     auto Add = BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::S_ADD_I32), StackPtrReg)
         .addReg(StackPtrReg)
         .addImm(RoundedSize * getScratchScaleFactor(ST))
         .setMIFlag(MachineInstr::FrameSetup);
     Add->getOperand(3).setIsDead(); // Mark SCC as dead.
   }
 
-  assert((!HasFP || (FuncInfo->SGPRForFPSaveRestoreCopy ||
-                     FuncInfo->FramePointerSaveIndex)) &&
+  bool FPSaved = FuncInfo->hasCustomSGPRSpillEntry(FramePtrReg);
+  assert((!HasFP || FPSaved) &&
          "Needed to save FP but didn't save it anywhere");
 
   // If we allow spilling to AGPRs we may have saved FP but then spill
   // everything into AGPRs instead of the stack.
-  assert((HasFP || (!FuncInfo->SGPRForFPSaveRestoreCopy &&
-                    !FuncInfo->FramePointerSaveIndex) ||
-                   EnableSpillVGPRToAGPR) &&
+  assert((HasFP || !FPSaved || EnableSpillVGPRToAGPR) &&
          "Saved FP but didn't need it");
 
-  assert((!HasBP || (FuncInfo->SGPRForBPSaveRestoreCopy ||
-                     FuncInfo->BasePointerSaveIndex)) &&
+  bool BPSaved = FuncInfo->hasCustomSGPRSpillEntry(BasePtrReg);
+  assert((!HasBP || BPSaved) &&
          "Needed to save BP but didn't save it anywhere");
 
-  assert((HasBP || (!FuncInfo->SGPRForBPSaveRestoreCopy &&
-                    !FuncInfo->BasePointerSaveIndex)) &&
-         "Saved BP but didn't need it");
+  assert((HasBP || !BPSaved) && "Saved BP but didn't need it");
 }
 
 void SIFrameLowering::emitEpilogue(MachineFunction &MF,
                                    MachineBasicBlock &MBB) const {
   const SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
   if (FuncInfo->isEntryFunction())
     return;
 
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIInstrInfo *TII = ST.getInstrInfo();
-  MachineRegisterInfo &MRI = MF.getRegInfo();
   const SIRegisterInfo &TRI = TII->getRegisterInfo();
   LivePhysRegs LiveRegs;
   // Get the insert location for the epilogue. If there were no terminators in
   // the block, get the last instruction.
   MachineBasicBlock::iterator MBBI = MBB.end();
   DebugLoc DL;
   if (!MBB.empty()) {
     MBBI = MBB.getLastNonDebugInstr();
     if (MBBI != MBB.end())
       DL = MBBI->getDebugLoc();
 
     MBBI = MBB.getFirstTerminator();
   }
 
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   uint32_t NumBytes = MFI.getStackSize();
   uint32_t RoundedSize = FuncInfo->isStackRealigned()
                              ? NumBytes + MFI.getMaxAlign().value()
                              : NumBytes;
   const Register StackPtrReg = FuncInfo->getStackPtrOffsetReg();
-  const Register FramePtrReg = FuncInfo->getFrameOffsetReg();
-  const Register BasePtrReg =
-      TRI.hasBasePointer(MF) ? TRI.getBaseRegister() : Register();
-
-  Optional<int> FPSaveIndex = FuncInfo->FramePointerSaveIndex;
-  Optional<int> BPSaveIndex = FuncInfo->BasePointerSaveIndex;
 
   if (RoundedSize != 0 && hasFP(MF)) {
     auto Add = BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::S_ADD_I32), StackPtrReg)
         .addReg(StackPtrReg)
         .addImm(-static_cast<int64_t>(RoundedSize * getScratchScaleFactor(ST)))
         .setMIFlag(MachineInstr::FrameDestroy);
     Add->getOperand(3).setIsDead(); // Mark SCC as dead.
   }
 
-  if (FuncInfo->SGPRForFPSaveRestoreCopy) {
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::COPY), FramePtrReg)
-        .addReg(FuncInfo->SGPRForFPSaveRestoreCopy)
-        .setMIFlag(MachineInstr::FrameDestroy);
-  }
-
-  if (FuncInfo->SGPRForBPSaveRestoreCopy) {
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::COPY), BasePtrReg)
-        .addReg(FuncInfo->SGPRForBPSaveRestoreCopy)
-        .setMIFlag(MachineInstr::FrameDestroy);
-  }
-
-  auto RestoreSGPRFromMemory = [&](Register Reg, const int FI) {
-    initLiveRegs(LiveRegs, TRI, FuncInfo, MF, MBB, MBBI, /*IsProlog*/ false);
-    MCPhysReg TmpVGPR = findScratchNonCalleeSaveRegister(
-        MRI, LiveRegs, AMDGPU::VGPR_32RegClass);
-    if (!TmpVGPR)
-      report_fatal_error("failed to find free scratch register");
-    buildEpilogRestore(ST, TRI, *FuncInfo, LiveRegs, MF, MBB, MBBI, DL, TmpVGPR,
-                       FI);
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::V_READFIRSTLANE_B32), Reg)
-        .addReg(TmpVGPR, RegState::Kill);
-  };
-
-  auto RestoreSGPRFromVGPRLane = [&](Register Reg, const int FI) {
-    assert(MFI.getStackID(FI) == TargetStackID::SGPRSpill);
-    ArrayRef<SIRegisterInfo::SpilledReg> Spill =
-        FuncInfo->getSGPRToVGPRCustomSpills(FI);
-    assert(Spill.size() == 1);
-    BuildMI(MBB, MBBI, DL, TII->get(AMDGPU::V_READLANE_B32), Reg)
-        .addReg(Spill[0].VGPR)
-        .addImm(Spill[0].Lane);
-  };
-
-  if (FPSaveIndex) {
-    const int FramePtrFI = *FPSaveIndex;
-    assert(!MFI.isDeadObjectIndex(FramePtrFI));
-    if (spilledToMemory(MF, FramePtrFI))
-      RestoreSGPRFromMemory(FramePtrReg, FramePtrFI);
-    else
-      RestoreSGPRFromVGPRLane(FramePtrReg, FramePtrFI);
-  }
-
-  if (BPSaveIndex) {
-    const int BasePtrFI = *BPSaveIndex;
-    assert(!MFI.isDeadObjectIndex(BasePtrFI));
-    if (spilledToMemory(MF, BasePtrFI))
-      RestoreSGPRFromMemory(BasePtrReg, BasePtrFI);
-    else
-      RestoreSGPRFromVGPRLane(BasePtrReg, BasePtrFI);
+  for (const auto &Spill : FuncInfo->getCustomSGPRSpills()) {
+    CustomSGPRSpillBuilder CSB(Spill.first, Spill.second, MBB, MBBI, DL, TII,
+                               TRI, LiveRegs);
+    CSB.restore();
   }
 
   Register ScratchExecCopy;
   for (const auto &Reg : FuncInfo->getWWMSpills()) {
     Register VGPR = Reg.first;
     int FI = Reg.second;
     if (!ScratchExecCopy)
       ScratchExecCopy =

 #ifndef NDEBUG
 static bool allSGPRSpillsAreDead(const MachineFunction &MF) {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
   for (int I = MFI.getObjectIndexBegin(), E = MFI.getObjectIndexEnd();
        I != E; ++I) {
     if (!MFI.isDeadObjectIndex(I) &&
         MFI.getStackID(I) == TargetStackID::SGPRSpill &&
-        (I != FuncInfo->FramePointerSaveIndex &&
-         I != FuncInfo->BasePointerSaveIndex)) {
+        !FuncInfo->checkIndexInCustomSGPRSpills(I)) {
       return false;
     }
   }
 
   return true;
 }
 #endif
 
if (UnusedLowVGPR && (TRI->getHWRegIndex(UnusedLowVGPR) <
       // freezeReservedRegs() so that getReservedRegs() can reserve this newly
       // identified VGPR (for AGPR copy).
       FuncInfo->setVGPRForAGPRCopy(UnusedLowVGPR);
       MRI.freezeReservedRegs(MF);
     }
   }
 }
 
+// The special SGPR spills like the one needed for FP, BP or any reserved
+// registers delayed until frame lowering.
+void SIFrameLowering::determineCustomSGPRSaves(MachineFunction &MF,
+                                               BitVector &SavedVGPRs) const {
+  MachineFrameInfo &FrameInfo = MF.getFrameInfo();
+  SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  const SIRegisterInfo *TRI = ST.getRegisterInfo();
+  LivePhysRegs LiveRegs;
+  LiveRegs.init(*TRI);
+  // Initially mark callee saved registers as used so we will not choose them
+  // while looking for scratch SGPRs.
+  const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs();
+  for (unsigned I = 0; CSRegs[I]; ++I)
+    LiveRegs.addReg(CSRegs[I]);
+
+  // hasFP only knows about stack objects that already exist. We're now
+  // determining the stack slots that will be created, so we have to predict
+  // them. Stack objects force FP usage with calls.
+  //
+  // Note a new VGPR CSR may be introduced if one is used for the spill, but we
+  // don't want to report it here.
+  //
+  // FIXME: Is this really hasReservedCallFrame?
+  const bool WillHaveFP =
+      FrameInfo.hasCalls() &&
+      (SavedVGPRs.any() || !allStackObjectsAreDead(FrameInfo));
+
+  if (WillHaveFP || hasFP(MF)) {
+    Register FramePtrReg = MFI->getFrameOffsetReg();
+    assert(!MFI->hasCustomSGPRSpillEntry(FramePtrReg) &&
+           "Re-reserving spill slot for FP");
+    getVGPRSpillLaneOrTempRegister(MF, LiveRegs, FramePtrReg);
+  }
+
+  if (TRI->hasBasePointer(MF)) {
+    Register BasePtrReg = TRI->getBaseRegister();
+    assert(!MFI->hasCustomSGPRSpillEntry(BasePtrReg) &&
+           "Re-reserving spill slot for BP");
+    getVGPRSpillLaneOrTempRegister(MF, LiveRegs, BasePtrReg);
+  }
+}
+
 // Only report VGPRs to generic code.
 void SIFrameLowering::determineCalleeSaves(MachineFunction &MF,
                                            BitVector &SavedVGPRs,
                                            RegScavenger *RS) const {
   TargetFrameLowering::determineCalleeSaves(MF, SavedVGPRs, RS);
   SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   if (MFI->isEntryFunction())
     return;
 
-  MachineFrameInfo &FrameInfo = MF.getFrameInfo();
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *TRI = ST.getRegisterInfo();
 
   for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {
       // WRITELANE instructions used for SGPR spills can overwrite the inactive
       // lanes of VGPRs and callee must spill and restore them even if they are
       // marked Caller-saved.
void SIFrameLowering::determineCalleeSaves(MachineFunction &MF,
   SavedVGPRs.clearBitsNotInMask(TRI->getAllVectorRegMask());
 
   // Do not save AGPRs prior to GFX90A because there was no easy way to do so.
   // In gfx908 there was do AGPR loads and stores and thus spilling also
   // require a temporary VGPR.
   if (!ST.hasGFX90AInsts())
     SavedVGPRs.clearBitsInMask(TRI->getAllAGPRRegMask());
 
-  // hasFP only knows about stack objects that already exist. We're now
-  // determining the stack slots that will be created, so we have to predict
-  // them. Stack objects force FP usage with calls.
-  //
-  // Note a new VGPR CSR may be introduced if one is used for the spill, but we
-  // don't want to report it here.
-  //
-  // FIXME: Is this really hasReservedCallFrame?
-  const bool WillHaveFP =
-      FrameInfo.hasCalls() &&
-      (SavedVGPRs.any() || !allStackObjectsAreDead(FrameInfo));
+  determineCustomSGPRSaves(MF, SavedVGPRs);
 
   // The Whole-Wave VGPRs need to be specially inserted in the prolog, so don't
   // allow the default insertion to handle them.
   for (auto &Reg : MFI->getWWMSpills())
     SavedVGPRs.reset(Reg.first);
-
-  LivePhysRegs LiveRegs;
-  LiveRegs.init(*TRI);
-
-  if (WillHaveFP || hasFP(MF)) {
-    assert(!MFI->SGPRForFPSaveRestoreCopy && !MFI->FramePointerSaveIndex &&
-           "Re-reserving spill slot for FP");
-    getVGPRSpillLaneOrTempRegister(MF, LiveRegs, MFI->SGPRForFPSaveRestoreCopy,
-                                   MFI->FramePointerSaveIndex, true);
-  }
-
-  if (TRI->hasBasePointer(MF)) {
-    if (MFI->SGPRForFPSaveRestoreCopy)
-      LiveRegs.addReg(MFI->SGPRForFPSaveRestoreCopy);
-
-    assert(!MFI->SGPRForBPSaveRestoreCopy &&
-           !MFI->BasePointerSaveIndex && "Re-reserving spill slot for BP");
-    getVGPRSpillLaneOrTempRegister(MF, LiveRegs, MFI->SGPRForBPSaveRestoreCopy,
-                                   MFI->BasePointerSaveIndex, false);
-  }
 }
 
 void SIFrameLowering::determineCalleeSavesSGPR(MachineFunction &MF,
                                                BitVector &SavedRegs,
                                                RegScavenger *RS) const {
   TargetFrameLowering::determineCalleeSaves(MF, SavedRegs, RS);
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   if (MFI->isEntryFunction())

 
 bool SIFrameLowering::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!
 
   const SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
-  if (!FuncInfo->SGPRForFPSaveRestoreCopy &&
-      !FuncInfo->SGPRForBPSaveRestoreCopy)
-    return false;
-
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *RI = ST.getRegisterInfo();
   Register FramePtrReg = FuncInfo->getFrameOffsetReg();
   Register BasePtrReg = RI->getBaseRegister();
+  Register SGPRForFPSaveRestoreCopy =
+      FuncInfo->getScratchSGPRCopyDstReg(FramePtrReg);
+  Register SGPRForBPSaveRestoreCopy =
+      FuncInfo->getScratchSGPRCopyDstReg(BasePtrReg);
+  if (!SGPRForFPSaveRestoreCopy && !SGPRForBPSaveRestoreCopy)
+    return false;
+
   unsigned NumModifiedRegs = 0;
 
-  if (FuncInfo->SGPRForFPSaveRestoreCopy)
+  if (SGPRForFPSaveRestoreCopy)
     NumModifiedRegs++;
-  if (FuncInfo->SGPRForBPSaveRestoreCopy)
+  if (SGPRForBPSaveRestoreCopy)
     NumModifiedRegs++;
 
   for (auto &CS : CSI) {
-    if (CS.getReg() == FramePtrReg && FuncInfo->SGPRForFPSaveRestoreCopy) {
-      CS.setDstReg(FuncInfo->SGPRForFPSaveRestoreCopy);
+    if (CS.getReg() == FramePtrReg && SGPRForFPSaveRestoreCopy) {
+      CS.setDstReg(SGPRForFPSaveRestoreCopy);
       if (--NumModifiedRegs)
         break;
-    } else if (CS.getReg() == BasePtrReg &&
-               FuncInfo->SGPRForBPSaveRestoreCopy) {
-      CS.setDstReg(FuncInfo->SGPRForBPSaveRestoreCopy);
+    } else if (CS.getReg() == BasePtrReg && SGPRForBPSaveRestoreCopy) {
+      CS.setDstReg(SGPRForBPSaveRestoreCopy);
       if (--NumModifiedRegs)
         break;
     }
   }
 
   return false;
 }
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

static void mapping(IO &YamlIO, SIMachineFunctionInfo &MFI) {
     YamlIO.mapOptional("scavengeFI", MFI.ScavengeFI);
     YamlIO.mapOptional("vgprForAGPRCopy", MFI.VGPRForAGPRCopy,
                        StringValue()); // Don't print out when it's empty.
   }
 };
 
 } // end namespace yaml
 
+// A CSR SGPR value can be preserved inside a callee using one of the following
+// methods.
+//   1. Copy to an unused scratch SGPR.
+//   2. Spill to a VGPR lane.
+//   3. Spill to memory via. a scratch VGPR.
+// class CustomSGPRSaveInfo represents the save/restore method used for an SGPR
+// at function prolog/epilog.
+enum class SGPRSaveKind : uint8_t {
+  COPY_TO_SCRATCH_SGPR,
+  SPILL_TO_VGPR_LANE,
+  SPILL_TO_MEM
+};
+
+class CustomSGPRSaveInfo {
+  SGPRSaveKind Kind;
+  union {
+    int Index;
+    Register Reg;
+  };
+
+public:
+  CustomSGPRSaveInfo(SGPRSaveKind K, int I) : Kind(K), Index(I) {}
+  CustomSGPRSaveInfo(SGPRSaveKind K, Register R) : Kind(K), Reg(R) {}
+  Register getReg() const { return Reg; }
+  int getIndex() const { return Index; }
+  SGPRSaveKind getKind() const { return Kind; }
+};
+
 /// This class keeps track of the SPI_SP_INPUT_ADDR config register, which
 /// tells the hardware which interpolation parameters to load.
 class SIMachineFunctionInfo final : public AMDGPUMachineFunction {
   friend class GCNTargetMachine;
 
   // Registers that may be reserved for spilling purposes. These may be the same
   // as the input registers.
   Register ScratchRSrcReg = AMDGPU::PRIVATE_RSRC_REG;
private:
 
   using ReservedRegSet = SmallSetVector<Register, 8>;
   // To track the VGPRs reserved for WWM instructions. They get stack slots
   // later during PrologEpilogInserter and get added into the superset WWMSpills
   // for actual spilling. A separate set makes the register reserved part and
   // the serialization easier.
   ReservedRegSet WWMReservedRegs;
 
+  using CustomSGPRSpillsMap = DenseMap<Register, CustomSGPRSaveInfo>;
+  // To track the custom SGPR spill method used for a CSR SGPR register during
+  // frame lowering. Even though the SGPR spills are handled during
+  // SILowerSGPRSpills pass, some special handling needed later during the
+  // PrologEpilogInserter.
+  CustomSGPRSpillsMap CustomSGPRSpills;
+
   DenseMap<int, VGPRSpillToAGPR> VGPRToAGPRSpills;
 
   // AGPRs used for VGPR spills.
   SmallVector<MCPhysReg, 32> SpillAGPR;
 
   // VGPRs used for AGPR spills.
   SmallVector<MCPhysReg, 32> SpillVGPR;
 
   // Emergency stack slot. Sometimes, we create this before finalizing the stack
   // frame, so save it here and add it to the RegScavenger later.
   Optional<int> ScavengeFI;
 
 private:
   Register VGPRForAGPRCopy;
 
 public:
   Register getVGPRForAGPRCopy() const {
     return VGPRForAGPRCopy;
   }
 
   void setVGPRForAGPRCopy(Register NewVGPRForAGPRCopy) {
     VGPRForAGPRCopy = NewVGPRForAGPRCopy;
   }
 
-public: // FIXME
-  /// If this is set, an SGPR used for save/restore of the register used for the
-  /// frame pointer.
-  Register SGPRForFPSaveRestoreCopy;
-  Optional<int> FramePointerSaveIndex;
-
-  /// If this is set, an SGPR used for save/restore of the register used for the
-  /// base pointer.
-  Register SGPRForBPSaveRestoreCopy;
-  Optional<int> BasePointerSaveIndex;
-
   bool isCalleeSavedReg(const MCPhysReg *CSRegs, MCPhysReg Reg);
 
 public:
   SIMachineFunctionInfo(const MachineFunction &MF);
   SIMachineFunctionInfo(const SIMachineFunctionInfo &MFI) = default;
 
   MachineFunctionInfo *
   clone(BumpPtrAllocator &Allocator, MachineFunction &DestMF,
return (I == SGPRToVGPRSpills.end())
                ? ArrayRef<SIRegisterInfo::SpilledReg>()
                : makeArrayRef(I->second);
   }
 
   ArrayRef<Register> getSGPRSpillVGPRs() const { return SpillVGPRs; }
   const WWMSpillsMap &getWWMSpills() const { return WWMSpills; }
   const ReservedRegSet &getWWMReservedRegs() const { return WWMReservedRegs; }
 
+  const CustomSGPRSpillsMap &getCustomSGPRSpills() const {
+    return CustomSGPRSpills;
+  }
+
+  void addToCustomSGPRSpills(Register Reg, CustomSGPRSaveInfo SI) {
+    CustomSGPRSpills.insert(std::make_pair(Reg, SI));
+  }
+
+  // Check if an entry created for \p Reg in the custom SGPR spills. Return true
+  // on success and false otherwise.
+  bool hasCustomSGPRSpillEntry(Register Reg) const {
+    return CustomSGPRSpills.find(Reg) != CustomSGPRSpills.end();
+  }
+
+  // Get the scratch SGPR if allocated to save/restore \p Reg.
+  Register getScratchSGPRCopyDstReg(Register Reg) const {
+    auto I = CustomSGPRSpills.find(Reg);
+    if (I != CustomSGPRSpills.end() &&
+        I->second.getKind() == SGPRSaveKind::COPY_TO_SCRATCH_SGPR)
+      return I->second.getReg();
+
+    return AMDGPU::NoRegister;
+  }
+
+  // Get all scratch SGPRs allocated to copy/restore the custom SGPR spills.
+  void getAllScratchSGPRCopyDstRegs(SmallVectorImpl<Register> &Regs) const {
+    for (const auto &SI : CustomSGPRSpills) {
+      if (SI.second.getKind() == SGPRSaveKind::COPY_TO_SCRATCH_SGPR)
+        Regs.push_back(SI.second.getReg());
+    }
+  }
+
+  // Check if \p FI is allocated for any custom spill to a VGPR lane.
+  bool checkIndexInCustomSGPRSpills(int FI) const {
+    return find_if(CustomSGPRSpills,
+                   [FI](const std::pair<Register, CustomSGPRSaveInfo> &SI) {
+                     return SI.second.getKind() ==
+                                SGPRSaveKind::SPILL_TO_VGPR_LANE &&
+                            SI.second.getIndex() == FI;
+                   }) != CustomSGPRSpills.end();
+  }
+
   ArrayRef<SIRegisterInfo::SpilledReg>
   getSGPRToVGPRCustomSpills(int FrameIndex) const {
     auto I = SGPRToVGPRCustomSpills.find(FrameIndex);
     return (I == SGPRToVGPRCustomSpills.end())
                ? ArrayRef<SIRegisterInfo::SpilledReg>()
                : makeArrayRef(I->second);
   }
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

bool SIMachineFunctionInfo::removeDeadFrameIndices(
   for (auto &R : make_early_inc_range(SGPRToVGPRSpills)) {
     MFI.RemoveStackObject(R.first);
     SGPRToVGPRSpills.erase(R.first);
   }
 
   bool HaveSGPRToMemory = false;
 
   if (ResetSGPRSpillStackIDs) {
-    // All other SPGRs must be allocated on the default stack, so reset the
+    // All other SGPRs must be allocated on the default stack, so reset the
     // stack ID.
-    for (int i = MFI.getObjectIndexBegin(), e = MFI.getObjectIndexEnd(); i != e;
-         ++i) {
-      if (i != FramePointerSaveIndex && i != BasePointerSaveIndex) {
-        if (MFI.getStackID(i) == TargetStackID::SGPRSpill) {
-          MFI.setStackID(i, TargetStackID::Default);
+    for (int I = MFI.getObjectIndexBegin(), E = MFI.getObjectIndexEnd(); I != E;
+         ++I) {
+      if (!checkIndexInCustomSGPRSpills(I)) {
+        if (MFI.getStackID(I) == TargetStackID::SGPRSpill) {
+          MFI.setStackID(I, TargetStackID::Default);
           HaveSGPRToMemory = true;
         }
       }
     }
   }
 
   for (auto &R : VGPRToAGPRSpills) {
     if (R.second.IsDead)