[MachineInstr] Allow setting PCSections in ExtraInfo

[melver]

Summary

Provide MachineInstr::setPCSection(), to propagate relevant metadata
through the backend. Use ExtraInfo to store the metadata.

Depends on D130875

[dvyukov]

I looked at this and it looks good to me, but I don't feel like the right person to approve this.

[melver]

Rebase.

[vitalybuka]

In D130876#3692805, @dvyukov wrote:

I looked at this and it looks good to me, but I don't feel like the right person to approve this.

same
if you guess better candidates please add them as "blocking".

[melver]

Rebase.

llvm/include/llvm/CodeGen/MachineFunction.h

public:
   uint32_t *allocateRegMask();
 
   ArrayRef<int> allocateShuffleMask(ArrayRef<int> Mask);
 
   /// 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,
-                    MDNode *HeapAllocMarker = nullptr, uint32_t CFIType = 0);
+  MachineInstr::ExtraInfo *createMIExtraInfo(
+      ArrayRef<MachineMemOperand *> MMOs, MCSymbol *PreInstrSymbol = nullptr,
+      MCSymbol *PostInstrSymbol = nullptr, MDNode *HeapAllocMarker = nullptr,
+      MDNode *PCSections = nullptr, uint32_t CFIType = 0);
 
   /// 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.

llvm/include/llvm/CodeGen/MachineInstr.h

private:
   class ExtraInfo final : TrailingObjects<ExtraInfo, MachineMemOperand *,
                                           MCSymbol *, MDNode *, uint32_t> {
   public:
     static ExtraInfo *create(BumpPtrAllocator &Allocator,
                              ArrayRef<MachineMemOperand *> MMOs,
                              MCSymbol *PreInstrSymbol = nullptr,
                              MCSymbol *PostInstrSymbol = nullptr,
                              MDNode *HeapAllocMarker = nullptr,
+                             MDNode *PCSections = nullptr,
                              uint32_t CFIType = 0) {
       bool HasPreInstrSymbol = PreInstrSymbol != nullptr;
       bool HasPostInstrSymbol = PostInstrSymbol != nullptr;
       bool HasHeapAllocMarker = HeapAllocMarker != nullptr;
       bool HasCFIType = CFIType != 0;
+      bool HasPCSections = PCSections != nullptr;
       auto *Result = new (Allocator.Allocate(
           totalSizeToAlloc<MachineMemOperand *, MCSymbol *, MDNode *, uint32_t>(
               MMOs.size(), HasPreInstrSymbol + HasPostInstrSymbol,
-              HasHeapAllocMarker, HasCFIType),
+              HasHeapAllocMarker + HasPCSections, HasCFIType),
           alignof(ExtraInfo)))
           ExtraInfo(MMOs.size(), HasPreInstrSymbol, HasPostInstrSymbol,
-                    HasHeapAllocMarker, HasCFIType);
+                    HasHeapAllocMarker, HasPCSections, HasCFIType);
 
       // 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;
+      if (HasPCSections)
+        Result->getTrailingObjects<MDNode *>()[HasHeapAllocMarker] =
+            PCSections;
       if (HasCFIType)
         Result->getTrailingObjects<uint32_t>()[0] = CFIType;
 
       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;
     }
 
+    MDNode *getPCSections() const {
+      return HasPCSections
+                 ? getTrailingObjects<MDNode *>()[HasHeapAllocMarker]
+                 : nullptr;
+    }
+
     uint32_t getCFIType() const {
       return HasCFIType ? getTrailingObjects<uint32_t>()[0] : 0;
     }
 
   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;
+    const bool HasPCSections;
     const bool HasCFIType;
 
     // 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;
