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

[akhuang]

Summary

Fixes some things from original commit at https://reviews.llvm.org/D69136. The main
change is that the heap alloc marker is always stored as ExtraInfo in the machine
instruction instead of in the PointerSumType because it cannot hold more than
4 pointer types.

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.

[rnk]

lgtm

llvm/unittests/CodeGen/MachineInstrTest.cpp

142

Seems like a reasonable fix.

144–145

Why is the "bogus" derived class needed anyway? It seems like you could make a regular MCContext and it would be fine.

[akhuang]

llvm/unittests/CodeGen/MachineInstrTest.cpp

144–145

Good point. It is just a regular MCContext.

[akhuang]

Remove bogus class wrapper

[sidneym]

I get a build error with this patch (b288f7d6bb8fdd21d27ba755302db194c181fdaf) It looks like what got committed is a little different from above. The committed patch partially changed updateCallSiteInfo with moveCallSiteInfo

/local/mnt/workspace/tmp/tipcheck/llvm-project/llvm/lib/CodeGen/MachineFunction.cpp:825:23: error: 
      out-of-line definition of 'moveCallSiteInfo' does not match any declaration in
      'llvm::MachineFunction'; did you mean 'updateCallSiteInfo'?
void MachineFunction::moveCallSiteInfo(const MachineInstr *Old,
                      ^~~~~~~~~~~~~~~~
                      updateCallSiteInfo

[akhuang]

In D69536#1744703, @sidneym wrote:

I get a build error with this patch (b288f7d6bb8fdd21d27ba755302db194c181fdaf) It looks like what got committed is a little different from above. The committed patch partially changed updateCallSiteInfo with moveCallSiteInfo

/local/mnt/workspace/tmp/tipcheck/llvm-project/llvm/lib/CodeGen/MachineFunction.cpp:825:23: error: 
      out-of-line definition of 'moveCallSiteInfo' does not match any declaration in
      'llvm::MachineFunction'; did you mean 'updateCallSiteInfo'?
void MachineFunction::moveCallSiteInfo(const MachineInstr *Old,
                      ^~~~~~~~~~~~~~~~
                      updateCallSiteInfo

Sorry, I guess I messed this up when merging.

llvm/include/llvm/CodeGen/MachineFunction.h

class MachineFunction {
   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;
 
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 *>()[0] = 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 *>()[0] : 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,
MCSymbol *getPostInstrSymbol() const {
     if (MCSymbol *S = Info.get<EIIK_PostInstrSymbol>())
       return S;
     if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
       return EI->getPostInstrSymbol();
 
     return nullptr;
   }
 
+  /// Helper to extract a heap alloc marker if one has been added.
+  MDNode *getHeapAllocMarker() const {
+    if (!Info)
+      return nullptr;
+    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 *MD);
+
   /// 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);
         }
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);
for (const auto &MBB : *MF) {
     }
   }
 
   // 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. Store heap alloc markers
+  // out of line because PointerSumType cannot hold more than 4 tag types with
+  // 32-bit pointers.
+  // FIXME: Maybe we should make the symbols in the extra info mutable?
+  else if (NumPointers > 1 || HasHeapAllocMarker) {
+    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
+    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)
+  // Do nothing if old and new symbols are the same.
+  if (Symbol == getPreInstrSymbol())
     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;
-    }
-
-    // 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;
-  }
 
-  // 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();
 }
 
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;
 };
 
+static MCAsmInfo AsmInfo = MCAsmInfo();
+
+std::unique_ptr<MCContext> createMCContext() {
+  return std::make_unique<MCContext>(
+      &AsmInfo, nullptr, nullptr, nullptr, nullptr, false);
+}
+
 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