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

324–325

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.

325

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

248

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

508–509

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

539–540

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

539–540

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

45–57

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

331–335

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 a landing pad to its index.
   DenseMap<const MachineBasicBlock *, unsigned> WasmLPadToIndexMap;
 
   /// 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;
-
   bool CallsEHReturn = false;
   bool CallsUnwindInit = false;
   bool HasEHScopes = false;
   bool HasEHFunclets = false;
 
   /// List of C++ TypeInfo used.
   std::vector<const GlobalValue *> TypeInfos;
 
   /// List of typeids encoding filters used.
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, MDNode *HeapAllocMarker = 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);
-
-  ArrayRef<std::tuple<MCSymbol *, MCSymbol *, const DIType *>>
-  getCodeViewHeapAllocSites() const {
-    return CodeViewHeapAllocSites;
-  }
-
   /// 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:
 
   /// Internal implementation detail class that provides out-of-line storage for
   /// extra info used by the machine instruction when this info cannot be stored
   /// in-line within the instruction itself.
   ///
   /// This has to be defined eagerly due to the implementation constraints of
   /// `PointerSumType` where it is used.
   class ExtraInfo final
-      : TrailingObjects<ExtraInfo, MachineMemOperand *, MCSymbol *> {
+      : TrailingObjects<ExtraInfo, MachineMemOperand *, MCSymbol *, MDNode *> {
   public:
     static ExtraInfo *create(BumpPtrAllocator &Allocator,
                              ArrayRef<MachineMemOperand *> MMOs,
                              MCSymbol *PreInstrSymbol = nullptr,
-                             MCSymbol *PostInstrSymbol = nullptr) {
+                             MCSymbol *PostInstrSymbol = nullptr,
+                             MDNode *HeapAllocMarker = nullptr) {
       bool HasPreInstrSymbol = PreInstrSymbol != nullptr;
       bool HasPostInstrSymbol = PostInstrSymbol != nullptr;
+      bool HasHeapAllocMarker = HeapAllocMarker != nullptr;
       auto *Result = new (Allocator.Allocate(
-          totalSizeToAlloc<MachineMemOperand *, MCSymbol *>(
-              MMOs.size(), HasPreInstrSymbol + HasPostInstrSymbol),
+          totalSizeToAlloc<MachineMemOperand *, MCSymbol *, MDNode *>(
+              MMOs.size(), HasPreInstrSymbol + HasPostInstrSymbol,
+              HasHeapAllocMarker),
           alignof(ExtraInfo)))
-          ExtraInfo(MMOs.size(), HasPreInstrSymbol, HasPostInstrSymbol);
+          ExtraInfo(MMOs.size(), HasPreInstrSymbol, HasPostInstrSymbol,
+                    HasHeapAllocMarker);
 
       // 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 (HasHeapAllocMarker)
+        Result->getTrailingObjects<MDNode *>()[HasPreInstrSymbol +
+                                               HasPostInstrSymbol] =
+            HeapAllocMarker;
 
       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;
     }
 
+    MDNode *getHeapAllocMarker() const {
+      return HasHeapAllocMarker
+                 ? getTrailingObjects<MDNode *>()[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 HasHeapAllocMarker;
 
     // Implement the `TrailingObjects` internal API.
     size_t numTrailingObjects(OverloadToken<MachineMemOperand *>) const {
       return NumMMOs;
     }
     size_t numTrailingObjects(OverloadToken<MCSymbol *>) const {
       return HasPreInstrSymbol + HasPostInstrSymbol;
     }
+    size_t numTrailingObjects(OverloadToken<MDNode *>) const {
+      return HasHeapAllocMarker;
+    }
 
     // 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 HasHeapAllocMarker)
         : NumMMOs(NumMMOs), HasPreInstrSymbol(HasPreInstrSymbol),
-          HasPostInstrSymbol(HasPostInstrSymbol) {}
+          HasPostInstrSymbol(HasPostInstrSymbol),
+          HasHeapAllocMarker(HasHeapAllocMarker) {}
   };
 
   /// 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_HeapAllocMarker,
+    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_HeapAllocMarker, MDNode *>,
                  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;
   }
 
+  MDNode *getHeapAllocMarker() const {
+    if (!Info)
+      return nullptr;
+    if (MDNode *S = Info.get<EIIK_HeapAllocMarker>())
+      return S;
+    if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
+      return EI->getHeapAllocMarker();
+
+    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
+  /// heap alloc site 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 setHeapAllocMarker(MachineFunction &MF, MDNode *DI);
+
   /// 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
private:
   bool hasPropertyInBundle(uint64_t Mask, QueryType Type) const;
 
   /// Implements the logic of getRegClassConstraintEffectForVReg for the
   /// this MI and the given operand index \p OpIdx.
   /// If the related operand does not constrained Reg, this returns CurRC.
   const TargetRegisterClass *getRegClassConstraintEffectForVRegImpl(
       unsigned OpIdx, Register Reg, const TargetRegisterClass *CurRC,
       const TargetInstrInfo *TII, const TargetRegisterInfo *TRI) const;
+
+  /// Stores extra instruction information inline or allocates as ExtraInfo
+  /// based on the number of pointers.
+  void setExtraInfo(MachineFunction &MF, ArrayRef<MachineMemOperand *> MMOs,
+                    MCSymbol *PreInstrSymbol, MCSymbol *PostInstrSymbol,
+                    MDNode *HeapAllocMarker);
 };
 
 /// Special DenseMapInfo traits to compare MachineInstr* by *value* of the
 /// instruction rather than by pointer value.
 /// The hashing and equality testing functions ignore definitions so this is
 /// useful for CSE, etc.
 struct MachineInstrExpressionTrait : DenseMapInfo<MachineInstr*> {
   static inline MachineInstr *getEmptyKey() {

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 (ShouldPrintDebugScopes) {
         for (const HandlerInfo &HI : Handlers) {
           NamedRegionTimer T(HI.TimerName, HI.TimerDescription,
                              HI.TimerGroupName, HI.TimerGroupDescription,
                              TimePassesIsEnabled);
           HI.Handler->beginInstruction(&MI);
         }
       }
 
       if (isVerbose())
         emitComments(MI, OutStreamer->GetCommentOS());
 
       switch (MI.getOpcode()) {
       case TargetOpcode::CFI_INSTRUCTION:
         emitCFIInstruction(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 (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.h

struct FunctionInfo {
     SmallVector<CVGlobalVariable, 1> Globals;
 
     std::unordered_map<const DILexicalBlockBase*, LexicalBlock> LexicalBlocks;
 
     // Lexical blocks containing local variables.
     SmallVector<LexicalBlock *, 1> ChildBlocks;
 
     std::vector<std::pair<MCSymbol *, MDNode *>> Annotations;
-    std::vector<std::tuple<MCSymbol *, MCSymbol *, const DIType *>>
+    std::vector<std::tuple<const MCSymbol *, const MCSymbol *, const DIType *>>
         HeapAllocSites;
 
     const MCSymbol *Begin = nullptr;
     const MCSymbol *End = nullptr;
     unsigned FuncId = 0;
     unsigned LastFileId = 0;
 
     /// Number of bytes allocated in the prologue for all local stack objects.

llvm/lib/CodeGen/AsmPrinter/CodeViewDebug.cpp

for (auto Annot : FI.Annotations) {
         StringRef Str = cast<MDString>(MD)->getString();
         assert(Str.data()[Str.size()] == '\0' && "non-nullterminated MDString");
         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 MCSymbol *BeginLabel = std::get<0>(HeapAllocSite);
+      const MCSymbol *EndLabel = std::get<1>(HeapAllocSite);
       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;
       }
     }
   }
 
   // Record beginning of function if we have a non-empty prologue.
   if (PrologEndLoc && !EmptyPrologue) {
     DebugLoc FnStartDL = PrologEndLoc.getFnDebugLoc();
     maybeRecordLocation(FnStartDL, MF);
   }
+
+  // Find heap alloc sites and emit labels around them.
+  for (const auto &MBB : *MF) {
+    for (const auto &MI : MBB) {
+      if (MI.getHeapAllocMarker()) {
+        requestLabelBeforeInsn(&MI);
+        requestLabelAfterInsn(&MI);
+      }
+    }
+  }
 }
 
 static bool shouldEmitUdt(const DIType *T) {
   if (!T)
     return false;
 
   // MSVC does not emit UDTs for typedefs that are scoped to classes.
   if (T->getTag() == dwarf::DW_TAG_typedef) {
void CodeViewDebug::endFunctionImpl(const MachineFunction *MF) {
   // Don't emit anything if we don't have any line tables.
   // Thunks are compiler-generated and probably won't have source correlation.
   if (!CurFn->HaveLineInfo && !GV.getSubprogram()->isThunk()) {
     FnDebugInfo.erase(&GV);
     CurFn = nullptr;
     return;
   }
 
+  // Find heap alloc sites and add to list.
+  for (const auto &MBB : *MF) {
+    for (const auto &MI : MBB) {
+      if (MDNode *MD = MI.getHeapAllocMarker()) {
+        CurFn->HeapAllocSites.push_back(std::make_tuple(getLabelBeforeInsn(&MI),
+                                                        getLabelAfterInsn(&MI),
+                                                        dyn_cast<DIType>(MD)));
+      }
+    }
+  }
+
   CurFn->Annotations = MF->getCodeViewAnnotations();
-  CurFn->HeapAllocSites = MF->getCodeViewHeapAllocSites();
 
   CurFn->End = Asm->getFunctionEnd();
 
   CurFn = nullptr;
 }
 
 // Usable locations are valid with non-zero line numbers. A line number of zero
 // corresponds to optimized code that doesn't have a distinct source location.

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, MDNode *HeapAllocMarker) {
   return MachineInstr::ExtraInfo::create(Allocator, MMOs, PreInstrSymbol,
-                                         PostInstrSymbol);
+                                         PostInstrSymbol, HeapAllocMarker);
 }
 
 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,
-                                               const MDNode *MD) {
-  MCSymbol *BeginLabel = Ctx.createTempSymbol("heapallocsite", true);
-  MCSymbol *EndLabel = Ctx.createTempSymbol("heapallocsite", true);
-  I->setPreInstrSymbol(*this, BeginLabel);
-  I->setPostInstrSymbol(*this, EndLabel);
-
-  const DIType *DI = dyn_cast<DIType>(MD);
-  CodeViewHeapAllocSites.push_back(std::make_tuple(BeginLabel, EndLabel, 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())
     return;
 

llvm/lib/CodeGen/MachineInstr.cpp

#endif
   // MachineOperand having a trivial destructor anyway, and adding a call here
   // wouldn't make it 'destructor-correct'.
 
   if (unsigned N = NumOperands - 1 - OpNo)
     moveOperands(Operands + OpNo, Operands + OpNo + 1, N, MRI);
   --NumOperands;
 }
 
-void MachineInstr::dropMemRefs(MachineFunction &MF) {
-  if (memoperands_empty())
-    return;
+void MachineInstr::setExtraInfo(MachineFunction &MF,
+                                ArrayRef<MachineMemOperand *> MMOs,
+                                MCSymbol *PreInstrSymbol,
+                                MCSymbol *PostInstrSymbol,
+                                MDNode *HeapAllocMarker) {
+  bool HasPreInstrSymbol = PreInstrSymbol != nullptr;
+  bool HasPostInstrSymbol = PostInstrSymbol != nullptr;
+  bool HasHeapAllocMarker = HeapAllocMarker != nullptr;
+  int NumPointers =
+      MMOs.size() + HasPreInstrSymbol + HasPostInstrSymbol + HasHeapAllocMarker;
 
-  // See if we can just drop all of our extra info.
-  if (!getPreInstrSymbol() && !getPostInstrSymbol()) {
+  // Drop all extra info if there is none.
+  if (NumPointers <= 0) {
     Info.clear();
     return;
   }
-  if (!getPostInstrSymbol()) {
-    Info.set<EIIK_PreInstrSymbol>(getPreInstrSymbol());
+
+  // If more than one pointer, then store out of line.
+  // FIXME: Maybe we should make the symbols in the extra info mutable?
+  else if (NumPointers > 1) {
+    Info.set<EIIK_OutOfLine>(MF.createMIExtraInfo(
+        MMOs, PreInstrSymbol, PostInstrSymbol, HeapAllocMarker));
     return;
   }
-  if (!getPreInstrSymbol()) {
-    Info.set<EIIK_PostInstrSymbol>(getPostInstrSymbol());
-    return;
-  }
+
+  // Otherwise store the single pointer inline.
+  if (HasPreInstrSymbol)
+    Info.set<EIIK_PreInstrSymbol>(PreInstrSymbol);
+  else if (HasPostInstrSymbol)
+    Info.set<EIIK_PostInstrSymbol>(PostInstrSymbol);
+  else if (HasHeapAllocMarker)
+    Info.set<EIIK_HeapAllocMarker>(HeapAllocMarker);
+  else
+    Info.set<EIIK_MMO>(MMOs[0]);
 
-  // Otherwise allocate a fresh extra info with just these symbols.
-  Info.set<EIIK_OutOfLine>(
-      MF.createMIExtraInfo({}, getPreInstrSymbol(), getPostInstrSymbol()));
+void MachineInstr::dropMemRefs(MachineFunction &MF) {
+  if (memoperands_empty())
+    return;
+
+  setExtraInfo(MF, {}, getPreInstrSymbol(), getPostInstrSymbol(),
+               getHeapAllocMarker());
 }
 
 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()) {
-    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()));
+  setExtraInfo(MF, MMOs, getPreInstrSymbol(), getPostInstrSymbol(),
+               getHeapAllocMarker());
 }
 
 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() &&
+      getHeapAllocMarker() == MI.getHeapAllocMarker()) {
     Info = MI.Info;
     return;
   }
 
   // Otherwise, fall back on a copy-based clone.
   setMemRefs(MF, MI.memoperands());
 }
 
for (const MachineInstr &MI : make_pointee_range(MIs.slice(1))) {
     // Otherwise accumulate these into our temporary buffer of the merged state.
     MergedMMOs.append(MI.memoperands_begin(), MI.memoperands_end());
   }
 
   setMemRefs(MF, MergedMMOs);
 }
 
 void MachineInstr::setPreInstrSymbol(MachineFunction &MF, MCSymbol *Symbol) {
-  MCSymbol *OldSymbol = getPreInstrSymbol();
-  if (OldSymbol == Symbol)
-    return;
-  if (OldSymbol && !Symbol) {
-    // We're removing a symbol rather than adding one. Try to clean up any
-    // extra info carried around.
-    if (Info.is<EIIK_PreInstrSymbol>()) {
-      Info.clear();
-      return;
-    }
-
-    if (memoperands_empty()) {
-      assert(getPostInstrSymbol() &&
-             "Should never have only a single symbol allocated out-of-line!");
-      Info.set<EIIK_PostInstrSymbol>(getPostInstrSymbol());
-      return;
-    }
+  // Do nothing if old and new symbols are the same.
+  if (Symbol == getPreInstrSymbol())
 
-    // Otherwise fallback on the generic update.
-  } else if (!Info || Info.is<EIIK_PreInstrSymbol>()) {
-    // If we don't have any other extra info, we can store this inline.
-    Info.set<EIIK_PreInstrSymbol>(Symbol);
+  // If there was only one symbol and we're removing it, just clear info.
+  if (!Symbol && Info.is<EIIK_PreInstrSymbol>()) {
+    Info.clear();
     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(), Symbol, getPostInstrSymbol()));
+  setExtraInfo(MF, memoperands(), Symbol, getPostInstrSymbol(),
+               getHeapAllocMarker());
 }
 
 void MachineInstr::setPostInstrSymbol(MachineFunction &MF, MCSymbol *Symbol) {
-  MCSymbol *OldSymbol = getPostInstrSymbol();
-  if (OldSymbol == Symbol)
+  // Do nothing if old and new symbols are the same.
+  if (Symbol == getPostInstrSymbol())
     return;
-  if (OldSymbol && !Symbol) {
-    // We're removing a symbol rather than adding one. Try to clean up any
-    // extra info carried around.
-    if (Info.is<EIIK_PostInstrSymbol>()) {
-      Info.clear();
-      return;
-    }
+
+  // If there was only one symbol and we're removing it, just clear info.
+  if (!Symbol && Info.is<EIIK_PostInstrSymbol>()) {
 
-    if (memoperands_empty()) {
-      assert(getPreInstrSymbol() &&
-             "Should never have only a single symbol allocated out-of-line!");
-      Info.set<EIIK_PreInstrSymbol>(getPreInstrSymbol());
-      return;
-    }
+  setExtraInfo(MF, memoperands(), getPreInstrSymbol(), Symbol,
+               getHeapAllocMarker());
 
-    // Otherwise fallback on the generic update.
-  } else if (!Info || Info.is<EIIK_PostInstrSymbol>()) {
-    // If we don't have any other extra info, we can store this inline.
-    Info.set<EIIK_PostInstrSymbol>(Symbol);
+void MachineInstr::setHeapAllocMarker(MachineFunction &MF, MDNode *Marker) {
+  // Do nothing if old and new symbols are the same.
+  if (Marker == getHeapAllocMarker())
+    return;
+
+  // If there was only one symbol and we're removing it, just clear info.
+  if (!Marker && Info.is<EIIK_HeapAllocMarker>()) {
+    Info.clear();
     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));
+  setExtraInfo(MF, memoperands(), getPreInstrSymbol(), getPostInstrSymbol(),
+               Marker);
 }
 
 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());
+  setHeapAllocMarker(MF, MI.getHeapAllocMarker());
 }
 
 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();
 }
 
 uint16_t MachineInstr::copyFlagsFromInstruction(const Instruction &I) {
   uint16_t MIFlags = 0;
   // Copy the wrapping flags.
   if (const OverflowingBinaryOperator *OB =
           dyn_cast<OverflowingBinaryOperator>(&I)) {
     if (OB->hasNoSignedWrap())
       MIFlags |= MachineInstr::MIFlag::NoSWrap;
     if (OB->hasNoUnsignedWrap())
       MIFlags |= MachineInstr::MIFlag::NoUWrap;
   }
 
   // Copy the exact flag.
   if (const PossiblyExactOperator *PE = dyn_cast<PossiblyExactOperator>(&I))
     if (PE->isExact())
       MIFlags |= MachineInstr::MIFlag::IsExact;
 
   // Copy the fast-math flags.
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 (MDNode *HeapAllocMarker = getHeapAllocMarker()) {
+    if (!FirstOp) {
+      FirstOp = false;
+      OS << ',';
+    }
+    OS << " heap-alloc-marker";
+  }
 
   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) {
   // Set all unused physreg defs as dead.
   assert(CLI.Call && "No call instruction specified.");
   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);
-  }
+  if (CLI.CS)
+    if (MDNode *MD = CLI.CS->getInstruction()->getMetadata("heapallocsite"))
+      CLI.Call->setHeapAllocMarker(*MF, MD);
 
   return true;
 }
 
 bool FastISel::lowerCall(const CallInst *CI) {
   ImmutableCallSite CS(CI);
 
   FunctionType *FuncTy = CS.getFunctionType();

llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp

for (SDNode *N = SU->getNode()->getGluedNode(); N; N = N->getGluedNode())
       GluedNodes.push_back(N);
     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 (MDNode *MD = DAG->getHeapAllocSite(N))
         if (NewInsn && NewInsn->isCall())
-          MF.addCodeViewHeapAllocSite(NewInsn, MD);
-      }
+          NewInsn->setHeapAllocMarker(MF, 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);
+        NewInsn->setHeapAllocMarker(MF, 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();
 ; };

 }
 
 ; Function Attrs: nounwind readnone speculatable willreturn
 declare void @llvm.dbg.value(metadata, metadata, metadata) #2
 
 
 ; Don't emit metadata for tail calls.
 ; CHECK-LABEL: call_tail:         # @call_tail
