Add a heap alloc site marker field to the ExtraInfo in MachineInstrs.

[akhuang]

Summary

Add instruction marker to MachineInstr ExtraInfo. This does almost the
same thing as Pre/PostInstrSymbols, except that it doesn't create a label until
printing instructions. This allows for labels to be put around instructions that
are deleted/duplicated somewhere.
Use this marker to track heap alloc site call instructions.

Also undo the workaround in r375137.

This is a fix for https://bugs.llvm.org/show_bug.cgi?id=43479#add_comment

[rnk]

llvm/include/llvm/CodeGen/MachineFunction.h

326–327

Now that CodeViewDebug is going to be in charge of filling in this data structure, I think it should be moved to CodeViewDebug::FunctionInfo. Actually, I see we already have them there, so we can probably remove this vector completely.

328

I don't think this map keyed on MachineInstr pointers will do the right thing in the presence of dead code elimination and code duplication.

For dead code elimination, it could have an ABA problem like the SelectionDAG nodes had before, where our solution was to clear the map after codegen. I think MachineInstr addresses can be recycled if a call is deleted and another instruction allocated at the same address.

For code duplication, are you sure this looks up the correct type? I would expect one of the instructions to not be present in here, and return a null DIType. For the taildup-heapallocsite.ll test case that I added, can you make sure that both S_HEAPALLOCSITE records use the same type indices?

I think you can store the DIType directly in the ExtraInfo struct and that will mean we don't need this anymore.

llvm/include/llvm/CodeGen/MachineInstr.h

245

This label doesn't seem like it's ever emitted. It's really just a special label with the name "heapallocsite". I think you can store the DIType* here, and use that as a heap allocation call site indicator. Another option would be to store metadata references here, which would make it extensible. We could, for example, have this point to a vector of pairs of fixed metadata kinds (MD_dbg, MD_tbaa, etc) and Metadata*. It would provide a natural pathway from moving any kind of metadata annotation from an LLVM IR instruction to a MachineInstr.

Whichever attachment you choose, it should avoid the need for the extra DenseMap, and it should get copied when tail duplication clones the instruction.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

1072–1078

I think it would be best to, as much as possible, avoid having codeview specific codepaths here.

llvm/lib/CodeGen/AsmPrinter/CodeViewDebug.cpp

1106–1109

It'll be nice to avoid this. :)

1413–1414

I think you can use this code pattern here to find all the heap allocation sites, and save them on FunctionInfo for later. This beginFunctionImpl method is sort of the place where the debug info handlers get a chance to look at the entire function before each instruction is emitted. We already do big up front passes over the whole function to create the lexical scope tree, which is what places labels around block scopes.

You should be able to reuse the requestLabelBeforeInsn / requestLabelAfterInsn APIs. The label names won't be as easy to FileCheck, but there will be fewer labels in the output, which is cleaner in some ways.

llvm/lib/CodeGen/MachineInstr.cpp

519

This doesn't track the new instrmarker, so now postinstr symbols and instrmarkers can't coexist on the same instruction.

549–550

I don't think this does the right thing. If I understand correctly, Info is a single pointer. The .set call overwrites the pointer, and then the second call to getPostInstrSymbol will always return null, and then the .set call will reset the pointer to null again.

This extra info code is subtle. Now that there are five possibilities, we might want to come up with some shared logic for this.

Actually, I'm surprised we didn't write unit tests for this in the first place. Oops. It's highly testable data manipulation. You can see examples in llvm/unittests/CodeGen/MachineInstrTest.cpp.

[rnk]

Thanks for putting up with these micro-optimized data structures, sorry they're opaque.

llvm/lib/CodeGen/MachineInstr.cpp

549–550

I meant to elaborate on what I meant by "shared logic": I was thinking that all of the "extra info" accessors should call a routine that takes:

• an MMO ArrayRef
• a pre instr symbol
• a post instr symbol
• the new thing, DIType*, Metatadata*, or an ArrayRef<pair<unsigned, Metadata*>>

It then checks each case in turn:

• only one MMO and nothing else: use the EIIK_MMO case
• only one pre instr symbol: use the EIIK_PreInstr case

...

• if no cases apply, make the extra info struct

[akhuang]

-Moved the logic for editing extra info for machine instrs into a separate function.
-Moved label creation code into CodeViewDebug.
-Remove now unused heapallocsite vector from MachineFunction.

I still need to write unit tests for the MachineInstr stuff. I've just been figuring
how to generate MCSymbols from the faked MachineFunction.

[akhuang]

Clang format

[rnk]

Nice, this looks good. :)

I think if we're going to do full due diligence for this system which we are starting to depend on, next we should serialize this data via MIR. In the MIParser we only handle MMOs, but we should handle these labels and heap alloc site markers here:
https://github.com/llvm/llvm-project/blob/18f805a/llvm/lib/CodeGen/MIRParser/MIParser.cpp#L953
That can be a separate change.