+      return HasHeapAllocMarker + HasPCSections;
     }
     size_t numTrailingObjects(OverloadToken<uint32_t>) const {
       return HasCFIType;
     }
 
     // Just a boring constructor to allow us to initialize the sizes. Always use
     // the `create` routine above.
     ExtraInfo(int NumMMOs, bool HasPreInstrSymbol, bool HasPostInstrSymbol,
-              bool HasHeapAllocMarker, bool HasCFIType)
+              bool HasHeapAllocMarker, bool HasPCSections, bool HasCFIType)
         : NumMMOs(NumMMOs), HasPreInstrSymbol(HasPreInstrSymbol),
           HasPostInstrSymbol(HasPostInstrSymbol),
-          HasHeapAllocMarker(HasHeapAllocMarker), HasCFIType(HasCFIType) {}
+          HasHeapAllocMarker(HasHeapAllocMarker), HasPCSections(HasPCSections),
+          HasCFIType(HasCFIType) {}
   };
 
   /// 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,
MDNode *getHeapAllocMarker() const {
     if (!Info)
       return nullptr;
     if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
       return EI->getHeapAllocMarker();
 
     return nullptr;
   }
 
+  /// Helper to extract PCSections metadata target sections.
+  MDNode *getPCSections() const {
+    if (!Info)
+      return nullptr;
+    if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
+      return EI->getPCSections();
+
+    return nullptr;
+  }
+
   /// Helper to extract a CFI type hash if one has been added.
   uint32_t getCFIType() const {
     if (!Info)
       return 0;
     if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
       return EI->getCFIType();
 
     return 0;

   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);
 
+  // Set metadata on instructions that say which sections to emit instruction
+  // addresses into.
+  void setPCSections(MachineFunction &MF, MDNode *MD);
+
   /// Set the CFI type for the instruction.
   void setCFIType(MachineFunction &MF, uint32_t Type);
 
   /// 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;
 
private:
   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, uint32_t CFIType);
+                    MDNode *HeapAllocMarker, MDNode *PCSections,
+                    uint32_t CFIType);
 };
 
 /// 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/MachineFunction.cpp

MachineFunction::getMachineMemOperand(const MachineMemOperand *MMO,
   return new (Allocator) MachineMemOperand(
       MMO->getPointerInfo(), Flags, MMO->getSize(), MMO->getBaseAlign(),
       MMO->getAAInfo(), MMO->getRanges(), MMO->getSyncScopeID(),
       MMO->getSuccessOrdering(), MMO->getFailureOrdering());
 }
 
 MachineInstr::ExtraInfo *MachineFunction::createMIExtraInfo(
     ArrayRef<MachineMemOperand *> MMOs, MCSymbol *PreInstrSymbol,
-    MCSymbol *PostInstrSymbol, MDNode *HeapAllocMarker, uint32_t CFIType) {
+    MCSymbol *PostInstrSymbol, MDNode *HeapAllocMarker, MDNode *PCSections,
+    uint32_t CFIType) {
   return MachineInstr::ExtraInfo::create(Allocator, MMOs, PreInstrSymbol,
                                          PostInstrSymbol, HeapAllocMarker,
-                                         CFIType);
+                                         PCSections, CFIType);
 }
 
 const char *MachineFunction::createExternalSymbolName(StringRef Name) {
   char *Dest = Allocator.Allocate<char>(Name.size() + 1);
   llvm::copy(Name, Dest);
   Dest[Name.size()] = 0;
   return Dest;
 }

llvm/lib/CodeGen/MachineInstr.cpp