-; CHECK-NOT: .Lheapallocsite
 ; CHECK: jmp alloc_foo
 
 ; CHECK-LABEL: call_virtual:      # @call_virtual
-; CHECK: .Lheapallocsite0:
-; CHECK: callq *{{.*}}%rax{{.*}}
-; CHECK: .Lheapallocsite1:
+; CHECK:       callq *{{.*}}%rax{{.*}}
+; CHECK-NEXT:  [[LABEL1:.Ltmp[0-9]+]]:
 
 ; CHECK-LABEL: call_multiple:     # @call_multiple
-; FastISel emits instructions in a different order.
-; DAG:   .Lheapallocsite2:
-; FAST:  .Lheapallocsite4:
-; CHECK: callq alloc_foo
-; DAG:   .Lheapallocsite3:
-; FAST:  .Lheapallocsite5:
-; DAG:   .Lheapallocsite4:
-; FAST:  .Lheapallocsite2:
-; CHECK: callq alloc_foo
-; DAG:   .Lheapallocsite5:
-; FAST:  .Lheapallocsite3:
+; CHECK:       callq alloc_foo
+; CHECK-NEXT:  [[LABEL3:.Ltmp[0-9]+]]:
+; CHECK:       callq alloc_foo
+; CHECK-NEXT:  [[LABEL5:.Ltmp[0-9]+]]:
 
 ; 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:  .secrel32 [[LABEL0:.Ltmp[0-9]+]]
