[NFC][CodeGen] Refactor: splitting register allocation method in RegAllocFast

[pratlucas]

Summary

The method responsible for allocating a virtual register to a physical
one in RegAllocFast contained the logic for finding which is the best
register in terms of spilling cost, making sure the register is free by
calling the spill methods and update its state and performing the
allocation itself.

This patch splits those resposabilities into separate methods.

[pratlucas]

Ping.

[pratlucas]

This change is no longer relevant after D52010.

llvm/lib/CodeGen/RegAllocFast.cpp

private:
     SmallVector<MachineInstr *, 32> Coalesced;
 
     using RegUnitSet = SparseSet<uint16_t, identity<uint16_t>>;
     /// Set of register units that are used in the current instruction, and so
     /// cannot be allocated.
     RegUnitSet UsedInInstr;
 
     void setPhysRegState(MCPhysReg PhysReg, unsigned NewState);
+    bool isPhysRegFree(MCPhysReg PhysReg) const;
 
     /// Mark a physreg as used in this instruction.
     void markRegUsedInInstr(MCPhysReg PhysReg) {
       for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid(); ++Units)
         UsedInInstr.insert(*Units);
     }
 
     /// Check if a physreg or any of its aliases are used in this instruction.
#endif
     LiveRegMap::iterator findLiveVirtReg(Register VirtReg) {
       return LiveVirtRegs.find(Register::virtReg2Index(VirtReg));
     }
 
     LiveRegMap::const_iterator findLiveVirtReg(Register VirtReg) const {
       return LiveVirtRegs.find(Register::virtReg2Index(VirtReg));
     }
 
+    MCPhysReg getBestRegForAllocation(const TargetRegisterClass &Class,
+                                      Register Hint0, Register Hint1) const;
+    MCPhysReg getFreeRegForAllocation(MachineInstr &MI,
+                                      const TargetRegisterClass &Class,
+                                      Register Hint0, Register Hint1);
     void allocVirtReg(MachineInstr &MI, LiveReg &LR, Register Hint);
     void allocVirtRegUndef(MachineOperand &MO);
     MCPhysReg defineVirtReg(MachineInstr &MI, unsigned OpNum, Register VirtReg,
                             Register Hint);
     LiveReg &reloadVirtReg(MachineInstr &MI, unsigned OpNum, Register VirtReg,
                            Register Hint);
     void spillAll(MachineBasicBlock::iterator MI, bool OnlyLiveOut);
     bool setPhysReg(MachineInstr &MI, MachineOperand &MO, MCPhysReg PhysReg);

 INITIALIZE_PASS(RegAllocFast, "regallocfast", "Fast Register Allocator", false,
                 false)
 
 void RegAllocFast::setPhysRegState(MCPhysReg PhysReg, unsigned NewState) {
   for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI)
     RegUnitStates[*UI] = NewState;
 }
 
+bool RegAllocFast::isPhysRegFree(MCPhysReg PhysReg) const {
+  for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
+    if (RegUnitStates[*UI] != regFree)
+      return false;
+  }
+  return true;
+}
+
 /// This allocates space for the specified virtual register to be held on the
 /// stack.
 int RegAllocFast::getStackSpaceFor(Register VirtReg) {
   // Find the location Reg would belong...
   int SS = StackSlotForVirtReg[VirtReg];
   // Already has space allocated?
   if (SS != -1)
     return SS;
default:
       spillVirtReg(MI, VirtReg);
       break;
     case regFree:
     case regReserved:
       break;
     }
   }
 
-  markRegUsedInInstr(PhysReg);
   setPhysRegState(PhysReg, NewState);
 }
 
 /// Return the cost of spilling clearing out PhysReg and aliases so it is free
 /// for allocation. Returns 0 when PhysReg is free or disabled with all aliases
 /// disabled - it can be allocated directly.
 /// \returns spillImpossible when PhysReg or an alias can't be spilled.
 unsigned RegAllocFast::calcSpillCost(MCPhysReg PhysReg) const {

 void RegAllocFast::assignVirtToPhysReg(LiveReg &LR, MCPhysReg PhysReg) {
   Register VirtReg = LR.VirtReg;
   LLVM_DEBUG(dbgs() << "Assigning " << printReg(VirtReg, TRI) << " to "
                     << printReg(PhysReg, TRI) << '\n');
   assert(LR.PhysReg == 0 && "Already assigned a physreg");
   assert(PhysReg != 0 && "Trying to assign no register");
   LR.PhysReg = PhysReg;
   setPhysRegState(PhysReg, VirtReg);
+  markRegUsedInInstr(PhysReg);
 }
 
 static bool isCoalescable(const MachineInstr &MI) {
   return MI.isFullCopy();
 }
 
 Register RegAllocFast::traceCopyChain(Register Reg) const {
   static const unsigned ChainLengthLimit = 3;
for (const MachineInstr &MI : MRI->def_instructions(VirtReg)) {
     }
 
     if (++C >= DefLimit)
       break;
   }
   return Register();
 }
 
