1 files changed, 5 insertions, 8 deletions
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 188021c2832..469ea67f0ac 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -2105,9 +2105,6 @@ unsigned RAGreedy::tryLastChanceRecoloring(LiveInterval &VirtReg,
   // Mark VirtReg as fixed, i.e., it will not be recolored pass this point in
   // this recoloring "session".
   FixedRegisters.insert(VirtReg.reg);
-  // Remember the ID of the last vreg in case the recoloring fails.
-  unsigned LastVReg =
-      TargetRegisterInfo::index2VirtReg(MRI->getNumVirtRegs() - 1);
   SmallVector<unsigned, 4> CurrentNewVRegs;
 
   Order.rewind();
@@ -2179,14 +2176,14 @@ unsigned RAGreedy::tryLastChanceRecoloring(LiveInterval &VirtReg,
     FixedRegisters = SaveFixedRegisters;
     Matrix->unassign(VirtReg);
 
-    // When we move a register from RS_Assign to RS_Split, we do not
-    // actually do anything with it. I.e., it should not end up in NewVRegs.
-    // For the other cases, since we created new live-ranges, we need to
-    // process them.
+    // For a newly created vreg which is also in RecoloringCandidates,
+    // don't add it to NewVRegs because its physical register will be restored
+    // below. Other vregs in CurrentNewVRegs are created by calling
+    // selectOrSplit and should be added into NewVRegs.
     for (SmallVectorImpl<unsigned>::iterator Next = CurrentNewVRegs.begin(),
                                              End = CurrentNewVRegs.end();
          Next != End; ++Next) {
-      if (*Next <= LastVReg && getStage(LIS->getInterval(*Next)) == RS_Split)
+      if (RecoloringCandidates.count(&LIS->getInterval(*Next)))
         continue;
       NewVRegs.push_back(*Next);
     }