+; CHECK-NEXT:  .secidx [[LABEL0]]
+; CHECK-NEXT:  .short [[LABEL1]]-[[LABEL0]]
 ; CHECK-NEXT:  .long 3
 ; CHECK:       .short  4446                    # Record kind: S_HEAPALLOCSITE
-; CHECK-NEXT:  .secrel32 .Lheapallocsite2
-; CHECK-NEXT:  .secidx .Lheapallocsite2
-; CHECK-NEXT:  .short .Lheapallocsite3-.Lheapallocsite2
+; CHECK-NEXT:  .secrel32 [[LABEL2:.Ltmp[0-9]+]]
+; CHECK-NEXT:  .secidx [[LABEL2]]
+; CHECK-NEXT:  .short [[LABEL3]]-[[LABEL2]]
 ; CHECK-NEXT:  .long 4096
 ; CHECK:       .short  4446                    # Record kind: S_HEAPALLOCSITE
-; CHECK-NEXT:  .secrel32 .Lheapallocsite4
-; CHECK-NEXT:  .secidx .Lheapallocsite4
-; CHECK-NEXT:  .short .Lheapallocsite5-.Lheapallocsite4
+; CHECK-NEXT:  .secrel32 [[LABEL4:.Ltmp[0-9]+]]
+; CHECK-NEXT:  .secidx [[LABEL4]]
+; CHECK-NEXT:  .short [[LABEL5]]-[[LABEL4]]
 ; CHECK-NEXT:  .long 4096
 
 attributes #0 = { nounwind "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-features"="+cx8,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
 attributes #1 = { "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-features"="+cx8,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
 attributes #2 = { nounwind readnone speculatable willreturn }
 attributes #3 = { nounwind }
 
 !llvm.dbg.cu = !{!0}

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:
cond.end:                                         ; preds = %entry, %cond.true
   call void @llvm.dbg.value(metadata i8* %call, metadata !15, metadata !DIExpression()), !dbg !17
   tail call void @f2(), !dbg !23
   ret void, !dbg !24
 }
 
 ; CHECK-LABEL: taildupit: # @taildupit
 ; CHECK: testq
 ; CHECK: je
