diff --git a/llvm/lib/CodeGen/MachinePipeliner.cpp b/llvm/lib/CodeGen/MachinePipeliner.cpp
index 945a10cd69842..bb827f76d36d2 100644
--- a/llvm/lib/CodeGen/MachinePipeliner.cpp
+++ b/llvm/lib/CodeGen/MachinePipeliner.cpp
@@ -1530,7 +1530,7 @@ class HighRegisterPressureDetector {
 
   DenseMap<MachineInstr *, RegisterOperands> ROMap;
 
-  using Instr2LastUsesTy = DenseMap<MachineInstr *, SmallDenseSet<Register, 4>>;
+  using Instr2LastUsesTy = DenseMap<MachineInstr *, SmallSet<VirtRegOrUnit, 4>>;
 
 public:
   using OrderedInstsTy = std::vector<MachineInstr *>;
@@ -1550,12 +1550,8 @@ class HighRegisterPressureDetector {
     }
   }
 
-  void dumpPSet(Register Reg) const {
-    dbgs() << "Reg=" << printReg(Reg, TRI, 0, &MRI) << " PSet=";
-    // FIXME: The static_cast is a bug compensating bugs in the callers.
-    VirtRegOrUnit VRegOrUnit =
-        Reg.isVirtual() ? VirtRegOrUnit(Reg)
-                        : VirtRegOrUnit(static_cast<MCRegUnit>(Reg.id()));
+  void dumpPSet(VirtRegOrUnit VRegOrUnit) const {
+    dbgs() << "Reg=" << printVRegOrUnit(VRegOrUnit, TRI) << " PSet=";
     for (auto PSetIter = MRI.getPressureSets(VRegOrUnit); PSetIter.isValid();
          ++PSetIter) {
       dbgs() << *PSetIter << ' ';
@@ -1564,11 +1560,7 @@ class HighRegisterPressureDetector {
   }
 
   void increaseRegisterPressure(std::vector<unsigned> &Pressure,
-                                Register Reg) const {
-    // FIXME: The static_cast is a bug compensating bugs in the callers.
-    VirtRegOrUnit VRegOrUnit =
-        Reg.isVirtual() ? VirtRegOrUnit(Reg)
-                        : VirtRegOrUnit(static_cast<MCRegUnit>(Reg.id()));
+                                VirtRegOrUnit VRegOrUnit) const {
     auto PSetIter = MRI.getPressureSets(VRegOrUnit);
     unsigned Weight = PSetIter.getWeight();
     for (; PSetIter.isValid(); ++PSetIter)
@@ -1576,8 +1568,8 @@ class HighRegisterPressureDetector {
   }
 
   void decreaseRegisterPressure(std::vector<unsigned> &Pressure,
-                                Register Reg) const {
-    auto PSetIter = MRI.getPressureSets(VirtRegOrUnit(Reg));
+                                VirtRegOrUnit VRegOrUnit) const {
+    auto PSetIter = MRI.getPressureSets(VRegOrUnit);
     unsigned Weight = PSetIter.getWeight();
     for (; PSetIter.isValid(); ++PSetIter) {
       auto &P = Pressure[*PSetIter];
@@ -1587,13 +1579,15 @@ class HighRegisterPressureDetector {
     }
   }
 
-  // Return true if Reg is reserved one, for example, stack pointer
-  bool isReservedRegister(Register Reg) const {
-    return Reg.isPhysical() && MRI.isReserved(Reg.asMCReg());
+  // Return true if \p VRegOrUnit is reserved one, for example, stack pointer
+  bool isReservedRegUnit(VirtRegOrUnit VRegOrUnit) const {
+    return !VRegOrUnit.isVirtualReg() &&
+           MRI.isReservedRegUnit(VRegOrUnit.asMCRegUnit());
   }
 
-  bool isDefinedInThisLoop(Register Reg) const {
-    return Reg.isVirtual() && MRI.getVRegDef(Reg)->getParent() == OrigMBB;
+  bool isDefinedInThisLoop(VirtRegOrUnit Reg) const {
+    return Reg.isVirtualReg() &&
+           MRI.getVRegDef(Reg.asVirtualReg())->getParent() == OrigMBB;
   }
 
   // Search for live-in variables. They are factored into the register pressure
@@ -1605,21 +1599,18 @@ class HighRegisterPressureDetector {
   //     a[i] += b[i] + c;
   // \endcode
   void computeLiveIn() {
-    DenseSet<Register> Used;
+    SmallSet<VirtRegOrUnit, 8> Used;
     for (auto &MI : *OrigMBB) {
       if (MI.isDebugInstr())
         continue;
       for (auto &Use : ROMap[&MI].Uses) {
-        // FIXME: The static_cast is a bug.
-        Register Reg =
-            Use.VRegOrUnit.isVirtualReg()
-                ? Use.VRegOrUnit.asVirtualReg()
-                : Register(static_cast<unsigned>(Use.VRegOrUnit.asMCRegUnit()));
+        VirtRegOrUnit Reg = Use.VRegOrUnit;
         // Ignore the variable that appears only on one side of phi instruction
         // because it's used only at the first iteration.
-        if (MI.isPHI() && Reg != getLoopPhiReg(MI, OrigMBB))
+        if (MI.isPHI() && Reg.isVirtualReg() &&
+            Reg.asVirtualReg() != getLoopPhiReg(MI, OrigMBB))
           continue;
-        if (isReservedRegister(Reg))
+        if (isReservedRegUnit(Reg))
           continue;
         if (isDefinedInThisLoop(Reg))
           continue;
@@ -1654,23 +1645,18 @@ class HighRegisterPressureDetector {
     // Following virtual register will be ignored
     //   - live-in one
     //   - defined but not used in the loop (potentially live-out)
-    DenseSet<Register> TargetRegs;
-    const auto UpdateTargetRegs = [this, &TargetRegs](Register Reg) {
+    SmallSet<VirtRegOrUnit, 8> TargetRegs;
+    const auto UpdateTargetRegs = [this, &TargetRegs](VirtRegOrUnit Reg) {
       if (isDefinedInThisLoop(Reg))
         TargetRegs.insert(Reg);
     };
     for (MachineInstr *MI : OrderedInsts) {
       if (MI->isPHI()) {
         Register Reg = getLoopPhiReg(*MI, OrigMBB);
-        UpdateTargetRegs(Reg);
+        UpdateTargetRegs(VirtRegOrUnit(Reg));
       } else {
         for (auto &Use : ROMap.find(MI)->getSecond().Uses) {
-          // FIXME: The static_cast is a bug.
-          Register Reg = Use.VRegOrUnit.isVirtualReg()
-                             ? Use.VRegOrUnit.asVirtualReg()
-                             : Register(static_cast<unsigned>(
-                                   Use.VRegOrUnit.asMCRegUnit()));
-          UpdateTargetRegs(Reg);
+          UpdateTargetRegs(Use.VRegOrUnit);
         }
       }
     }
@@ -1679,15 +1665,11 @@ class HighRegisterPressureDetector {
       return Stages[MI] + MI->isPHI();
     };
 
-    DenseMap<Register, MachineInstr *> LastUseMI;
+    std::map<VirtRegOrUnit, MachineInstr *> LastUseMI;
     for (MachineInstr *MI : llvm::reverse(OrderedInsts)) {
       for (auto &Use : ROMap.find(MI)->getSecond().Uses) {
-        // FIXME: The static_cast is a bug.
-        Register Reg =
-            Use.VRegOrUnit.isVirtualReg()
-                ? Use.VRegOrUnit.asVirtualReg()
-                : Register(static_cast<unsigned>(Use.VRegOrUnit.asMCRegUnit()));
-        if (!TargetRegs.contains(Reg))
+        VirtRegOrUnit Reg = Use.VRegOrUnit;
+        if (!TargetRegs.contains(Use.VRegOrUnit))
           continue;
         auto [Ite, Inserted] = LastUseMI.try_emplace(Reg, MI);
         if (!Inserted) {
@@ -1721,7 +1703,7 @@ class HighRegisterPressureDetector {
   computeMaxSetPressure(const OrderedInstsTy &OrderedInsts,
                         Instr2StageTy &Stages,
                         const unsigned StageCount) const {
-    using RegSetTy = SmallDenseSet<Register, 16>;
+    using RegSetTy = SmallSet<VirtRegOrUnit, 16>;
 
     // Indexed by #Iter. To treat "local" variables of each stage separately, we
     // manage the liveness of the registers independently by iterations.
@@ -1740,34 +1722,29 @@ class HighRegisterPressureDetector {
     });
 
     const auto InsertReg = [this, &CurSetPressure](RegSetTy &RegSet,
-                                                   VirtRegOrUnit VRegOrUnit) {
-      // FIXME: The static_cast is a bug.
-      Register Reg =
-          VRegOrUnit.isVirtualReg()
-              ? VRegOrUnit.asVirtualReg()
-              : Register(static_cast<unsigned>(VRegOrUnit.asMCRegUnit()));
-      if (!Reg.isValid() || isReservedRegister(Reg))
+                                                   VirtRegOrUnit Reg) {
+      if (isReservedRegUnit(Reg))
         return;
 
       bool Inserted = RegSet.insert(Reg).second;
       if (!Inserted)
         return;
 
-      LLVM_DEBUG(dbgs() << "insert " << printReg(Reg, TRI, 0, &MRI) << "\n");
+      LLVM_DEBUG(dbgs() << "insert " << printVRegOrUnit(Reg, TRI) << "\n");
       increaseRegisterPressure(CurSetPressure, Reg);
       LLVM_DEBUG(dumpPSet(Reg));
     };
 
     const auto EraseReg = [this, &CurSetPressure](RegSetTy &RegSet,
-                                                  Register Reg) {
-      if (!Reg.isValid() || isReservedRegister(Reg))
+                                                  VirtRegOrUnit Reg) {
+      if (isReservedRegUnit(Reg))
         return;
 
       // live-in register
       if (!RegSet.contains(Reg))
         return;
 
-      LLVM_DEBUG(dbgs() << "erase " << printReg(Reg, TRI, 0, &MRI) << "\n");
+      LLVM_DEBUG(dbgs() << "erase " << printVRegOrUnit(Reg, TRI) << "\n");
       RegSet.erase(Reg);
       decreaseRegisterPressure(CurSetPressure, Reg);
       LLVM_DEBUG(dumpPSet(Reg));