if (unsigned N = NumOperands - 1 - OpNo)
     moveOperands(Operands + OpNo, Operands + OpNo + 1, N, MRI);
   --NumOperands;
 }
 
 void MachineInstr::setExtraInfo(MachineFunction &MF,
                                 ArrayRef<MachineMemOperand *> MMOs,
                                 MCSymbol *PreInstrSymbol,
                                 MCSymbol *PostInstrSymbol,
-                                MDNode *HeapAllocMarker, uint32_t CFIType) {
+                                MDNode *HeapAllocMarker, MDNode *PCSections,
+                                uint32_t CFIType) {
   bool HasPreInstrSymbol = PreInstrSymbol != nullptr;
   bool HasPostInstrSymbol = PostInstrSymbol != nullptr;
   bool HasHeapAllocMarker = HeapAllocMarker != nullptr;
+  bool HasPCSections = PCSections != nullptr;
   bool HasCFIType = CFIType != 0;
   int NumPointers = MMOs.size() + HasPreInstrSymbol + HasPostInstrSymbol +
-                    HasHeapAllocMarker + HasCFIType;
+                    HasHeapAllocMarker + HasPCSections + HasCFIType;
 
   // Drop all extra info if there is none.
   if (NumPointers <= 0) {
     Info.clear();
     return;
   }
 
   // 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 || HasCFIType) {
-    Info.set<EIIK_OutOfLine>(MF.createMIExtraInfo(
-        MMOs, PreInstrSymbol, PostInstrSymbol, HeapAllocMarker, CFIType));
+  else if (NumPointers > 1 || HasHeapAllocMarker || HasPCSections ||
+           HasCFIType) {
+    Info.set<EIIK_OutOfLine>(
+        MF.createMIExtraInfo(MMOs, PreInstrSymbol, PostInstrSymbol,
+                             HeapAllocMarker, PCSections, CFIType));
     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]);
 }
 
 void MachineInstr::dropMemRefs(MachineFunction &MF) {
   if (memoperands_empty())
     return;
 
   setExtraInfo(MF, {}, getPreInstrSymbol(), getPostInstrSymbol(),
-               getHeapAllocMarker(), getCFIType());
+               getHeapAllocMarker(), getPCSections(), getCFIType());
 }
 
 void MachineInstr::setMemRefs(MachineFunction &MF,
                               ArrayRef<MachineMemOperand *> MMOs) {
   if (MMOs.empty()) {
     dropMemRefs(MF);
     return;
   }
 
   setExtraInfo(MF, MMOs, getPreInstrSymbol(), getPostInstrSymbol(),
-               getHeapAllocMarker(), getCFIType());
+               getHeapAllocMarker(), getPCSections(), getCFIType());
 }
 
 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() &&
-      getHeapAllocMarker() == MI.getHeapAllocMarker()) {
+      getHeapAllocMarker() == MI.getHeapAllocMarker() &&
+      getPCSections() == MI.getPCSections()) {
     Info = MI.Info;
     return;
   }
 
   // Otherwise, fall back on a copy-based clone.
   setMemRefs(MF, MI.memoperands());
 }
 
void MachineInstr::setPreInstrSymbol(MachineFunction &MF, MCSymbol *Symbol) {
 
   // If there was only one symbol and we're removing it, just clear info.
   if (!Symbol && Info.is<EIIK_PreInstrSymbol>()) {
     Info.clear();
     return;
   }
 
   setExtraInfo(MF, memoperands(), Symbol, getPostInstrSymbol(),
-               getHeapAllocMarker(), getCFIType());
+               getHeapAllocMarker(), getPCSections(), getCFIType());
 }
 
 void MachineInstr::setPostInstrSymbol(MachineFunction &MF, MCSymbol *Symbol) {
   // Do nothing if old and new symbols are the same.
   if (Symbol == getPostInstrSymbol())
     return;
 
   // If there was only one symbol and we're removing it, just clear info.
   if (!Symbol && Info.is<EIIK_PostInstrSymbol>()) {
     Info.clear();
     return;
   }
 
   setExtraInfo(MF, memoperands(), getPreInstrSymbol(), Symbol,
-               getHeapAllocMarker(), getCFIType());
+               getHeapAllocMarker(), getPCSections(), getCFIType());
 }
 
 void MachineInstr::setHeapAllocMarker(MachineFunction &MF, MDNode *Marker) {
   // Do nothing if old and new symbols are the same.
   if (Marker == getHeapAllocMarker())
     return;
 
   setExtraInfo(MF, memoperands(), getPreInstrSymbol(), getPostInstrSymbol(),
-               Marker, getCFIType());
+               Marker, getPCSections(), getCFIType());
+}
+
+void MachineInstr::setPCSections(MachineFunction &MF, MDNode *PCSections) {
+  // Do nothing if old and new symbols are the same.
+  if (PCSections == getPCSections())
+    return;
+
+  setExtraInfo(MF, memoperands(), getPreInstrSymbol(), getPostInstrSymbol(),
+               getHeapAllocMarker(), PCSections, getCFIType());
 }
 
 void MachineInstr::setCFIType(MachineFunction &MF, uint32_t Type) {
   // Do nothing if old and new types are the same.
   if (Type == getCFIType())
     return;
 
   setExtraInfo(MF, memoperands(), getPreInstrSymbol(), getPostInstrSymbol(),
-               getHeapAllocMarker(), Type);
+               getHeapAllocMarker(), getPCSections(), Type);
 }
 
 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());
