RegAllocFast: Record internal state based on register units

Record internal state based on register units. This is often more
efficient as there are typically fewer register units to update
compared to iterating over all the aliases of a register.

Original patch by Matthias Braun, but I've been rebasing and fixing it
for almost 2 years and fixed a few bugs causing intermediate failures
to make this patch independent of the changes in
https://reviews.llvm.org/D52010.

Author: arsenm

llvm/lib/CodeGen/RegAllocFast.cpp

@@ -106,13 +106,8 @@ namespace {
     /// that it is alive across blocks.
     BitVector MayLiveAcrossBlocks;
 
-    /// State of a physical register.
-    enum RegState {
-      /// A disabled register is not available for allocation, but an alias may
-      /// be in use. A register can only be moved out of the disabled state if
-      /// all aliases are disabled.
-      regDisabled,
-
+    /// State of a register unit.
+    enum RegUnitState {
       /// A free register is not currently in use and can be allocated
       /// immediately without checking aliases.
       regFree,
@@ -126,8 +121,8 @@ namespace {
       /// register. In that case, LiveVirtRegs contains the inverse mapping.
     };
 
-    /// Maps each physical register to a RegState enum or a virtual register.
-    std::vector<unsigned> PhysRegState;
+    /// Maps each physical register to a RegUnitState enum or virtual register.
+    std::vector<unsigned> RegUnitStates;
 
     SmallVector<Register, 16> VirtDead;
     SmallVector<MachineInstr *, 32> Coalesced;
@@ -138,6 +133,7 @@ namespace {
     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) {
@@ -189,14 +185,15 @@ namespace {
     bool isLastUseOfLocalReg(const MachineOperand &MO) const;
 
     void addKillFlag(const LiveReg &LRI);
+    bool verifyRegStateMapping(const LiveReg &LR) const;
     void killVirtReg(LiveReg &LR);
     void killVirtReg(Register VirtReg);
     void spillVirtReg(MachineBasicBlock::iterator MI, LiveReg &LR);
     void spillVirtReg(MachineBasicBlock::iterator MI, Register VirtReg);
 
     void usePhysReg(MachineOperand &MO);
     void definePhysReg(MachineBasicBlock::iterator MI, MCPhysReg PhysReg,
-                       RegState NewState);
+                       unsigned NewState);
     unsigned calcSpillCost(MCPhysReg PhysReg) const;
     void assignVirtToPhysReg(LiveReg &, MCPhysReg PhysReg);
 
@@ -229,7 +226,8 @@ namespace {
     bool mayLiveOut(Register VirtReg);
     bool mayLiveIn(Register VirtReg);
 
-    void dumpState();
+    void printRegUnitState(unsigned State) const;
+    void dumpState() const;
   };
 
 } // end anonymous namespace
@@ -240,7 +238,16 @@ INITIALIZE_PASS(RegAllocFast, "regallocfast", "Fast Register Allocator", false,
                 false)
 
 void RegAllocFast::setPhysRegState(MCPhysReg PhysReg, unsigned NewState) {
-  PhysRegState[PhysReg] = 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
@@ -384,12 +391,21 @@ void RegAllocFast::addKillFlag(const LiveReg &LR) {
   }
 }
 
+bool RegAllocFast::verifyRegStateMapping(const LiveReg &LR) const {
+  for (MCRegUnitIterator UI(LR.PhysReg, TRI); UI.isValid(); ++UI) {
+    if (RegUnitStates[*UI] != LR.VirtReg)
+      return false;
+  }
+
+  return true;
+}
+
 /// Mark virtreg as no longer available.
 void RegAllocFast::killVirtReg(LiveReg &LR) {
+  assert(verifyRegStateMapping(LR) && "Broken RegState mapping");
   addKillFlag(LR);
-  assert(PhysRegState[LR.PhysReg] == LR.VirtReg &&
-         "Broken RegState mapping");
-  setPhysRegState(LR.PhysReg, regFree);
+  MCPhysReg PhysReg = LR.PhysReg;
+  setPhysRegState(PhysReg, regFree);
   LR.PhysReg = 0;
 }
 
@@ -416,15 +432,17 @@ void RegAllocFast::spillVirtReg(MachineBasicBlock::iterator MI,
 
 /// Do the actual work of spilling.
 void RegAllocFast::spillVirtReg(MachineBasicBlock::iterator MI, LiveReg &LR) {
-  assert(PhysRegState[LR.PhysReg] == LR.VirtReg && "Broken RegState mapping");
+  assert(verifyRegStateMapping(LR) && "Broken RegState mapping");
+
+  MCPhysReg PhysReg = LR.PhysReg;
 
   if (LR.Dirty) {
     // If this physreg is used by the instruction, we want to kill it on the
     // instruction, not on the spill.
     bool SpillKill = MachineBasicBlock::iterator(LR.LastUse) != MI;
     LR.Dirty = false;
 
-    spill(MI, LR.VirtReg, LR.PhysReg, SpillKill);
+    spill(MI, LR.VirtReg, PhysReg, SpillKill);
 
     if (SpillKill)
       LR.LastUse = nullptr; // Don't kill register again
@@ -460,53 +478,16 @@ void RegAllocFast::usePhysReg(MachineOperand &MO) {
   assert(PhysReg.isPhysical() && "Bad usePhysReg operand");
 
   markRegUsedInInstr(PhysReg);
-  switch (PhysRegState[PhysReg]) {
-  case regDisabled:
-    break;
-  case regReserved:
-    PhysRegState[PhysReg] = regFree;
-    LLVM_FALLTHROUGH;
-  case regFree:
-    MO.setIsKill();
-    return;
-  default:
-    // The physreg was allocated to a virtual register. That means the value we
-    // wanted has been clobbered.
-    llvm_unreachable("Instruction uses an allocated register");
-  }
 
-  // Maybe a superregister is reserved?
-  for (MCRegAliasIterator AI(PhysReg, TRI, false); AI.isValid(); ++AI) {
-    MCPhysReg Alias = *AI;
-    switch (PhysRegState[Alias]) {
-    case regDisabled:
-      break;
+  for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
+    switch (RegUnitStates[*UI]) {
     case regReserved:
-      // Either PhysReg is a subregister of Alias and we mark the
-      // whole register as free, or PhysReg is the superregister of
-      // Alias and we mark all the aliases as disabled before freeing
-      // PhysReg.
-      // In the latter case, since PhysReg was disabled, this means that
-      // its value is defined only by physical sub-registers. This check
-      // is performed by the assert of the default case in this loop.
-      // Note: The value of the superregister may only be partial
-      // defined, that is why regDisabled is a valid state for aliases.
-      assert((TRI->isSuperRegister(PhysReg, Alias) ||
-              TRI->isSuperRegister(Alias, PhysReg)) &&
-             "Instruction is not using a subregister of a reserved register");
+      RegUnitStates[*UI] = regFree;
       LLVM_FALLTHROUGH;
     case regFree:
-      if (TRI->isSuperRegister(PhysReg, Alias)) {
-        // Leave the superregister in the working set.
-        setPhysRegState(Alias, regFree);
-        MO.getParent()->addRegisterKilled(Alias, TRI, true);
-        return;
-      }
-      // Some other alias was in the working set - clear it.
-      setPhysRegState(Alias, regDisabled);
       break;
     default:
-      llvm_unreachable("Instruction uses an alias of an allocated register");
+      llvm_unreachable("Unexpected reg unit state");
     }
   }
 
@@ -519,38 +500,20 @@ void RegAllocFast::usePhysReg(MachineOperand &MO) {
 /// similar to defineVirtReg except the physreg is reserved instead of
 /// allocated.
 void RegAllocFast::definePhysReg(MachineBasicBlock::iterator MI,
-                                 MCPhysReg PhysReg, RegState NewState) {
-  markRegUsedInInstr(PhysReg);
-  switch (Register VirtReg = PhysRegState[PhysReg]) {
-  case regDisabled:
-    break;
-  default:
-    spillVirtReg(MI, VirtReg);
-    LLVM_FALLTHROUGH;
-  case regFree:
-  case regReserved:
-    setPhysRegState(PhysReg, NewState);
-    return;
-  }
-
-  // This is a disabled register, disable all aliases.
-  setPhysRegState(PhysReg, NewState);
-  for (MCRegAliasIterator AI(PhysReg, TRI, false); AI.isValid(); ++AI) {
-    MCPhysReg Alias = *AI;
-    switch (Register VirtReg = PhysRegState[Alias]) {
-    case regDisabled:
-      break;
+                                 MCPhysReg PhysReg, unsigned NewState) {
+  for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
+    switch (unsigned VirtReg = RegUnitStates[*UI]) {
     default:
       spillVirtReg(MI, VirtReg);
-      LLVM_FALLTHROUGH;
+      break;
     case regFree:
     case regReserved:
-      setPhysRegState(Alias, regDisabled);
-      if (TRI->isSuperRegister(PhysReg, Alias))
-        return;
       break;
     }
   }
+
+  markRegUsedInInstr(PhysReg);
+  setPhysRegState(PhysReg, NewState);
 }
 
 /// Return the cost of spilling clearing out PhysReg and aliases so it is free
@@ -563,46 +526,24 @@ unsigned RegAllocFast::calcSpillCost(MCPhysReg PhysReg) const {
                       << " is already used in instr.\n");
     return spillImpossible;
   }
-  switch (Register VirtReg = PhysRegState[PhysReg]) {
-  case regDisabled:
-    break;
-  case regFree:
-    return 0;
-  case regReserved:
-    LLVM_DEBUG(dbgs() << printReg(VirtReg, TRI) << " corresponding "
-                      << printReg(PhysReg, TRI) << " is reserved already.\n");
-    return spillImpossible;
-  default: {
-    LiveRegMap::const_iterator LRI = findLiveVirtReg(VirtReg);
-    assert(LRI != LiveVirtRegs.end() && LRI->PhysReg &&
-           "Missing VirtReg entry");
-    return LRI->Dirty ? spillDirty : spillClean;
-  }
-  }
 
-  // This is a disabled register, add up cost of aliases.
-  LLVM_DEBUG(dbgs() << printReg(PhysReg, TRI) << " is disabled.\n");
-  unsigned Cost = 0;
-  for (MCRegAliasIterator AI(PhysReg, TRI, false); AI.isValid(); ++AI) {
-    MCPhysReg Alias = *AI;
-    switch (Register VirtReg = PhysRegState[Alias]) {
-    case regDisabled:
-      break;
+  for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
+    switch (unsigned VirtReg = RegUnitStates[*UI]) {
     case regFree:
-      ++Cost;
       break;
     case regReserved:
+      LLVM_DEBUG(dbgs() << printReg(VirtReg, TRI) << " corresponding "
+                        << printReg(PhysReg, TRI) << " is reserved already.\n");
       return spillImpossible;
     default: {
       LiveRegMap::const_iterator LRI = findLiveVirtReg(VirtReg);
       assert(LRI != LiveVirtRegs.end() && LRI->PhysReg &&
              "Missing VirtReg entry");
-      Cost += LRI->Dirty ? spillDirty : spillClean;
-      break;
+      return LRI->Dirty ? spillDirty : spillClean;
     }
     }
   }
-  return Cost;
+  return 0;
 }
 
 /// This method updates local state so that we know that PhysReg is the
@@ -909,9 +850,17 @@ void RegAllocFast::handleThroughOperands(MachineInstr &MI,
     if (!Reg || !Reg.isPhysical())
       continue;
     markRegUsedInInstr(Reg);
-    for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI) {
-      if (ThroughRegs.count(PhysRegState[*AI]))
-        definePhysReg(MI, *AI, regFree);
+
+    for (MCRegUnitIterator UI(Reg, TRI); UI.isValid(); ++UI) {
+      if (!ThroughRegs.count(RegUnitStates[*UI]))
+        continue;
+
+      // Need to spill any aliasing registers.
+      for (MCRegUnitRootIterator RI(*UI, TRI); RI.isValid(); ++RI) {
+        for (MCSuperRegIterator SI(*RI, TRI, true); SI.isValid(); ++SI) {
+          definePhysReg(MI, *SI, regFree);
+        }
+      }
     }
   }
 
@@ -975,37 +924,40 @@ void RegAllocFast::handleThroughOperands(MachineInstr &MI,
 }
 
 #ifndef NDEBUG
-void RegAllocFast::dumpState() {
-  for (unsigned Reg = 1, E = TRI->getNumRegs(); Reg != E; ++Reg) {
-    if (PhysRegState[Reg] == regDisabled) continue;
-    dbgs() << " " << printReg(Reg, TRI);
-    switch(PhysRegState[Reg]) {
+
+void RegAllocFast::dumpState() const {
+  for (unsigned Unit = 1, UnitE = TRI->getNumRegUnits(); Unit != UnitE;
+       ++Unit) {
+    switch (unsigned VirtReg = RegUnitStates[Unit]) {
     case regFree:
       break;
     case regReserved:
-      dbgs() << "*";
+      dbgs() << " " << printRegUnit(Unit, TRI) << "[P]";
       break;
     default: {
-      dbgs() << '=' << printReg(PhysRegState[Reg]);
-      LiveRegMap::iterator LRI = findLiveVirtReg(PhysRegState[Reg]);
-      assert(LRI != LiveVirtRegs.end() && LRI->PhysReg &&
-             "Missing VirtReg entry");
-      if (LRI->Dirty)
-        dbgs() << "*";
-      assert(LRI->PhysReg == Reg && "Bad inverse map");
+      dbgs() << ' ' << printRegUnit(Unit, TRI) << '=' << printReg(VirtReg);
+      LiveRegMap::const_iterator I = findLiveVirtReg(VirtReg);
+      assert(I != LiveVirtRegs.end() && "have LiveVirtRegs entry");
+      if (I->Dirty)
+        dbgs() << "[D]";
+      assert(TRI->hasRegUnit(I->PhysReg, Unit) && "inverse mapping present");
       break;
     }
     }
   }
   dbgs() << '\n';
   // Check that LiveVirtRegs is the inverse.
-  for (LiveRegMap::iterator i = LiveVirtRegs.begin(),
-       e = LiveVirtRegs.end(); i != e; ++i) {
-    if (!i->PhysReg)
-      continue;
-    assert(i->VirtReg.isVirtual() && "Bad map key");
-    assert(Register::isPhysicalRegister(i->PhysReg) && "Bad map value");
-    assert(PhysRegState[i->PhysReg] == i->VirtReg && "Bad inverse map");
+  for (const LiveReg &LR : LiveVirtRegs) {
+    Register VirtReg = LR.VirtReg;
+    assert(VirtReg.isVirtual() && "Bad map key");
+    MCPhysReg PhysReg = LR.PhysReg;
+    if (PhysReg != 0) {
+      assert(Register::isPhysicalRegister(PhysReg) &&
+             "mapped to physreg");
+      for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
+        assert(RegUnitStates[*UI] == VirtReg && "inverse map valid");
+      }
+    }
   }
 }
 #endif
@@ -1247,7 +1199,7 @@ void RegAllocFast::allocateBasicBlock(MachineBasicBlock &MBB) {
   this->MBB = &MBB;
   LLVM_DEBUG(dbgs() << "\nAllocating " << MBB);
 
-  PhysRegState.assign(TRI->getNumRegs(), regDisabled);
+  RegUnitStates.assign(TRI->getNumRegUnits(), regFree);
   assert(LiveVirtRegs.empty() && "Mapping not cleared from last block?");
 
   MachineBasicBlock::iterator MII = MBB.begin();

llvm/test/CodeGen/AArch64/arm64-fast-isel-conversion-fallback.ll

@@ -4,8 +4,8 @@
 define i32 @fptosi_wh(half %a) nounwind ssp {
 entry:
 ; CHECK-LABEL: fptosi_wh
-; CHECK: fcvt s1, h0
-; CHECK: fcvtzs [[REG:w[0-9]+]], s1
+; CHECK: fcvt s0, h0
+; CHECK: fcvtzs [[REG:w[0-9]+]], s0
 ; CHECK: mov w0, [[REG]]
   %conv = fptosi half %a to i32
   ret i32 %conv
@@ -15,8 +15,8 @@ entry:
 define i32 @fptoui_swh(half %a) nounwind ssp {
 entry:
 ; CHECK-LABEL: fptoui_swh
-; CHECK: fcvt s1, h0
-; CHECK: fcvtzu [[REG:w[0-9]+]], s1
+; CHECK: fcvt s0, h0
+; CHECK: fcvtzu [[REG:w[0-9]+]], s0
 ; CHECK: mov w0, [[REG]]
   %conv = fptoui half %a to i32
   ret i32 %conv