-; CHECK: .Lheapallocsite0:
 ; CHECK: callq alloc
-; CHECK: .Lheapallocsite1:
+; CHECK-NEXT: [[L1:.Ltmp[0-9]+]]
 ; CHECK: jmp f2 # TAILCALL
-; CHECK-NOT: Lheapallocsite
 ; CHECK: callq alloc
-; CHECK-NOT: Lheapallocsite
+; CHECK-NEXT: [[L3:.Ltmp[0-9]+]]
 ; CHECK: jmp f2 # TAILCALL
 
+; CHECK-LABEL: .short 4423                    # Record kind: S_GPROC32_ID
+; CHECK:       .short 4446                    # Record kind: S_HEAPALLOCSITE
+; CHECK-NEXT:  .secrel32 [[L0:.Ltmp[0-9]+]]
+; CHECK-NEXT:  .secidx [[L0]]
+; CHECK-NEXT:  .short [[L1]]-[[L0]]
+; CHECK-NEXT:  .long 3
+; CHECK:       .short 4446                    # Record kind: S_HEAPALLOCSITE
+; CHECK-NEXT:  .secrel32 [[L2:.Ltmp[0-9]+]]
+; CHECK-NEXT:  .secidx [[L2]]
+; CHECK-NEXT:  .short [[L3]]-[[L2]]
+; CHECK-NEXT:  .long 3
+
 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}

