RegisterScavenger: Allow fail without spill

[arsenm]

Summary

AMDGPU wants to use this in some contexts where
the spilling is either impossible, or a worse alternative
to doing something else.

[arsenm]

ping

[qcolombet]

LGTM. Nitpick below.

include/llvm/CodeGen/RegisterScavenging.h

161

Add the same comment as what you put on scavengeRegisterBackwards.

[arsenm]

r354816

include/llvm/CodeGen/RegisterScavenging.h

public:
   }
 
   /// Make a register of the specific register class
   /// available and do the appropriate bookkeeping. SPAdj is the stack
   /// adjustment due to call frame, it's passed along to eliminateFrameIndex().
   /// Returns the scavenged register.
   /// This is deprecated as it depends on the quality of the kill flags being
   /// present; Use scavengeRegisterBackwards() instead!
   unsigned scavengeRegister(const TargetRegisterClass *RC,
-                            MachineBasicBlock::iterator I, int SPAdj);
-  unsigned scavengeRegister(const TargetRegisterClass *RegClass, int SPAdj) {
-    return scavengeRegister(RegClass, MBBI, SPAdj);
+                            MachineBasicBlock::iterator I, int SPAdj,
+                            bool AllowSpill = true);
+  unsigned scavengeRegister(const TargetRegisterClass *RegClass, int SPAdj,
+                            bool AllowSpill = true) {
+    return scavengeRegister(RegClass, MBBI, SPAdj, AllowSpill);
   }
 
   /// Make a register of the specific register class available from the current
   /// position backwards to the place before \p To. If \p RestoreAfter is true
   /// this includes the instruction following the current position.
   /// SPAdj is the stack adjustment due to call frame, it's passed along to
   /// eliminateFrameIndex().
   /// Returns the scavenged register.
+  ///
+  /// If \p AllowSpill is false, fail if a spill is required to make the
+  /// register available, and return NoRegister.
   unsigned scavengeRegisterBackwards(const TargetRegisterClass &RC,
                                      MachineBasicBlock::iterator To,
-                                     bool RestoreAfter, int SPAdj);
+                                     bool RestoreAfter, int SPAdj,
+                                     bool AllowSpill = true);
 
   /// Tell the scavenger a register is used.
   void setRegUsed(unsigned Reg, LaneBitmask LaneMask = LaneBitmask::getAll());
 
 private:
   /// Returns true if a register is reserved. It is never "unused".
   bool isReserved(unsigned Reg) const { return MRI->isReserved(Reg); }
 

lib/CodeGen/RegisterScavenging.cpp

if (!TRI->saveScavengerRegister(*MBB, Before, UseMI, &RC, Reg)) {
     FIOperandNum = getFrameIndexOperandNum(*II);
     TRI->eliminateFrameIndex(II, SPAdj, FIOperandNum, this);
   }
   return Scavenged[SI];
 }
 
 unsigned RegScavenger::scavengeRegister(const TargetRegisterClass *RC,
                                         MachineBasicBlock::iterator I,
-                                        int SPAdj) {
+                                        int SPAdj, bool AllowSpill) {
   MachineInstr &MI = *I;
   const MachineFunction &MF = *MI.getMF();
   // Consider all allocatable registers in the register class initially
   BitVector Candidates = TRI->getAllocatableSet(MF, RC);
 
   // Exclude all the registers being used by the instruction.
   for (const MachineOperand &MO : MI.operands()) {
     if (MO.isReg() && MO.getReg() != 0 && !(MO.isUse() && MO.isUndef()) &&
unsigned RegScavenger::scavengeRegister(const TargetRegisterClass *RC,
   unsigned SReg = findSurvivorReg(I, Candidates, 25, UseMI);
 
   // If we found an unused register there is no reason to spill it.
   if (!isRegUsed(SReg)) {
     LLVM_DEBUG(dbgs() << "Scavenged register: " << printReg(SReg, TRI) << "\n");
     return SReg;
   }
 
+  if (!AllowSpill)
+    return 0;
+
   ScavengedInfo &Scavenged = spill(SReg, *RC, SPAdj, I, UseMI);
   Scavenged.Restore = &*std::prev(UseMI);
 
   LLVM_DEBUG(dbgs() << "Scavenged register (with spill): "
                     << printReg(SReg, TRI) << "\n");
 
   return SReg;
 }
 
 unsigned RegScavenger::scavengeRegisterBackwards(const TargetRegisterClass &RC,
                                                  MachineBasicBlock::iterator To,
-                                                 bool RestoreAfter, int SPAdj) {
+                                                 bool RestoreAfter, int SPAdj,
+                                                 bool AllowSpill) {
   const MachineBasicBlock &MBB = *To->getParent();
   const MachineFunction &MF = *MBB.getParent();
 
   // Find the register whose use is furthest away.
   MachineBasicBlock::iterator UseMI;
   ArrayRef<MCPhysReg> AllocationOrder = RC.getRawAllocationOrder(MF);
   std::pair<MCPhysReg, MachineBasicBlock::iterator> P =
       findSurvivorBackwards(*MRI, MBBI, To, LiveUnits, AllocationOrder,
                             RestoreAfter);
   MCPhysReg Reg = P.first;
   MachineBasicBlock::iterator SpillBefore = P.second;
   assert(Reg != 0 && "No register left to scavenge!");
   // Found an available register?
-  if (SpillBefore != MBB.end()) {
-    MachineBasicBlock::iterator ReloadAfter =
-      RestoreAfter ? std::next(MBBI) : MBBI;
-    MachineBasicBlock::iterator ReloadBefore = std::next(ReloadAfter);
-    if (ReloadBefore != MBB.end())
-      LLVM_DEBUG(dbgs() << "Reload before: " << *ReloadBefore << '\n');
-    ScavengedInfo &Scavenged = spill(Reg, RC, SPAdj, SpillBefore, ReloadBefore);
-    Scavenged.Restore = &*std::prev(SpillBefore);
-    LiveUnits.removeReg(Reg);
-    LLVM_DEBUG(dbgs() << "Scavenged register with spill: " << printReg(Reg, TRI)
-                      << " until " << *SpillBefore);
-  } else {
-    LLVM_DEBUG(dbgs() << "Scavenged free register: " << printReg(Reg, TRI)
-                      << '\n');
-  }
+  if (SpillBefore == MBB.end()) {
+    LLVM_DEBUG(dbgs() << "Scavenged free register: " << printReg(Reg, TRI)
+               << '\n');
+    return Reg;
+  }
+
+  if (!AllowSpill)
+    return 0;
+
   return Reg;
 }
 
 /// Allocate a register for the virtual register \p VReg. The last use of
 /// \p VReg is around the current position of the register scavenger \p RS.
 /// \p ReserveAfter controls whether the scavenged register needs to be reserved
 /// after the current instruction, otherwise it will only be reserved before the
 /// current instruction.