diff --git a/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp b/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp
--- a/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp
@@ -910,9 +910,13 @@
       if (HasDbg)
         ProcessSourceNode(N, DAG, Emitter, VRBaseMap, Orders, Seen, NewInsn);
 
+      // if (NewInsn && NewInsn->isCall())
+        // NewInsn->dump();
+
       if (MDNode *MD = DAG->getHeapAllocSite(N)) {
-        if (NewInsn && NewInsn->isCall())
+        if (NewInsn && NewInsn->isCall()) {
           MF.addCodeViewHeapAllocSite(NewInsn, MD);
+        }
       }
 
       GluedNodes.pop_back();
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -4058,10 +4058,9 @@
     MF.getFrameInfo().setHasTailCall();
     SDValue Ret = DAG.getNode(X86ISD::TC_RETURN, dl, NodeTys, Ops);
     DAG.addCallSiteInfo(Ret.getNode(), std::move(CSInfo));
-    if (CLI.CS && CLI.CS->getMetadata("heapallocsite")) {
-      DAG.addHeapAllocSite(Chain.getNode(),
+    if (CLI.CS && CLI.CS->getMetadata("heapallocsite"))
+      DAG.addHeapAllocSite(Ret.getNode(),
                            CLI.CS->getMetadata("heapallocsite"));
-    }
 
     return Ret;
   }
@@ -4075,10 +4074,9 @@
   DAG.addCallSiteInfo(Chain.getNode(), std::move(CSInfo));
 
   // Save heapallocsite metadata.
-  if (CLI.CS && CLI.CS->getMetadata("heapallocsite")) {
+  if (CLI.CS && CLI.CS->getMetadata("heapallocsite"))
     DAG.addHeapAllocSite(Chain.getNode(),
                          CLI.CS->getMetadata("heapallocsite"));
-  }
 
   // Create the CALLSEQ_END node.
   unsigned NumBytesForCalleeToPop;