+  setPCSections(MF, MI.getPCSections());
 }
 
 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 (MDNode *HeapAllocMarker = getHeapAllocMarker()) {
     if (!FirstOp) {
       FirstOp = false;
       OS << ',';
     }
     OS << " heap-alloc-marker ";
     HeapAllocMarker->printAsOperand(OS, MST);
   }
+  if (MDNode *PCSections = getPCSections()) {
+    if (!FirstOp) {
+      FirstOp = false;
+      OS << ',';
+    }
+    OS << " pcsections ";
+    PCSections->printAsOperand(OS, MST);
+  }
   if (uint32_t CFIType = getCFIType()) {
     if (!FirstOp)
       OS << ',';
     OS << " cfi-type " << CFIType;
   }
 
   if (DebugInstrNum) {
     if (!FirstOp)

llvm/unittests/CodeGen/MachineInstrTest.cpp

TEST(MachineInstrExtraInfo, AddExtraInfo) {
   auto MAI = MCAsmInfo();
   auto MC = createMCContext(&MAI);
   auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
                                       MachineMemOperand::MOLoad, 8, Align(8));
   SmallVector<MachineMemOperand *, 2> MMOs;
   MMOs.push_back(MMO);
   MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
   MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
-  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+  MDNode *HAM = MDNode::getDistinct(Ctx, None);
+  MDNode *PCS = MDNode::getDistinct(Ctx, None);
 
   ASSERT_TRUE(MI->memoperands_empty());
   ASSERT_FALSE(MI->getPreInstrSymbol());
   ASSERT_FALSE(MI->getPostInstrSymbol());
   ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_FALSE(MI->getPCSections());
 
   MI->setMemRefs(*MF, MMOs);
   ASSERT_TRUE(MI->memoperands().size() == 1);
   ASSERT_FALSE(MI->getPreInstrSymbol());
   ASSERT_FALSE(MI->getPostInstrSymbol());
   ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_FALSE(MI->getPCSections());
 
   MI->setPreInstrSymbol(*MF, Sym1);
   ASSERT_TRUE(MI->memoperands().size() == 1);
   ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
   ASSERT_FALSE(MI->getPostInstrSymbol());
   ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_FALSE(MI->getPCSections());
 
   MI->setPostInstrSymbol(*MF, Sym2);
   ASSERT_TRUE(MI->memoperands().size() == 1);
   ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
   ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
   ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_FALSE(MI->getPCSections());
 
-  MI->setHeapAllocMarker(*MF, MDN);
+  MI->setHeapAllocMarker(*MF, HAM);
   ASSERT_TRUE(MI->memoperands().size() == 1);
   ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
   ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
-  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == HAM);
+  ASSERT_FALSE(MI->getPCSections());
+
+  MI->setPCSections(*MF, PCS);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == HAM);
+  ASSERT_TRUE(MI->getPCSections() == PCS);
 }
 
 TEST(MachineInstrExtraInfo, ChangeExtraInfo) {
   LLVMContext Ctx;
   Module Mod("Module", Ctx);
   auto MF = createMachineFunction(Ctx, Mod);
   MCInstrDesc MCID = {0, 0, 0, 0, 0, 0, 0, nullptr, nullptr, nullptr};
 
   auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
   auto MAI = MCAsmInfo();
   auto MC = createMCContext(&MAI);
   auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
                                       MachineMemOperand::MOLoad, 8, Align(8));
   SmallVector<MachineMemOperand *, 2> MMOs;
   MMOs.push_back(MMO);
   MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
   MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