llvm/test/CodeGen/X86/taildup-heapallocsite.ll

47

Let's CHECK for two S_HEAPALLOCSITE records as well.

[akhuang]

• Add unit tests for MachineInstr extra info
• Fix logic for the assert in setExtraInfo

[akhuang]

Sounds good. I also realized that the logic for the "Should never have only a single symbol allocated out-of-line!" is wrong now. I added a parameter to the function to keep track of whether the assert should happen. Or do you think it can just be removed?

[rnk]

llvm/lib/CodeGen/MachineInstr.cpp

336–340

I think you could get rid of all this and the parameter. It seems trivial with the new code structure.

[akhuang]

-remove unnecessary assert, rebase

llvm/include/llvm/CodeGen/MachineFunction.h

class MachineFunction {
 
   /// Map of invoke call site index values to associated begin EH_LABEL.
   DenseMap<MCSymbol*, unsigned> CallSiteMap;
 
   /// CodeView label annotations.
   std::vector<std::pair<MCSymbol *, MDNode *>> CodeViewAnnotations;
 
   /// CodeView heapallocsites.
   std::vector<std::tuple<MCSymbol *, MCSymbol *, const DIType *>>
       CodeViewHeapAllocSites;
+  DenseMap<MachineInstr *, const DIType *> HeapAllocSiteInstrMap;
 
   bool CallsEHReturn = false;
   bool CallsUnwindInit = false;
   bool HasEHScopes = false;
   bool HasEHFunclets = false;
 
   /// List of C++ TypeInfo used.
   std::vector<const GlobalValue *> TypeInfos;
public:
 
   /// Allocate and initialize a register mask with @p NumRegister bits.
   uint32_t *allocateRegMask();
 
   /// Allocate and construct an extra info structure for a `MachineInstr`.
   ///
   /// This is allocated on the function's allocator and so lives the life of
   /// the function.
-  MachineInstr::ExtraInfo *
-  createMIExtraInfo(ArrayRef<MachineMemOperand *> MMOs,
-                    MCSymbol *PreInstrSymbol = nullptr,
-                    MCSymbol *PostInstrSymbol = nullptr);
+  MachineInstr::ExtraInfo *createMIExtraInfo(
+      ArrayRef<MachineMemOperand *> MMOs, MCSymbol *PreInstrSymbol = nullptr,
+      MCSymbol *PostInstrSymbol = nullptr, MCSymbol *InstrMarker = nullptr);
 
   /// Allocate a string and populate it with the given external symbol name.
   const char *createExternalSymbolName(StringRef Name);
 
   //===--------------------------------------------------------------------===//
   // Label Manipulation.
 
   /// getJTISymbol - Return the MCSymbol for the specified non-empty jump table.
public:
   void addCodeViewAnnotation(MCSymbol *Label, MDNode *MD) {
     CodeViewAnnotations.push_back({Label, MD});
   }
 
   ArrayRef<std::pair<MCSymbol *, MDNode *>> getCodeViewAnnotations() const {
     return CodeViewAnnotations;
   }
 
-  /// Record heapallocsites
-  void addCodeViewHeapAllocSite(MachineInstr *I, const MDNode *MD);
+  /// Map heapallocsite instructions to debug type and add a heapallocsite
+  /// marker to the instruction.
+  void setCodeViewHeapAllocSite(MachineInstr *I, const MDNode *MD);
+
+  /// Record heapallocsites to be emitted in CodeView debug info.
+  void addCodeViewHeapAllocSite(MCSymbol *Begin, MCSymbol *End,
+                                const DIType *DI);
 
   ArrayRef<std::tuple<MCSymbol *, MCSymbol *, const DIType *>>
   getCodeViewHeapAllocSites() const {
     return CodeViewHeapAllocSites;
   }
 
+  const DIType *getHeapAllocSiteType(MachineInstr *I) const {
+    auto It = HeapAllocSiteInstrMap.find(I);
+    if (It == HeapAllocSiteInstrMap.end())
+      return nullptr;
+    return It->second;
+  }
+
   /// Return a reference to the C++ typeinfo for the current function.
   const std::vector<const GlobalValue *> &getTypeInfos() const {
     return TypeInfos;
   }
 
