1 files changed, 10 insertions, 7 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp
index 735dcef..80aeefe 100644
--- a/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp
@@ -916,13 +916,16 @@ SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
     bool IsLocal = (Relocate->getParent() == StatepointInstr->getParent());
 
     RecordType Record;
-    if (IsLocal && LowerAsVReg.count(SDV)) {
-      // Result is already stored in StatepointLowering
-      Record.type = RecordType::SDValueNode;
-    } else if (LowerAsVReg.count(SDV)) {
-      Record.type = RecordType::VReg;
-      assert(VirtRegs.count(SDV));
-      Record.payload.Reg = VirtRegs[SDV];
+    if (LowerAsVReg.count(SDV)) {
+      if (IsLocal) {
+        // Result is already stored in StatepointLowering
+        Record.type = RecordType::SDValueNode;
+      } else {
+        Record.type = RecordType::VReg;
+        auto It = VirtRegs.find(SDV);
+        assert(It != VirtRegs.end());
+        Record.payload.Reg = It->second;
+      }
     } else if (Loc.getNode()) {
       Record.type = RecordType::Spill;
       Record.payload.FI = cast<FrameIndexSDNode>(Loc)->getIndex();