-  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+  MDNode *HAM = MDNode::getDistinct(Ctx, None);
+  MDNode *PCS = MDNode::getDistinct(Ctx, None);
 
   MI->setMemRefs(*MF, MMOs);
   MI->setPreInstrSymbol(*MF, Sym1);
   MI->setPostInstrSymbol(*MF, Sym2);
-  MI->setHeapAllocMarker(*MF, MDN);
+  MI->setHeapAllocMarker(*MF, HAM);
+  MI->setPCSections(*MF, PCS);
 
   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);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == HAM);
+  ASSERT_TRUE(MI->getPCSections() == PCS);
 
   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);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == HAM);
+  ASSERT_TRUE(MI->getPCSections() == PCS);
 }
 
 TEST(MachineInstrExtraInfo, RemoveExtraInfo) {
   LLVMContext Ctx;
   Module Mod("Module", Ctx);
   auto MF = createMachineFunction(Ctx, Mod);
   MCInstrDesc MCID = {0, 0, 0, 0, 0, 0, 0, nullptr, nullptr, nullptr};
 
   auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
   auto MAI = MCAsmInfo();
   auto MC = createMCContext(&MAI);
   auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
                                       MachineMemOperand::MOLoad, 8, Align(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);
-  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+  MDNode *HAM = MDNode::getDistinct(Ctx, None);
+  MDNode *PCS = MDNode::getDistinct(Ctx, None);
 
   MI->setMemRefs(*MF, MMOs);
   MI->setPreInstrSymbol(*MF, Sym1);
   MI->setPostInstrSymbol(*MF, Sym2);
-  MI->setHeapAllocMarker(*MF, MDN);
+  MI->setHeapAllocMarker(*MF, HAM);
+  MI->setPCSections(*MF, PCS);
 
   MI->setPostInstrSymbol(*MF, nullptr);
   ASSERT_TRUE(MI->memoperands().size() == 2);
   ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
   ASSERT_FALSE(MI->getPostInstrSymbol());
-  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == HAM);
+  ASSERT_TRUE(MI->getPCSections() == PCS);
 
   MI->setHeapAllocMarker(*MF, nullptr);
   ASSERT_TRUE(MI->memoperands().size() == 2);
   ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
   ASSERT_FALSE(MI->getPostInstrSymbol());
   ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_TRUE(MI->getPCSections() == PCS);
+
+  MI->setPCSections(*MF, nullptr);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_FALSE(MI->getPCSections());
 
   MI->setPreInstrSymbol(*MF, nullptr);
   ASSERT_TRUE(MI->memoperands().size() == 2);
   ASSERT_FALSE(MI->getPreInstrSymbol());
   ASSERT_FALSE(MI->getPostInstrSymbol());
   ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_FALSE(MI->getPCSections());
 
   MI->setMemRefs(*MF, {});
   ASSERT_TRUE(MI->memoperands_empty());
   ASSERT_FALSE(MI->getPreInstrSymbol());
   ASSERT_FALSE(MI->getPostInstrSymbol());
   ASSERT_FALSE(MI->getHeapAllocMarker());
+  ASSERT_FALSE(MI->getPCSections());
 }
 
 TEST(MachineInstrDebugValue, AddDebugValueOperand) {
   LLVMContext Ctx;
   Module Mod("Module", Ctx);
   auto MF = createMachineFunction(Ctx, Mod);
 
   for (const unsigned short Opcode :