   /// Return a reference to the typeids encoding filters used in the current
   /// function.
   const std::vector<unsigned> &getFilterIds() const {

llvm/include/llvm/CodeGen/MachineInstr.h

private:
   /// This has to be defined eagerly due to the implementation constraints of
   /// `PointerSumType` where it is used.
   class ExtraInfo final
       : TrailingObjects<ExtraInfo, MachineMemOperand *, MCSymbol *> {
   public:
     static ExtraInfo *create(BumpPtrAllocator &Allocator,
                              ArrayRef<MachineMemOperand *> MMOs,
                              MCSymbol *PreInstrSymbol = nullptr,
-                             MCSymbol *PostInstrSymbol = nullptr) {
+                             MCSymbol *PostInstrSymbol = nullptr,
+                             MCSymbol *InstrMarker = nullptr) {
       bool HasPreInstrSymbol = PreInstrSymbol != nullptr;
       bool HasPostInstrSymbol = PostInstrSymbol != nullptr;
+      bool HasInstrMarker = InstrMarker != nullptr;
       auto *Result = new (Allocator.Allocate(
           totalSizeToAlloc<MachineMemOperand *, MCSymbol *>(
-              MMOs.size(), HasPreInstrSymbol + HasPostInstrSymbol),
-          alignof(ExtraInfo)))
-          ExtraInfo(MMOs.size(), HasPreInstrSymbol, HasPostInstrSymbol);
+              MMOs.size(),
+              HasPreInstrSymbol + HasPostInstrSymbol + HasInstrMarker),
+          alignof(ExtraInfo))) ExtraInfo(MMOs.size(), HasPreInstrSymbol,
+                                         HasPostInstrSymbol, HasInstrMarker);
 
       // Copy the actual data into the trailing objects.
       std::copy(MMOs.begin(), MMOs.end(),
                 Result->getTrailingObjects<MachineMemOperand *>());
 
       if (HasPreInstrSymbol)
         Result->getTrailingObjects<MCSymbol *>()[0] = PreInstrSymbol;
       if (HasPostInstrSymbol)
         Result->getTrailingObjects<MCSymbol *>()[HasPreInstrSymbol] =
             PostInstrSymbol;
+      if (HasInstrMarker)
+        Result->getTrailingObjects<MCSymbol *>()[HasPreInstrSymbol +
+                                                 HasPostInstrSymbol] =
+            InstrMarker;
 
       return Result;
     }
 
     ArrayRef<MachineMemOperand *> getMMOs() const {
       return makeArrayRef(getTrailingObjects<MachineMemOperand *>(), NumMMOs);
     }
 
     MCSymbol *getPreInstrSymbol() const {
       return HasPreInstrSymbol ? getTrailingObjects<MCSymbol *>()[0] : nullptr;
     }
 
     MCSymbol *getPostInstrSymbol() const {
       return HasPostInstrSymbol
                  ? getTrailingObjects<MCSymbol *>()[HasPreInstrSymbol]
                  : nullptr;
     }
 
+    MCSymbol *getInstrMarker() const {
+      return HasInstrMarker
+                 ? getTrailingObjects<MCSymbol *>()[HasPreInstrSymbol +
+                                                    HasPostInstrSymbol]
+                 : nullptr;
+    }
+
   private:
     friend TrailingObjects;
 
     // Description of the extra info, used to interpret the actual optional
     // data appended.
     //
     // Note that this is not terribly space optimized. This leaves a great deal
     // of flexibility to fit more in here later.
     const int NumMMOs;
     const bool HasPreInstrSymbol;
     const bool HasPostInstrSymbol;
+    const bool HasInstrMarker;
 
     // Implement the `TrailingObjects` internal API.
     size_t numTrailingObjects(OverloadToken<MachineMemOperand *>) const {
       return NumMMOs;
     }
     size_t numTrailingObjects(OverloadToken<MCSymbol *>) const {
-      return HasPreInstrSymbol + HasPostInstrSymbol;
+      return HasPreInstrSymbol + HasPostInstrSymbol + HasInstrMarker;
     }
 
     // Just a boring constructor to allow us to initialize the sizes. Always use
     // the `create` routine above.
-    ExtraInfo(int NumMMOs, bool HasPreInstrSymbol, bool HasPostInstrSymbol)
+    ExtraInfo(int NumMMOs, bool HasPreInstrSymbol, bool HasPostInstrSymbol,
+              bool HasInstrMarker)
         : NumMMOs(NumMMOs), HasPreInstrSymbol(HasPreInstrSymbol),
-          HasPostInstrSymbol(HasPostInstrSymbol) {}
+          HasPostInstrSymbol(HasPostInstrSymbol),
+          HasInstrMarker(HasInstrMarker) {}
   };
 
   /// Enumeration of the kinds of inline extra info available. It is important
   /// that the `MachineMemOperand` inline kind has a tag value of zero to make
   /// it accessible as an `ArrayRef`.
   enum ExtraInfoInlineKinds {
     EIIK_MMO = 0,
     EIIK_PreInstrSymbol,
     EIIK_PostInstrSymbol,
-    EIIK_OutOfLine
+    EIIK_InstrMarker,
+    EIIK_OutOfLine,
   };
 
   // We store extra information about the instruction here. The common case is
   // expected to be nothing or a single pointer (typically a MMO or a symbol).
   // We work to optimize this common case by storing it inline here rather than
   // requiring a separate allocation, but we fall back to an allocation when
   // multiple pointers are needed.
   PointerSumType<ExtraInfoInlineKinds,
                  PointerSumTypeMember<EIIK_MMO, MachineMemOperand *>,
                  PointerSumTypeMember<EIIK_PreInstrSymbol, MCSymbol *>,
                  PointerSumTypeMember<EIIK_PostInstrSymbol, MCSymbol *>,
