[RegScavenger] Simplify state tracking for backwards scavenging (#71202)

Track the live register state immediately before, instead of after,
MBBI. This makes it simple to track the state at the start or end of a
basic block without a separate (and poorly named) Tracking flag.

This changes the API of the backward(MachineBasicBlock::iterator I)
method, which now recedes to the state just before, instead of just
after, *I. Some clients are simplified by this change.

There is one small functional change shown in the lit tests where
multiple spilled registers all need to be reloaded before the same
instruction. The reloads will now be inserted in the opposite order.
This should not affect correctness.

Author: jayfoad

llvm/include/llvm/CodeGen/RegisterScavenging.h

@@ -38,9 +38,6 @@ class RegScavenger {
   MachineBasicBlock *MBB = nullptr;
   MachineBasicBlock::iterator MBBI;
 
-  /// True if RegScavenger is currently tracking the liveness of registers.
-  bool Tracking = false;
-
   /// Information on scavenged registers (held in a spill slot).
   struct ScavengedInfo {
     ScavengedInfo(int FI = -1) : FrameIndex(FI) {}
@@ -95,21 +92,12 @@ class RegScavenger {
   /// method gives precise results even in the absence of kill flags.
   void backward();
 
-  /// Call backward() as long as the internal iterator does not point to \p I.
+  /// Call backward() to update internal register state to just before \p *I.
   void backward(MachineBasicBlock::iterator I) {
     while (MBBI != I)
       backward();
   }
 
-  /// Move the internal MBB iterator but do not update register states.
-  void skipTo(MachineBasicBlock::iterator I) {
-    if (I == MachineBasicBlock::iterator(nullptr))
-      Tracking = false;
-    MBBI = I;
-  }
-
-  MachineBasicBlock::iterator getCurrentPosition() const { return MBBI; }
-
   /// Return if a specific register is currently used.
   bool isRegUsed(Register Reg, bool includeReserved = true) const;
 

llvm/lib/CodeGen/PrologEpilogInserter.cpp

@@ -1499,7 +1499,7 @@ void PEI::replaceFrameIndicesBackward(MachineBasicBlock *BB,
 
     // Step backwards to get the liveness state at (immedately after) MI.
     if (LocalRS)
-      LocalRS->backward(MI);
+      LocalRS->backward(I);
 
     bool RemovedMI = false;
     for (const auto &[Idx, Op] : enumerate(MI.operands())) {
@@ -1515,11 +1515,6 @@ void PEI::replaceFrameIndicesBackward(MachineBasicBlock *BB,
         break;
     }
 
-    // Refresh the scavenger's internal iterator in case MI was removed or more
-    // instructions were inserted after it.
-    if (LocalRS)
-      LocalRS->skipTo(std::prev(I));
-
     if (!RemovedMI)
       --I;
   }

llvm/lib/CodeGen/RegisterScavenging.cpp

@@ -65,30 +65,22 @@ void RegScavenger::init(MachineBasicBlock &MBB) {
     SI.Reg = 0;
     SI.Restore = nullptr;
   }
-
-  Tracking = false;
 }
 
 void RegScavenger::enterBasicBlock(MachineBasicBlock &MBB) {
   init(MBB);
   LiveUnits.addLiveIns(MBB);
+  MBBI = MBB.begin();
 }
 
 void RegScavenger::enterBasicBlockEnd(MachineBasicBlock &MBB) {
   init(MBB);
   LiveUnits.addLiveOuts(MBB);
-
-  // Move internal iterator at the last instruction of the block.
-  if (!MBB.empty()) {
-    MBBI = std::prev(MBB.end());
-    Tracking = true;
-  }
+  MBBI = MBB.end();
 }
 
 void RegScavenger::backward() {
-  assert(Tracking && "Must be tracking to determine kills and defs");
-
-  const MachineInstr &MI = *MBBI;
+  const MachineInstr &MI = *--MBBI;
   LiveUnits.stepBackward(MI);
 
   // Expire scavenge spill frameindex uses.
@@ -98,12 +90,6 @@ void RegScavenger::backward() {
       I.Restore = nullptr;
     }
   }
-
-  if (MBBI == MBB->begin()) {
-    MBBI = MachineBasicBlock::iterator(nullptr);
-    Tracking = false;
-  } else
-    --MBBI;
 }
 
 bool RegScavenger::isRegUsed(Register Reg, bool includeReserved) const {
@@ -317,9 +303,8 @@ Register RegScavenger::scavengeRegisterBackwards(const TargetRegisterClass &RC,
   // 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);
+  std::pair<MCPhysReg, MachineBasicBlock::iterator> P = findSurvivorBackwards(
+      *MRI, std::prev(MBBI), To, LiveUnits, AllocationOrder, RestoreAfter);
   MCPhysReg Reg = P.first;
   MachineBasicBlock::iterator SpillBefore = P.second;
   // Found an available register?
@@ -334,9 +319,8 @@ Register RegScavenger::scavengeRegisterBackwards(const TargetRegisterClass &RC,
 
   assert(Reg != 0 && "No register left to scavenge!");
 
-  MachineBasicBlock::iterator ReloadAfter =
-    RestoreAfter ? std::next(MBBI) : MBBI;
-  MachineBasicBlock::iterator ReloadBefore = std::next(ReloadAfter);
+  MachineBasicBlock::iterator ReloadBefore =
+      RestoreAfter ? std::next(MBBI) : MBBI;
   if (ReloadBefore != MBB.end())
     LLVM_DEBUG(dbgs() << "Reload before: " << *ReloadBefore << '\n');
   ScavengedInfo &Scavenged = spill(Reg, RC, SPAdj, SpillBefore, ReloadBefore);
@@ -414,9 +398,9 @@ static bool scavengeFrameVirtualRegsInBlock(MachineRegisterInfo &MRI,
   unsigned InitialNumVirtRegs = MRI.getNumVirtRegs();
   bool NextInstructionReadsVReg = false;
   for (MachineBasicBlock::iterator I = MBB.end(); I != MBB.begin(); ) {
-    --I;
-    // Move RegScavenger to the position between *I and *std::next(I).
+    // Move RegScavenger to the position between *std::prev(I) and *I.
     RS.backward(I);
+    --I;
 
     // Look for unassigned vregs in the uses of *std::next(I).
     if (NextInstructionReadsVReg) {

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -3426,7 +3426,7 @@ void AArch64FrameLowering::processFunctionBeforeFrameFinalized(
   // function.
   DebugLoc DL;
   RS->enterBasicBlockEnd(MBB);
-  RS->backward(std::prev(MBBI));
+  RS->backward(MBBI);
   Register DstReg = RS->FindUnusedReg(&AArch64::GPR64commonRegClass);
   assert(DstReg && "There must be a free register after frame setup");
   BuildMI(MBB, MBBI, DL, TII.get(AArch64::MOVi64imm), DstReg).addImm(-2);

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

@@ -299,7 +299,7 @@ bool AArch64SpeculationHardening::instrumentControlFlow(
     if (I == MBB.begin())
       RS.enterBasicBlock(MBB);
     else
-      RS.backward(std::prev(I));
+      RS.backward(I);
     // FIXME: The below just finds *a* unused register. Maybe code could be
     // optimized more if this looks for the register that isn't used for the
     // longest time around this place, to enable more scheduling freedom. Not

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

@@ -1376,7 +1376,7 @@ void SIFrameLowering::processFunctionBeforeFrameFinalized(
                                                 TRI->isAGPR(MRI, VReg))) {
             assert(RS != nullptr);
             RS->enterBasicBlockEnd(MBB);
-            RS->backward(MI);
+            RS->backward(std::next(MI.getIterator()));
             TRI->eliminateFrameIndex(MI, 0, FIOp, RS);
             SpillFIs.set(FI);
             continue;

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

@@ -681,7 +681,7 @@ static void indirectCopyToAGPR(const SIInstrInfo &TII,
   }
 
   RS.enterBasicBlockEnd(MBB);
-  RS.backward(MI);
+  RS.backward(std::next(MI));
 
   // Ideally we want to have three registers for a long reg_sequence copy
   // to hide 2 waitstates between v_mov_b32 and accvgpr_write.

llvm/lib/Target/Mips/Mips16InstrInfo.cpp

@@ -341,7 +341,7 @@ unsigned Mips16InstrInfo::loadImmediate(unsigned FrameReg, int64_t Imm,
   int SpReg = 0;
 
   rs.enterBasicBlockEnd(MBB);
-  rs.backward(II);
+  rs.backward(std::next(II));
   //
   // We need to know which registers can be used, in the case where there
   // are not enough free registers. We exclude all registers that

llvm/lib/Target/PowerPC/PPCFrameLowering.cpp

@@ -466,7 +466,7 @@ PPCFrameLowering::findScratchRegister(MachineBasicBlock *MBB,
       RS.enterBasicBlock(*MBB);
     } else {
       RS.enterBasicBlockEnd(*MBB);
-      RS.backward(std::prev(MBBI));
+      RS.backward(MBBI);
     }
   } else {
     // The scratch register will be used at the start of the block.

llvm/lib/Target/RISCV/RISCVMakeCompressible.cpp

@@ -271,7 +271,7 @@ static Register analyzeCompressibleUses(MachineInstr &FirstMI,
 
   RegScavenger RS;
   RS.enterBasicBlockEnd(MBB);
-  RS.backward(MIs.back()->getIterator());
+  RS.backward(std::next(MIs.back()->getIterator()));
   return RS.scavengeRegisterBackwards(*RCToScavenge, FirstMI.getIterator(),
                                       /*RestoreAfter=*/false, /*SPAdj=*/0,
                                       /*AllowSpill=*/false);

llvm/test/CodeGen/XCore/scavenging.ll

@@ -87,8 +87,8 @@ declare void @g(ptr, ptr)
 ; CHECK: ldw r2, cp[[[INDEX4]]]
 ; r4 & r5 used by InsertSPConstInst() to emit STW_l3r instruction.
 ; CHECK: stw r0, r1[r2]
-; CHECK: ldw r2, sp[0]
 ; CHECK: ldw r1, sp[1]
+; CHECK: ldw r2, sp[0]
 ; CHECK: ldaw r0, sp[0]
 ; scavenge r2 using SR spill slot
 ; CHECK: stw r2, sp[1]