[RDF] RegisterRef/RegisterId improvements. NFC (#168030)

RegisterId can represent a physical register, a MCRegUnit, or
an index into a side structure that stores register masks. These 3
types were encoded by using the physical reg, stack slot, and
virtual register encoding partitions from the Register class.
This encoding scheme alias wasn't well contained so
Register::index2StackSlot and Register::stackSlotIndex appeared
in multiple places.

This patch gives RegisterRef its own encoding defines and separates
it from Register.

I've removed the generic idx() method in favor of getAsMCReg(),
getAsMCRegUnit(), and getMaskIdx() for some degree of type safety.

Some places used the RegisterId field of RegisterRef directly as a
register. Those have been updated to use getAsMCReg.
    
Some special cases for RegisterId 0 have been removed as it can
be treated like a MCRegister by existing code.
    
I think I want to rename the Reg field of RegisterRef to Id, but
I'll do that in another patch.
    
Additionally, callers of the RegisterRef constructor need to be
audited for implicit conversions from Register/MCRegister
to unsigned.

Author: topperc

llvm/include/llvm/CodeGen/RDFRegisters.h

@@ -86,6 +86,11 @@ template <typename T, unsigned N = 32> struct IndexedSet {
 };
 
 struct RegisterRef {
+private:
+  static constexpr RegisterId MaskFlag = 1u << 30;
+  static constexpr RegisterId UnitFlag = 1u << 31;
+
+public:
   RegisterId Reg = 0;
   LaneBitmask Mask = LaneBitmask::getNone(); // Only for registers.
 
@@ -99,7 +104,20 @@ struct RegisterRef {
   constexpr bool isUnit() const { return isUnitId(Reg); }
   constexpr bool isMask() const { return isMaskId(Reg); }
 
-  constexpr unsigned idx() const { return toIdx(Reg); }
+  constexpr MCRegister asMCReg() const {
+    assert(isReg());
+    return Reg;
+  }
+
+  constexpr MCRegUnit asMCRegUnit() const {
+    assert(isUnit());
+    return Reg & ~UnitFlag;
+  }
+
+  constexpr unsigned asMaskIdx() const {
+    assert(isMask());
+    return Reg & ~MaskFlag;
+  }
 
   constexpr operator bool() const {
     return !isReg() || (Reg != 0 && Mask.any());
@@ -110,26 +128,15 @@ struct RegisterRef {
            std::hash<LaneBitmask::Type>{}(Mask.getAsInteger());
   }
 
-  static constexpr bool isRegId(unsigned Id) {
-    return Register::isPhysicalRegister(Id);
-  }
-  static constexpr bool isUnitId(unsigned Id) {
-    return Register::isVirtualRegister(Id);
+  static constexpr bool isRegId(RegisterId Id) {
+    return !(Id & UnitFlag) && !(Id & MaskFlag);
   }
-  static constexpr bool isMaskId(unsigned Id) { return Register(Id).isStack(); }
+  static constexpr bool isUnitId(RegisterId Id) { return Id & UnitFlag; }
+  static constexpr bool isMaskId(RegisterId Id) { return Id & MaskFlag; }
 
-  static constexpr RegisterId toUnitId(unsigned Idx) {
-    return Idx | Register::VirtualRegFlag;
-  }
+  static constexpr RegisterId toUnitId(unsigned Idx) { return Idx | UnitFlag; }
 
-  static constexpr unsigned toIdx(RegisterId Id) {
-    // Not using virtReg2Index or stackSlot2Index, because they are
-    // not constexpr.
-    if (isUnitId(Id))
-      return Id & ~Register::VirtualRegFlag;
-    // RegId and MaskId are unchanged.
-    return Id;
-  }
+  static constexpr RegisterId toMaskId(unsigned Idx) { return Idx | MaskFlag; }
 
   bool operator<(RegisterRef) const = delete;
   bool operator==(RegisterRef) const = delete;
@@ -141,11 +148,11 @@ struct PhysicalRegisterInfo {
                        const MachineFunction &mf);
 
   RegisterId getRegMaskId(const uint32_t *RM) const {
-    return Register::index2StackSlot(RegMasks.find(RM));
+    return RegisterRef::toMaskId(RegMasks.find(RM));
   }
 
   const uint32_t *getRegMaskBits(RegisterId R) const {
-    return RegMasks.get(Register(R).stackSlotIndex());
+    return RegMasks.get(RegisterRef(R).asMaskIdx());
   }
 
   bool alias(RegisterRef RA, RegisterRef RB) const;
@@ -158,7 +165,7 @@ struct PhysicalRegisterInfo {
   }
 
   const BitVector &getMaskUnits(RegisterId MaskId) const {
-    return MaskInfos[Register(MaskId).stackSlotIndex()].Units;
+    return MaskInfos[RegisterRef(MaskId).asMaskIdx()].Units;
   }
 
   std::set<RegisterId> getUnits(RegisterRef RR) const;
@@ -167,7 +174,7 @@ struct PhysicalRegisterInfo {
     return AliasInfos[U].Regs;
   }
 
-  RegisterRef mapTo(RegisterRef RR, unsigned R) const;
+  RegisterRef mapTo(RegisterRef RR, RegisterId R) const;
   const TargetRegisterInfo &getTRI() const { return TRI; }
 
   bool equal_to(RegisterRef A, RegisterRef B) const;

llvm/lib/CodeGen/RDFGraph.cpp

@@ -1827,7 +1827,7 @@ bool DataFlowGraph::hasUntrackedRef(Stmt S, bool IgnoreReserved) const {
   for (Ref R : S.Addr->members(*this)) {
     Ops.push_back(&R.Addr->getOp());
     RegisterRef RR = R.Addr->getRegRef(*this);
-    if (IgnoreReserved && RR.isReg() && ReservedRegs[RR.idx()])
+    if (IgnoreReserved && RR.isReg() && ReservedRegs[RR.asMCReg().id()])
       continue;
     if (!isTracked(RR))
       return true;

llvm/lib/CodeGen/RDFRegisters.cpp

@@ -126,13 +126,10 @@ std::set<RegisterId> PhysicalRegisterInfo::getAliasSet(RegisterId Reg) const {
 std::set<RegisterId> PhysicalRegisterInfo::getUnits(RegisterRef RR) const {
   std::set<RegisterId> Units;
 
-  if (RR.Reg == 0)
-    return Units; // Empty
-
   if (RR.isReg()) {
     if (RR.Mask.none())
       return Units; // Empty
-    for (MCRegUnitMaskIterator UM(RR.idx(), &TRI); UM.isValid(); ++UM) {
+    for (MCRegUnitMaskIterator UM(RR.asMCReg(), &TRI); UM.isValid(); ++UM) {
       auto [U, M] = *UM;
       if ((M & RR.Mask).any())
         Units.insert(U);
@@ -142,7 +139,7 @@ std::set<RegisterId> PhysicalRegisterInfo::getUnits(RegisterRef RR) const {
 
   assert(RR.isMask());
   unsigned NumRegs = TRI.getNumRegs();
-  const uint32_t *MB = getRegMaskBits(RR.idx());
+  const uint32_t *MB = getRegMaskBits(RR.Reg);
   for (unsigned I = 0, E = (NumRegs + 31) / 32; I != E; ++I) {
     uint32_t C = ~MB[I]; // Clobbered regs
     if (I == 0)          // Reg 0 should be ignored
@@ -162,12 +159,13 @@ std::set<RegisterId> PhysicalRegisterInfo::getUnits(RegisterRef RR) const {
   return Units;
 }
 
-RegisterRef PhysicalRegisterInfo::mapTo(RegisterRef RR, unsigned R) const {
+RegisterRef PhysicalRegisterInfo::mapTo(RegisterRef RR, RegisterId R) const {
   if (RR.Reg == R)
     return RR;
-  if (unsigned Idx = TRI.getSubRegIndex(R, RR.Reg))
+  if (unsigned Idx = TRI.getSubRegIndex(RegisterRef(R).asMCReg(), RR.asMCReg()))
     return RegisterRef(R, TRI.composeSubRegIndexLaneMask(Idx, RR.Mask));
-  if (unsigned Idx = TRI.getSubRegIndex(RR.Reg, R)) {
+  if (unsigned Idx =
+          TRI.getSubRegIndex(RR.asMCReg(), RegisterRef(R).asMCReg())) {
     const RegInfo &RI = RegInfos[R];
     LaneBitmask RCM =
         RI.RegClass ? RI.RegClass->LaneMask : LaneBitmask::getAll();
@@ -187,8 +185,8 @@ bool PhysicalRegisterInfo::equal_to(RegisterRef A, RegisterRef B) const {
     return A.Mask == B.Mask;
 
   // Compare reg units lexicographically.
-  MCRegUnitMaskIterator AI(A.Reg, &getTRI());
-  MCRegUnitMaskIterator BI(B.Reg, &getTRI());
+  MCRegUnitMaskIterator AI(A.asMCReg(), &getTRI());
+  MCRegUnitMaskIterator BI(B.asMCReg(), &getTRI());
   while (AI.isValid() && BI.isValid()) {
     auto [AReg, AMask] = *AI;
     auto [BReg, BMask] = *BI;
@@ -225,8 +223,8 @@ bool PhysicalRegisterInfo::less(RegisterRef A, RegisterRef B) const {
     return A.Reg < B.Reg;
 
   // Compare reg units lexicographically.
-  llvm::MCRegUnitMaskIterator AI(A.Reg, &getTRI());
-  llvm::MCRegUnitMaskIterator BI(B.Reg, &getTRI());
+  llvm::MCRegUnitMaskIterator AI(A.asMCReg(), &getTRI());
+  llvm::MCRegUnitMaskIterator BI(B.asMCReg(), &getTRI());
   while (AI.isValid() && BI.isValid()) {
     auto [AReg, AMask] = *AI;
     auto [BReg, BMask] = *BI;
@@ -252,18 +250,17 @@ bool PhysicalRegisterInfo::less(RegisterRef A, RegisterRef B) const {
 }
 
 void PhysicalRegisterInfo::print(raw_ostream &OS, RegisterRef A) const {
-  if (A.Reg == 0 || A.isReg()) {
-    if (0 < A.idx() && A.idx() < TRI.getNumRegs())
-      OS << TRI.getName(A.idx());
+  if (A.isReg()) {
+    MCRegister Reg = A.asMCReg();
+    if (Reg && Reg.id() < TRI.getNumRegs())
+      OS << TRI.getName(Reg);
     else
-      OS << printReg(A.idx(), &TRI);
+      OS << printReg(Reg, &TRI);
     OS << PrintLaneMaskShort(A.Mask);
   } else if (A.isUnit()) {
-    OS << printRegUnit(A.idx(), &TRI);
+    OS << printRegUnit(A.asMCRegUnit(), &TRI);
   } else {
-    assert(A.isMask());
-    // RegMask SS flag is preserved by idx().
-    unsigned Idx = Register(A.idx()).stackSlotIndex();
+    unsigned Idx = A.asMaskIdx();
     const char *Fmt = Idx < 0x10000 ? "%04x" : "%08x";
     OS << "M#" << format(Fmt, Idx);
   }
@@ -280,7 +277,7 @@ bool RegisterAggr::hasAliasOf(RegisterRef RR) const {
   if (RR.isMask())
     return Units.anyCommon(PRI.getMaskUnits(RR.Reg));
 
-  for (MCRegUnitMaskIterator U(RR.Reg, &PRI.getTRI()); U.isValid(); ++U) {
+  for (MCRegUnitMaskIterator U(RR.asMCReg(), &PRI.getTRI()); U.isValid(); ++U) {
     auto [Unit, LaneMask] = *U;
     if ((LaneMask & RR.Mask).any())
       if (Units.test(Unit))
@@ -295,7 +292,7 @@ bool RegisterAggr::hasCoverOf(RegisterRef RR) const {
     return T.reset(Units).none();
   }
 
-  for (MCRegUnitMaskIterator U(RR.Reg, &PRI.getTRI()); U.isValid(); ++U) {
+  for (MCRegUnitMaskIterator U(RR.asMCReg(), &PRI.getTRI()); U.isValid(); ++U) {
     auto [Unit, LaneMask] = *U;
     if ((LaneMask & RR.Mask).any())
       if (!Units.test(Unit))
@@ -310,7 +307,7 @@ RegisterAggr &RegisterAggr::insert(RegisterRef RR) {
     return *this;
   }
 
-  for (MCRegUnitMaskIterator U(RR.Reg, &PRI.getTRI()); U.isValid(); ++U) {
+  for (MCRegUnitMaskIterator U(RR.asMCReg(), &PRI.getTRI()); U.isValid(); ++U) {
     auto [Unit, LaneMask] = *U;
     if ((LaneMask & RR.Mask).any())
       Units.set(Unit);