2 files changed, 8 insertions, 7 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
index c8304f875f7..6d7260d7aee 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
@@ -1236,9 +1236,10 @@ bool FastISel::lowerCallTo(CallLoweringInfo &CLI) {
     updateValueMap(CLI.CS->getInstruction(), CLI.ResultReg, CLI.NumResultRegs);
 
   // Set labels for heapallocsite call.
-  if (CLI.CS)
-    if (MDNode *MD = CLI.CS->getInstruction()->getMetadata("heapallocsite"))
-      CLI.Call->setHeapAllocMarker(*MF, MD);
+  if (CLI.CS && CLI.CS->getInstruction()->hasMetadata("heapallocsite")) {
+    const MDNode *MD = CLI.CS->getInstruction()->getMetadata("heapallocsite");
+    MF->addCodeViewHeapAllocSite(CLI.Call, MD);
+  }
 
   return true;
 }
diff --git a/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp b/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp
index 0e4d783e350..d4c1fb36475 100644
--- a/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp
@@ -910,9 +910,10 @@ EmitSchedule(MachineBasicBlock::iterator &InsertPos) {
       if (HasDbg)
         ProcessSourceNode(N, DAG, Emitter, VRBaseMap, Orders, Seen, NewInsn);
 
-      if (MDNode *MD = DAG->getHeapAllocSite(N))
+      if (MDNode *MD = DAG->getHeapAllocSite(N)) {
         if (NewInsn && NewInsn->isCall())
-          NewInsn->setHeapAllocMarker(MF, MD);
+          MF.addCodeViewHeapAllocSite(NewInsn, MD);
+      }
 
       GluedNodes.pop_back();
     }
@@ -922,10 +923,9 @@ EmitSchedule(MachineBasicBlock::iterator &InsertPos) {
     if (HasDbg)
       ProcessSourceNode(SU->getNode(), DAG, Emitter, VRBaseMap, Orders, Seen,
                         NewInsn);
-
     if (MDNode *MD = DAG->getHeapAllocSite(SU->getNode())) {
       if (NewInsn && NewInsn->isCall())
-        NewInsn->setHeapAllocMarker(MF, MD);
+        MF.addCodeViewHeapAllocSite(NewInsn, MD);
     }
   }