-/// Allocates a physical register for VirtReg.
-void RegAllocFast::allocVirtReg(MachineInstr &MI, LiveReg &LR, Register Hint0) {
-  const Register VirtReg = LR.VirtReg;
-
-  assert(Register::isVirtualRegister(VirtReg) &&
-         "Can only allocate virtual registers");
-
-  const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
-  LLVM_DEBUG(dbgs() << "Search register for " << printReg(VirtReg)
-                    << " in class " << TRI->getRegClassName(&RC)
-                    << " with hint " << printReg(Hint0, TRI) << '\n');
-
-  // Take hint when possible.
+/// Finds out which is the best register available to allocate a virtual
+/// register of the specified class, taking the hints whenever possible.
+MCPhysReg
+RegAllocFast::getBestRegForAllocation(const TargetRegisterClass &Class,
+                                      Register Hint0, Register Hint1) const {
+  // Try hints first
   if (Hint0.isPhysical() && MRI->isAllocatable(Hint0) &&
-      RC.contains(Hint0)) {
-    // Ignore the hint if we would have to spill a dirty register.
-    unsigned Cost = calcSpillCost(Hint0);
-    if (Cost < spillDirty) {
-      LLVM_DEBUG(dbgs() << "\tPreferred Register 1: " << printReg(Hint0, TRI)
-                        << '\n');
-      if (Cost)
-        definePhysReg(MI, Hint0, regFree);
-      assignVirtToPhysReg(LR, Hint0);
-      return;
-    } else {
-      LLVM_DEBUG(dbgs() << "\tPreferred Register 1: " << printReg(Hint0, TRI)
-                        << "occupied\n");
-    }
-  } else {
-    Hint0 = Register();
-  }
-
-  // Try other hint.
-  Register Hint1 = traceCopies(VirtReg);
+      Class.contains(Hint0) && (calcSpillCost(Hint0) < spillDirty))
+    return Hint0;
   if (Hint1.isPhysical() && MRI->isAllocatable(Hint1) &&
-      RC.contains(Hint1) && !isRegUsedInInstr(Hint1)) {
-    // Ignore the hint if we would have to spill a dirty register.
-    unsigned Cost = calcSpillCost(Hint1);
-    if (Cost < spillDirty) {
-      LLVM_DEBUG(dbgs() << "\tPreferred Register 0: " << printReg(Hint1, TRI)
-                        << '\n');
-      if (Cost)
-        definePhysReg(MI, Hint1, regFree);
-      assignVirtToPhysReg(LR, Hint1);
-      return;
-    } else {
-      LLVM_DEBUG(dbgs() << "\tPreferred Register 0: " << printReg(Hint1, TRI)
-                        << "occupied\n");
-    }
-  } else {
-    Hint1 = Register();
-  }
+      Class.contains(Hint1) && (calcSpillCost(Hint1) < spillDirty))
+    return Hint1;
 
+  // Otherwise, get reg with best cost following allocation order
   MCPhysReg BestReg = 0;
   unsigned BestCost = spillImpossible;
-  ArrayRef<MCPhysReg> AllocationOrder = RegClassInfo.getOrder(&RC);
+  auto AllocationOrder = RegClassInfo.getOrder(&Class);
   for (MCPhysReg PhysReg : AllocationOrder) {
-    LLVM_DEBUG(dbgs() << "\tRegister: " << printReg(PhysReg, TRI) << ' ');
     unsigned Cost = calcSpillCost(PhysReg);
-    LLVM_DEBUG(dbgs() << "Cost: " << Cost << " BestCost: " << BestCost << '\n');
-    // Immediate take a register with cost 0.
-    if (Cost == 0) {
-      assignVirtToPhysReg(LR, PhysReg);
-      return;
-    }
-
-    if (PhysReg == Hint1 || PhysReg == Hint0)
+    if (PhysReg == Hint0 || PhysReg == Hint1)
       Cost -= spillPrefBonus;
-
     if (Cost < BestCost) {
       BestReg = PhysReg;
       BestCost = Cost;
     }
-  }
+    if (Cost == 0)
+      break;
+  }
+  return BestReg;
 
