[CodeGen] Use RegUnits in RegisterClassInfo::getLastCalleeSavedAlias

[foad]

Summary

Change the implementation of getLastCalleeSavedAlias to use RegUnits
instead of register aliases. This is much faster on targets like AMDGPU
which define a very large number of overlapping register tuples.

No functional change intended. If PhysReg overlaps multiple CSRs then
getLastCalleeSavedAlias(PhysReg) could conceivably return a different
arbitrary one, but currently it is only used for some debug printing
anyway.

[qcolombet]

llvm/include/llvm/CodeGen/RegisterClassInfo.h

121

Cannot this get out of range?
I.e., I believe we could assert that *UI is bigger than CalleeSavedAliases.size().

I am wondering why the check was there in the previous implementation.

[foad]

llvm/include/llvm/CodeGen/RegisterClassInfo.h

121

I don't know! I just found that when I removed the check, everything still worked.

[qcolombet]

llvm/include/llvm/CodeGen/RegisterClassInfo.h

121

Oh boy!
Looking at the history of the change, I don't see any explanation why we needed that check. (Unless I am mistaken, the check has been introduced with https://reviews.llvm.org/D28566.)

Could it be possible to call getLastCalleeSavedAlias before CalleeSavedAliases has been initialized?

[arsenm]

llvm/include/llvm/CodeGen/RegisterClassInfo.h

121

My guess would be to handle calling conventions with 0 callee saved registers. Given that would cover the most common case for AMDGPU this is probably ok

[arsenm]

llvm/include/llvm/CodeGen/RegisterClassInfo.h

60

SmallVector<4> seems like a weird way to map this. I'd expect all sets of registers to be BitVectors?

[foad]

llvm/include/llvm/CodeGen/RegisterClassInfo.h

60

Not sure what you mean. We just want an array here, so we can lookup the CSR corresponding to a particular regunit as quickly as possible. (I agree the "4" is weird.)

llvm/include/llvm/CodeGen/RegisterClassInfo.h

class RegisterClassInfo {
 
   const MachineFunction *MF = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
 
   // Callee saved registers of last MF.
   // Used only to determine if an update for CalleeSavedAliases is necessary.
   SmallVector<MCPhysReg, 16> LastCalleeSavedRegs;
 
-  // Map register alias to the callee saved Register.
+  // Map regunit to the callee saved Register.
   SmallVector<MCPhysReg, 4> CalleeSavedAliases;
 
   // Indicate if a specified callee saved register be in the allocation order
   // exactly as written in the tablegen descriptions or listed later.
   BitVector IgnoreCSRForAllocOrder;
 
   // Reserved registers in the current MF.
   BitVector Reserved;
 
public:
   /// Register classes like GR32_NOSP are not proper sub-classes because %esp
   /// is not allocatable.  Similarly, tGPR is not a proper sub-class in Thumb
   /// mode because the GPR super-class is not legal.
   bool isProperSubClass(const TargetRegisterClass *RC) const {
     return get(RC).ProperSubClass;
   }
 
   /// getLastCalleeSavedAlias - Returns the last callee saved register that
-  /// overlaps PhysReg, or NoRegister if Reg doesn't overlap a
+  /// overlaps PhysReg, or NoRegister if PhysReg doesn't overlap a
   /// CalleeSavedAliases.
   MCRegister getLastCalleeSavedAlias(MCRegister PhysReg) const {
-    if (PhysReg.id() < CalleeSavedAliases.size())
-      return CalleeSavedAliases[PhysReg];
-    return MCRegister::NoRegister;
+    MCRegister CSR;
+    for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
+      CSR = CalleeSavedAliases[*UI];
+      if (CSR)
+        break;
+    }
+    return CSR;
   }
 
   /// Get the minimum register cost in RC's allocation order.
   /// This is the smallest value in RegCosts[Reg] for all
   /// the registers in getOrder(RC).
   uint8_t getMinCost(const TargetRegisterClass *RC) const {
     return get(RC).MinCost;
   }

llvm/lib/CodeGen/RegisterClassInfo.cpp

if (!Update) {
     }
   }
 
   // Get the callee saved registers.
   if (CSRChanged) {
     LastCalleeSavedRegs.clear();
     // Build a CSRAlias map. Every CSR alias saves the last
     // overlapping CSR.
-    CalleeSavedAliases.assign(TRI->getNumRegs(), 0);
+    CalleeSavedAliases.assign(TRI->getNumRegUnits(), 0);
     for (const MCPhysReg *I = CSR; *I; ++I) {
-      for (MCRegAliasIterator AI(*I, TRI, true); AI.isValid(); ++AI)
-        CalleeSavedAliases[*AI] = *I;
+      for (MCRegUnitIterator UI(*I, TRI); UI.isValid(); ++UI)
+        CalleeSavedAliases[*UI] = *I;
       LastCalleeSavedRegs.push_back(*I);
     }
 
     Update = true;
   }
 
   // Even if CSR list is same, we could have had a different allocation order
   // if ignoreCSRForAllocationOrder is evaluated differently.
void RegisterClassInfo::compute(const TargetRegisterClass *RC) const {
   ArrayRef<MCPhysReg> RawOrder = RC->getRawAllocationOrder(*MF);
   for (unsigned PhysReg : RawOrder) {
     // Remove reserved registers from the allocation order.
     if (Reserved.test(PhysReg))
       continue;
     uint8_t Cost = RegCosts[PhysReg];
     MinCost = std::min(MinCost, Cost);
 
-    if (CalleeSavedAliases[PhysReg] &&
+    if (getLastCalleeSavedAlias(PhysReg) &&
         !STI.ignoreCSRForAllocationOrder(*MF, PhysReg))
       // PhysReg aliases a CSR, save it for later.
       CSRAlias.push_back(PhysReg);
     else {
       if (Cost != LastCost)
         LastCostChange = N;
       RCI.Order[N++] = PhysReg;
       LastCost = Cost;