+                 PointerSumTypeMember<EIIK_InstrMarker, MCSymbol *>,
                  PointerSumTypeMember<EIIK_OutOfLine, ExtraInfo *>>
       Info;
 
   DebugLoc debugLoc;                    // Source line information.
 
   // Intrusive list support
   friend struct ilist_traits<MachineInstr>;
   friend struct ilist_callback_traits<MachineBasicBlock>;
MCSymbol *getPostInstrSymbol() const {
     if (MCSymbol *S = Info.get<EIIK_PostInstrSymbol>())
       return S;
     if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
       return EI->getPostInstrSymbol();
 
     return nullptr;
   }
 
+  MCSymbol *getInstrMarker() const {
+    if (!Info)
+      return nullptr;
+    if (MCSymbol *S = Info.get<EIIK_InstrMarker>())
+      return S;
+    if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
+      return EI->getInstrMarker();
+
+    return nullptr;
+  }
+
   /// API for querying MachineInstr properties. They are the same as MCInstrDesc
   /// queries but they are bundle aware.
 
   enum QueryType {
     IgnoreBundle,    // Ignore bundles
     AnyInBundle,     // Return true if any instruction in bundle has property
     AllInBundle      // Return true if all instructions in bundle have property
   };

   ///
   /// FIXME: This is not fully implemented yet.
   void setPostInstrSymbol(MachineFunction &MF, MCSymbol *Symbol);
 
   /// Clone another MachineInstr's pre- and post- instruction symbols and
   /// replace ours with it.
   void cloneInstrSymbols(MachineFunction &MF, const MachineInstr &MI);
 
+  /// Set a marker on instructions that denotes where we should create and emit
+  /// labels. This waits until after instruction selection and optimizations to
+  /// create the label, so it should still work if the instruction is removed
+  /// or duplicated.
+  void setInstrMarker(MachineFunction &MF, MCSymbol *Symbol);
+
   /// Return the MIFlags which represent both MachineInstrs. This
   /// should be used when merging two MachineInstrs into one. This routine does
   /// not modify the MIFlags of this MachineInstr.
   uint16_t mergeFlagsWith(const MachineInstr& Other) const;
 
   static uint16_t copyFlagsFromInstruction(const Instruction &I);
 
   /// Copy all flags to MachineInst MIFlags

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