-  if (!BestReg) {
-    // Nothing we can do: Report an error and keep going with an invalid
-    // allocation.
+/// Gets an avaliable register of the specified class for allocation, making
+/// sure it is free and ready to be used and spilling when necessary.
+MCPhysReg
+RegAllocFast::getFreeRegForAllocation(MachineInstr &MI,
+                                      const TargetRegisterClass &Class,
+                                      Register Hint0, Register Hint1) {
+  MCPhysReg PhysReg = getBestRegForAllocation(Class, Hint0, Hint1);
+
+  if (!PhysReg) {
+    // No register available
+    // Report an error and keep going with an invalid allocation
     if (MI.isInlineAsm())
       MI.emitError("inline assembly requires more registers than available");
     else
       MI.emitError("ran out of registers during register allocation");
-    definePhysReg(MI, *AllocationOrder.begin(), regFree);
-    assignVirtToPhysReg(LR, *AllocationOrder.begin());
-    return;
+    PhysReg = RegClassInfo.getOrder(&Class).front();
   }
 
-  definePhysReg(MI, BestReg, regFree);
-  assignVirtToPhysReg(LR, BestReg);
+  if (!isPhysRegFree(PhysReg))
+    definePhysReg(MI, PhysReg, regFree);
+  return PhysReg;
+}
+
+/// Allocates a physical register for VirtReg.
+void RegAllocFast::allocVirtReg(MachineInstr &MI, LiveReg &LR, Register Hint0) {
+  const Register VirtReg = LR.VirtReg;
+
+  assert(Register::isVirtualRegister(VirtReg) &&
+         "Can only allocate virtual registers");
+
+  const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
+  LLVM_DEBUG(dbgs() << "Search register for " << printReg(VirtReg)
+                    << " in class " << TRI->getRegClassName(&RC)
+                    << " with hint " << printReg(Hint0, TRI) << '\n');
+
+  Register Hint1 = traceCopies(VirtReg);
+  MCPhysReg PhysReg = getFreeRegForAllocation(MI, RC, Hint0, Hint1);
+  assignVirtToPhysReg(LR, PhysReg);
 }
 
 void RegAllocFast::allocVirtRegUndef(MachineOperand &MO) {
   assert(MO.isUndef() && "expected undef use");
   Register VirtReg = MO.getReg();
   assert(Register::isVirtualRegister(VirtReg) && "Expected virtreg");
 
   LiveRegMap::const_iterator LRI = findLiveVirtReg(VirtReg);
if (Register::isVirtualRegister(Reg)) {
       continue;
     }
     if (!MRI->isAllocatable(Reg)) continue;
     if (MO.isUse()) {
       usePhysReg(MO);
     } else if (MO.isEarlyClobber()) {
       definePhysReg(MI, Reg,
                     (MO.isImplicit() || MO.isDead()) ? regFree : regReserved);
+      markRegUsedInInstr(Reg);
       hasEarlyClobbers = true;
     } else
       hasPhysDefs = true;
   }
 
   // The instruction may have virtual register operands that must be allocated
   // the same register at use-time and def-time: early clobbers and tied
   // operands. If there are also physical defs, these registers must avoid
void RegAllocFast::allocateInstruction(MachineInstr &MI) {
   for (unsigned I = 0; I != DefOpEnd; ++I) {
     const MachineOperand &MO = MI.getOperand(I);
     if (!MO.isReg() || !MO.isDef() || !MO.getReg() || MO.isEarlyClobber())
       continue;
     Register Reg = MO.getReg();
 
     if (!Reg || !Reg.isPhysical() || !MRI->isAllocatable(Reg))
       continue;
+    markRegUsedInInstr(Reg);
     definePhysReg(MI, Reg, MO.isDead() ? regFree : regReserved);
   }
 
   // Fourth scan.
   // Allocate defs and collect dead defs.
   for (unsigned I = 0; I != DefOpEnd; ++I) {
     const MachineOperand &MO = MI.getOperand(I);
     if (!MO.isReg() || !MO.isDef() || !MO.getReg() || MO.isEarlyClobber())