diff --git a/llvm/utils/TableGen/Common/CodeGenRegisters.cpp b/llvm/utils/TableGen/Common/CodeGenRegisters.cpp
index 5ec9b35379fa4..4d24eb3de1ed9 100644
--- a/llvm/utils/TableGen/Common/CodeGenRegisters.cpp
+++ b/llvm/utils/TableGen/Common/CodeGenRegisters.cpp
@@ -1849,26 +1849,21 @@ static void computeUberWeights(MutableArrayRef<UberRegSet> UberSets,
   // Skip the first unallocatable set.
   for (UberRegSet &S : UberSets.drop_front()) {
     // Initialize all unit weights in this set, and remember the max units/reg.
-    const CodeGenRegister *Reg = nullptr;
-    unsigned MaxWeight = 0, Weight = 0;
-    for (RegUnitIterator UnitI(S.Regs); UnitI.isValid(); ++UnitI) {
-      if (Reg != UnitI.getReg()) {
-        if (Weight > MaxWeight)
-          MaxWeight = Weight;
-        Reg = UnitI.getReg();
-        Weight = 0;
-      }
-      if (!RegBank.getRegUnit(*UnitI).Artificial) {
-        unsigned UWeight = RegBank.getRegUnit(*UnitI).Weight;
-        if (!UWeight) {
-          UWeight = 1;
-          RegBank.increaseRegUnitWeight(*UnitI, UWeight);
+    unsigned MaxWeight = 0;
+    for (const CodeGenRegister *R : S.Regs) {
+      unsigned Weight = 0;
+      for (unsigned U : R->getRegUnits()) {
+        if (!RegBank.getRegUnit(U).Artificial) {
+          unsigned UWeight = RegBank.getRegUnit(U).Weight;
+          if (!UWeight) {
+            UWeight = 1;
+            RegBank.increaseRegUnitWeight(U, UWeight);
+          }
+          Weight += UWeight;
         }
-        Weight += UWeight;
       }
+      MaxWeight = std::max(MaxWeight, Weight);
     }
-    if (Weight > MaxWeight)
-      MaxWeight = Weight;
     if (S.Weight != MaxWeight) {
       LLVM_DEBUG({
         dbgs() << "UberSet " << &S - UberSets.begin() << " Weight "