for (auto &MBB : *MF) {
     for (auto &MI : MBB) {
       // Print the assembly for the instruction.
       if (!MI.isPosition() && !MI.isImplicitDef() && !MI.isKill() &&
           !MI.isDebugInstr()) {
         HasAnyRealCode = true;
         ++NumInstsInFunction;
       }
 
-      // If there is a pre-instruction symbol, emit a label for it here. If the
-      // instruction was duplicated and the label has already been emitted,
-      // don't re-emit the same label.
-      // FIXME: Consider strengthening that to an assertion.
+      // If there is a pre-instruction symbol, emit a label for it here.
       if (MCSymbol *S = MI.getPreInstrSymbol())
-        if (S->isUndefined())
-          OutStreamer->EmitLabel(S);
 
+      // If there is a marker, create and emit a label before the instruction.
+      MCSymbol *PreInstrLabel = nullptr;
+      if (MCSymbol *S = MI.getInstrMarker()) {
+        PreInstrLabel = createTempSymbol(S->getName());
+        OutStreamer->EmitLabel(PreInstrLabel);
+      }
+
       if (ShouldPrintDebugScopes) {
         for (const HandlerInfo &HI : Handlers) {
           NamedRegionTimer T(HI.TimerName, HI.TimerDescription,
                              HI.TimerGroupName, HI.TimerGroupDescription,
                              TimePassesIsEnabled);
           HI.Handler->beginInstruction(&MI);
         }
       }
for (auto &MI : MBB) {
       case TargetOpcode::KILL:
         if (isVerbose()) emitKill(&MI, *this);
         break;
       default:
         EmitInstruction(&MI);
         break;
       }
 
-      // If there is a post-instruction symbol, emit a label for it here.  If
-      // the instruction was duplicated and the label has already been emitted,
-      // don't re-emit the same label.
-      // FIXME: Consider strengthening that to an assertion.
+      // If there is a post-instruction symbol, emit a label for it here.
       if (MCSymbol *S = MI.getPostInstrSymbol())
-        if (S->isUndefined())
-          OutStreamer->EmitLabel(S);
 
+      // If there is a marker, create and emit a label after the instruction.
+      if (MCSymbol *S = MI.getInstrMarker()) {
+        StringRef LabelName = S->getName();
+        MCSymbol *PostInstrLabel = createTempSymbol(LabelName);
+        OutStreamer->EmitLabel(PostInstrLabel);
+
+        // If this is a heapallocsite symbol, add to the list of CodeView
+        // records.
+        if (LabelName == "heapallocsite") {
+          assert(PreInstrLabel && "Unmatched heap alloc site label");
+          const DIType *DI = MI.getMF()->getHeapAllocSiteType(&MI);
+          MI.getMF()->addCodeViewHeapAllocSite(PreInstrLabel, PostInstrLabel,
+                                               DI);
+        }
+      }
+
       if (ShouldPrintDebugScopes) {
         for (const HandlerInfo &HI : Handlers) {
           NamedRegionTimer T(HI.TimerName, HI.TimerDescription,
                              HI.TimerGroupName, HI.TimerGroupDescription,
                              TimePassesIsEnabled);
           HI.Handler->endInstruction();
         }
       }

llvm/lib/CodeGen/AsmPrinter/CodeViewDebug.cpp

for (auto Annot : FI.Annotations) {
         OS.EmitBytes(StringRef(Str.data(), Str.size() + 1));
       }
       endSymbolRecord(AnnotEnd);
     }
 
     for (auto HeapAllocSite : FI.HeapAllocSites) {
       MCSymbol *BeginLabel = std::get<0>(HeapAllocSite);
       MCSymbol *EndLabel = std::get<1>(HeapAllocSite);
-
-      // The labels might not be defined if the instruction was replaced
-      // somewhere in the codegen pipeline.
-      if (!BeginLabel->isDefined() || !EndLabel->isDefined())
-        continue;
-
       const DIType *DITy = std::get<2>(HeapAllocSite);
       MCSymbol *HeapAllocEnd = beginSymbolRecord(SymbolKind::S_HEAPALLOCSITE);
       OS.AddComment("Call site offset");
       OS.EmitCOFFSecRel32(BeginLabel, /*Offset=*/0);
       OS.AddComment("Call site section index");
       OS.EmitCOFFSectionIndex(BeginLabel);
       OS.AddComment("Call instruction length");
       OS.emitAbsoluteSymbolDiff(EndLabel, BeginLabel, 2);
void CodeViewDebug::beginFunctionImpl(const MachineFunction *MF) {
   OS.EmitCVFuncIdDirective(CurFn->FuncId);
 
   // Find the end of the function prolog.  First known non-DBG_VALUE and
   // non-frame setup location marks the beginning of the function body.
   // FIXME: is there a simpler a way to do this? Can we just search
   // for the first instruction of the function, not the last of the prolog?
   DebugLoc PrologEndLoc;
   bool EmptyPrologue = true;
   for (const auto &MBB : *MF) {
     for (const auto &MI : MBB) {
       if (!MI.isMetaInstruction() && !MI.getFlag(MachineInstr::FrameSetup) &&
           MI.getDebugLoc()) {
         PrologEndLoc = MI.getDebugLoc();
         break;
       } else if (!MI.isMetaInstruction()) {
         EmptyPrologue = false;
       }
     }

llvm/lib/CodeGen/MachineFunction.cpp

 MachineFunction::getMachineMemOperand(const MachineMemOperand *MMO,
                                       MachineMemOperand::Flags Flags) {
   return new (Allocator) MachineMemOperand(
       MMO->getPointerInfo(), Flags, MMO->getSize(), MMO->getBaseAlignment(),
       MMO->getAAInfo(), MMO->getRanges(), MMO->getSyncScopeID(),
       MMO->getOrdering(), MMO->getFailureOrdering());
 }
 
-MachineInstr::ExtraInfo *
-MachineFunction::createMIExtraInfo(ArrayRef<MachineMemOperand *> MMOs,
-                                   MCSymbol *PreInstrSymbol,
-                                   MCSymbol *PostInstrSymbol) {
+MachineInstr::ExtraInfo *MachineFunction::createMIExtraInfo(
+    ArrayRef<MachineMemOperand *> MMOs, MCSymbol *PreInstrSymbol,
+    MCSymbol *PostInstrSymbol, MCSymbol *InstrMarker) {
   return MachineInstr::ExtraInfo::create(Allocator, MMOs, PreInstrSymbol,
-                                         PostInstrSymbol);
+                                         PostInstrSymbol, InstrMarker);
 }
 
 const char *MachineFunction::createExternalSymbolName(StringRef Name) {
   char *Dest = Allocator.Allocate<char>(Name.size() + 1);
   llvm::copy(Name, Dest);
   Dest[Name.size()] = 0;
   return Dest;
 }
try_next:;
   int FilterID = -(1 + FilterIds.size());
   FilterIds.reserve(FilterIds.size() + TyIds.size() + 1);
   FilterIds.insert(FilterIds.end(), TyIds.begin(), TyIds.end());
   FilterEnds.push_back(FilterIds.size());
   FilterIds.push_back(0); // terminator
   return FilterID;
 }
 
-void MachineFunction::addCodeViewHeapAllocSite(MachineInstr *I,
+void MachineFunction::setCodeViewHeapAllocSite(MachineInstr *I,
                                                const MDNode *MD) {
-  MCSymbol *BeginLabel = Ctx.createTempSymbol("heapallocsite", true);
-  MCSymbol *EndLabel = Ctx.createTempSymbol("heapallocsite", true);
-  I->setPreInstrSymbol(*this, BeginLabel);
-  I->setPostInstrSymbol(*this, EndLabel);
-
+  MCSymbol *Sym = Ctx.getOrCreateSymbol("heapallocsite");
+  I->setInstrMarker(*this, Sym);
   const DIType *DI = dyn_cast<DIType>(MD);
-  CodeViewHeapAllocSites.push_back(std::make_tuple(BeginLabel, EndLabel, DI));
+  HeapAllocSiteInstrMap[I] = DI;
+}
+
+void MachineFunction::addCodeViewHeapAllocSite(MCSymbol *Begin, MCSymbol *End,
+                                               const DIType *DI) {
+  CodeViewHeapAllocSites.push_back(std::make_tuple(Begin, End, DI));
 }
 
 void MachineFunction::moveCallSiteInfo(const MachineInstr *Old,
                                        const MachineInstr *New) {
   assert(New->isCall() && "Call site info refers only to call instructions!");
 
   CallSiteInfoMap::iterator CSIt = getCallSiteInfo(Old);
   if (CSIt == CallSitesInfo.end())

llvm/lib/CodeGen/MachineInstr.cpp

void MachineInstr::dropMemRefs(MachineFunction &MF) {
   if (!getPostInstrSymbol()) {
     Info.set<EIIK_PreInstrSymbol>(getPreInstrSymbol());
     return;
   }
   if (!getPreInstrSymbol()) {
     Info.set<EIIK_PostInstrSymbol>(getPostInstrSymbol());
     return;
   }
+  if (!getInstrMarker()) {
+    Info.set<EIIK_InstrMarker>(getInstrMarker());
+    return;
+  }
 
   // Otherwise allocate a fresh extra info with just these symbols.
-  Info.set<EIIK_OutOfLine>(
-      MF.createMIExtraInfo({}, getPreInstrSymbol(), getPostInstrSymbol()));
+  Info.set<EIIK_OutOfLine>(MF.createMIExtraInfo(
+      {}, getPreInstrSymbol(), getPostInstrSymbol(), getInstrMarker()));
 }
 
 void MachineInstr::setMemRefs(MachineFunction &MF,
                               ArrayRef<MachineMemOperand *> MMOs) {
   if (MMOs.empty()) {
     dropMemRefs(MF);
     return;
   }
 
   // Try to store a single MMO inline.
-  if (MMOs.size() == 1 && !getPreInstrSymbol() && !getPostInstrSymbol()) {
+  if (MMOs.size() == 1 && !getPreInstrSymbol() && !getPostInstrSymbol() &&
+      !getInstrMarker()) {
     Info.set<EIIK_MMO>(MMOs[0]);
     return;
   }
 
   // Otherwise create an extra info struct with all of our info.
-  Info.set<EIIK_OutOfLine>(
-      MF.createMIExtraInfo(MMOs, getPreInstrSymbol(), getPostInstrSymbol()));
+  Info.set<EIIK_OutOfLine>(MF.createMIExtraInfo(
+      MMOs, getPreInstrSymbol(), getPostInstrSymbol(), getInstrMarker()));
 }
 
 void MachineInstr::addMemOperand(MachineFunction &MF,
                                  MachineMemOperand *MO) {
   SmallVector<MachineMemOperand *, 2> MMOs;
   MMOs.append(memoperands_begin(), memoperands_end());
   MMOs.push_back(MO);
   setMemRefs(MF, MMOs);
 }
 
 void MachineInstr::cloneMemRefs(MachineFunction &MF, const MachineInstr &MI) {
   if (this == &MI)
     // Nothing to do for a self-clone!
     return;
 
   assert(&MF == MI.getMF() &&
          "Invalid machine functions when cloning memory refrences!");
   // See if we can just steal the extra info already allocated for the
   // instruction. We can do this whenever the pre- and post-instruction symbols
   // are the same (including null).
   if (getPreInstrSymbol() == MI.getPreInstrSymbol() &&
-      getPostInstrSymbol() == MI.getPostInstrSymbol()) {
+      getPostInstrSymbol() == MI.getPostInstrSymbol() &&
+      getInstrMarker() == MI.getInstrMarker()) {
     Info = MI.Info;
     return;
   }
 
   // Otherwise, fall back on a copy-based clone.
   setMemRefs(MF, MI.memoperands());
 }
 
if (OldSymbol && !Symbol) {
     // If we don't have any other extra info, we can store this inline.
     Info.set<EIIK_PostInstrSymbol>(Symbol);
     return;
   }
 
   // Otherwise, allocate a full new set of extra info.
   // FIXME: Maybe we should make the symbols in the extra info mutable?
   Info.set<EIIK_OutOfLine>(
       MF.createMIExtraInfo(memoperands(), getPreInstrSymbol(), Symbol));
 }
 
 void MachineInstr::cloneInstrSymbols(MachineFunction &MF,
                                      const MachineInstr &MI) {
   if (this == &MI)
     // Nothing to do for a self-clone!
     return;
 
   assert(&MF == MI.getMF() &&
          "Invalid machine functions when cloning instruction symbols!");
 
   setPreInstrSymbol(MF, MI.getPreInstrSymbol());
   setPostInstrSymbol(MF, MI.getPostInstrSymbol());
+  setInstrMarker(MF, MI.getInstrMarker());
+}
+
+void MachineInstr::setInstrMarker(MachineFunction &MF, MCSymbol *Marker) {
+  MCSymbol *OldMarker = getInstrMarker();
+  if (OldMarker == Marker)
+    return;
+  if (OldMarker && !Marker) {
+    // We're removing a symbol rather than adding one. Try to clean up any
+    // extra info carried around.
+    if (Info.is<EIIK_InstrMarker>()) {
+      Info.clear();
+      return;
+    }
+
+    if (memoperands_empty()) {
+      Info.set<EIIK_PreInstrSymbol>(getPreInstrSymbol());
+      Info.set<EIIK_PostInstrSymbol>(getPostInstrSymbol());
+      return;
+    }
+
+    // Otherwise fallback on the generic update.
+  } else if (!Info || Info.is<EIIK_InstrMarker>()) {
+    // If we don't have any other extra info, we can store this inline.
+    Info.set<EIIK_InstrMarker>(Marker);
+    return;
+  }
+  // Otherwise, allocate a full new set of extra info.
+  Info.set<EIIK_OutOfLine>(MF.createMIExtraInfo(
+      memoperands(), getPreInstrSymbol(), getPostInstrSymbol(), Marker));
 }
 
 uint16_t MachineInstr::mergeFlagsWith(const MachineInstr &Other) const {
   // For now, the just return the union of the flags. If the flags get more
   // complicated over time, we might need more logic here.
   return getFlags() | Other.getFlags();
 }
 
void MachineInstr::print(raw_ostream &OS, ModuleSlotTracker &MST,
   if (MCSymbol *PostInstrSymbol = getPostInstrSymbol()) {
     if (!FirstOp) {
       FirstOp = false;
       OS << ',';
     }
     OS << " post-instr-symbol ";
     MachineOperand::printSymbol(OS, *PostInstrSymbol);
   }
+  if (MCSymbol *InstrMarker = getInstrMarker()) {
+    if (!FirstOp) {
+      FirstOp = false;
+      OS << ',';
+    }
+    OS << " instr-marker ";
+    MachineOperand::printSymbol(OS, *InstrMarker);
+  }
 
   if (!SkipDebugLoc) {
     if (const DebugLoc &DL = getDebugLoc()) {
       if (!FirstOp)
         OS << ',';
       OS << " debug-location ";
       DL->printAsOperand(OS, MST);
     }

llvm/lib/CodeGen/SelectionDAG/FastISel.cpp

bool FastISel::lowerCallTo(CallLoweringInfo &CLI) {
   CLI.Call->setPhysRegsDeadExcept(CLI.InRegs, TRI);
 
   if (CLI.NumResultRegs && CLI.CS)
     updateValueMap(CLI.CS->getInstruction(), CLI.ResultReg, CLI.NumResultRegs);
 
   // Set labels for heapallocsite call.
   if (CLI.CS && CLI.CS->getInstruction()->hasMetadata("heapallocsite")) {
     const MDNode *MD = CLI.CS->getInstruction()->getMetadata("heapallocsite");
-    MF->addCodeViewHeapAllocSite(CLI.Call, MD);
+    MF->setCodeViewHeapAllocSite(CLI.Call, MD);
   }
 
   return true;
 }
 
 bool FastISel::lowerCall(const CallInst *CI) {
   ImmutableCallSite CS(CI);
 

llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp

while (!GluedNodes.empty()) {
       SDNode *N = GluedNodes.back();
       auto NewInsn = EmitNode(N, SU->OrigNode != SU, SU->isCloned, VRBaseMap);
       // Remember the source order of the inserted instruction.
       if (HasDbg)
         ProcessSourceNode(N, DAG, Emitter, VRBaseMap, Orders, Seen, NewInsn);
 
       if (MDNode *MD = DAG->getHeapAllocSite(N)) {
         if (NewInsn && NewInsn->isCall())
-          MF.addCodeViewHeapAllocSite(NewInsn, MD);
+          MF.setCodeViewHeapAllocSite(NewInsn, MD);
       }
 
       GluedNodes.pop_back();
     }
     auto NewInsn =
         EmitNode(SU->getNode(), SU->OrigNode != SU, SU->isCloned, VRBaseMap);
     // Remember the source order of the inserted instruction.
     if (HasDbg)
       ProcessSourceNode(SU->getNode(), DAG, Emitter, VRBaseMap, Orders, Seen,
                         NewInsn);
     if (MDNode *MD = DAG->getHeapAllocSite(SU->getNode())) {
       if (NewInsn && NewInsn->isCall())
-        MF.addCodeViewHeapAllocSite(NewInsn, MD);
+        MF.setCodeViewHeapAllocSite(NewInsn, MD);
     }
   }
 
   // Insert all the dbg_values which have not already been inserted in source
   // order sequence.
   if (HasDbg) {
     MachineBasicBlock::iterator BBBegin = BB->getFirstNonPHI();
 

llvm/test/CodeGen/X86/label-heapallocsite.ll

-; RUN: llc < %s | FileCheck --check-prefixes=DAG,CHECK %s
-; RUN: llc -O0 < %s | FileCheck --check-prefixes=FAST,CHECK %s
+; RUN: llc < %s | FileCheck --check-prefixes=CHECK %s
+; RUN: llc -O0 < %s | FileCheck --check-prefixes=CHECK %s
 
 ; Source to regenerate:
 ; $ clang -cc1 -triple x86_64-windows-msvc t.cpp -debug-info-kind=limited \
 ;      -gcodeview -O2 -fms-extensions -emit-llvm -o t.ll
 ;
 ; extern "C" struct Foo {
 ;   __declspec(allocator) virtual void *alloc();
 ; };

 ; CHECK: jmp alloc_foo
 
 ; CHECK-LABEL: call_virtual:      # @call_virtual
 ; CHECK: .Lheapallocsite0:
 ; CHECK: callq *{{.*}}%rax{{.*}}
 ; CHECK: .Lheapallocsite1:
 
 ; CHECK-LABEL: call_multiple:     # @call_multiple
-; FastISel emits instructions in a different order.
-; DAG:   .Lheapallocsite2:
-; FAST:  .Lheapallocsite4:
+; CHECK:   .Lheapallocsite2:
 ; CHECK: callq alloc_foo
-; DAG:   .Lheapallocsite3:
-; FAST:  .Lheapallocsite5:
-; DAG:   .Lheapallocsite4:
-; FAST:  .Lheapallocsite2:
+; CHECK:   .Lheapallocsite3:
+; CHECK:   .Lheapallocsite4:
 ; CHECK: callq alloc_foo
-; DAG:   .Lheapallocsite5:
-; FAST:  .Lheapallocsite3:
+; CHECK:   .Lheapallocsite5:
 
 ; CHECK-LABEL: .short  4423                    # Record kind: S_GPROC32_ID
 ; CHECK:       .short  4446                    # Record kind: S_HEAPALLOCSITE
 ; CHECK-NEXT:  .secrel32 .Lheapallocsite0
 ; CHECK-NEXT:  .secidx .Lheapallocsite0
 ; CHECK-NEXT:  .short .Lheapallocsite1-.Lheapallocsite0
 ; CHECK-NEXT:  .long 3
 ; CHECK:       .short  4446                    # Record kind: S_HEAPALLOCSITE

llvm/test/CodeGen/X86/taildup-heapallocsite.ll

 ; RUN: llc < %s -tail-dup-placement-threshold=4 | FileCheck %s
 
 ; Based on test case from PR43695:
 ; __declspec(allocator) void *alloc(unsigned int size);
 ; void f2();
 ; void f1(unsigned int *size_ptr) {
 ;     void *hg = alloc(size_ptr ? *size_ptr : 1UL);
 ;     f2();
 ; }
 
-; In this case, block placement duplicates the heap allocation site. For now,
-; LLVM drops the labels from one call site. Eventually, we should track both.
+; In this case, block placement duplicates the heap allocation site.
 
 ; ModuleID = 't.cpp'
 source_filename = "t.cpp"
 target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-pc-windows-msvc19.22.27905"
 
 define dso_local void @taildupit(i32* readonly %size_ptr) !dbg !8 {
 entry:

 
 ; CHECK-LABEL: taildupit: # @taildupit
 ; CHECK: testq
 ; CHECK: je
 ; CHECK: .Lheapallocsite0:
 ; CHECK: callq alloc
 ; CHECK: .Lheapallocsite1:
 ; CHECK: jmp f2 # TAILCALL
-; CHECK-NOT: Lheapallocsite
+; CHECK: Lheapallocsite2
 ; CHECK: callq alloc
-; CHECK-NOT: Lheapallocsite
+; CHECK: Lheapallocsite3
 ; CHECK: jmp f2 # TAILCALL
 
 declare dso_local i8* @alloc(i32)
 
 declare dso_local void @f2()
 
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!3, !4, !5, !6}