llvm/unittests/CodeGen/MachineInstrTest.cpp

 //===- MachineInstrTest.cpp -----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/ModuleSlotTracker.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 

 };
 
 class BogusTargetMachine : public LLVMTargetMachine {
 public:
   BogusTargetMachine()
       : LLVMTargetMachine(Target(), "", Triple(""), "", "", TargetOptions(),
                           Reloc::Static, CodeModel::Small, CodeGenOpt::Default),
         ST(*this) {}
+
   ~BogusTargetMachine() override {}
 
   const TargetSubtargetInfo *getSubtargetImpl(const Function &) const override {
     return &ST;
   }
 
 private:
   BogusSubtarget ST;
 };
 
+class BogusMCContext : public MCContext {
+public:
+  BogusMCContext(MCAsmInfo *mai)
+      : MCContext(mai, nullptr, nullptr, nullptr, nullptr, false) {}
+};
+
+std::unique_ptr<BogusMCContext> createMCContext() {
+  MCAsmInfo mai = MCAsmInfo();
+  return std::make_unique<BogusMCContext>(&mai);
+}
+
 std::unique_ptr<BogusTargetMachine> createTargetMachine() {
   return std::make_unique<BogusTargetMachine>();
 }
 
 std::unique_ptr<MachineFunction> createMachineFunction() {
   LLVMContext Ctx;
   Module M("Module", Ctx);
   auto Type = FunctionType::get(Type::getVoidTy(Ctx), false);
TEST(MachineInstrSpan, DistanceEnd) {
   MachineInstrSpan MIS(MII, MBB);
   ASSERT_TRUE(MIS.empty());
 
   auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
   MBB->insert(MII, MI);
   ASSERT_TRUE(std::distance(MIS.begin(), MII) == 1);
 }
 
+TEST(MachineInstrExtraInfo, AddExtraInfo) {
+  auto MF = createMachineFunction();
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  auto MC = createMCContext();
+  auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
+                                      MachineMemOperand::MOLoad, 8, 8);
+  SmallVector<MachineMemOperand *, 2> MMOs;
+  MMOs.push_back(MMO);
+  MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
+  MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
+  LLVMContext Ctx;
+  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+
+  ASSERT_TRUE(MI->memoperands_empty());
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setMemRefs(*MF, MMOs);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setPreInstrSymbol(*MF, Sym1);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setPostInstrSymbol(*MF, Sym2);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setHeapAllocMarker(*MF, MDN);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+}
+
+TEST(MachineInstrExtraInfo, ChangeExtraInfo) {
+  auto MF = createMachineFunction();
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  auto MC = createMCContext();
+  auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
+                                      MachineMemOperand::MOLoad, 8, 8);
+  SmallVector<MachineMemOperand *, 2> MMOs;
+  MMOs.push_back(MMO);
+  MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
+  MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
+  LLVMContext Ctx;
+  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+
+  MI->setMemRefs(*MF, MMOs);
+  MI->setPreInstrSymbol(*MF, Sym1);
+  MI->setPostInstrSymbol(*MF, Sym2);
+  MI->setHeapAllocMarker(*MF, MDN);
+
+  MMOs.push_back(MMO);
+
+  MI->setMemRefs(*MF, MMOs);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+
+  MI->setPostInstrSymbol(*MF, Sym1);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+}
+
+TEST(MachineInstrExtraInfo, RemoveExtraInfo) {
+  auto MF = createMachineFunction();
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  auto MC = createMCContext();
+  auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
+                                      MachineMemOperand::MOLoad, 8, 8);
+  SmallVector<MachineMemOperand *, 2> MMOs;
+  MMOs.push_back(MMO);
+  MMOs.push_back(MMO);
+  MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
+  MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
+  LLVMContext Ctx;
+  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+
+  MI->setMemRefs(*MF, MMOs);
+  MI->setPreInstrSymbol(*MF, Sym1);
+  MI->setPostInstrSymbol(*MF, Sym2);
+  MI->setHeapAllocMarker(*MF, MDN);
+
+  MI->setPostInstrSymbol(*MF, nullptr);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+
+  MI->setHeapAllocMarker(*MF, nullptr);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setPreInstrSymbol(*MF, nullptr);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setMemRefs(*MF, {});
+  ASSERT_TRUE(MI->memoperands_empty());
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+}
+
 static_assert(is_trivially_copyable<MCOperand>::value, "trivially copyable");
 
 } // end namespace