[MIR] Support printing and parsing pcsections

[melver]

Summary

Adds support for printing and parsing PC sections metadata in MIR.

llvm/lib/CodeGen/MIRParser/MILexer.h

enum TokenKind {
     kw_liveins,
     kw_successors,
     kw_floatpred,
     kw_intpred,
     kw_shufflemask,
     kw_pre_instr_symbol,
     kw_post_instr_symbol,
     kw_heap_alloc_marker,
+    kw_pcsections,
     kw_cfi_type,
     kw_bbsections,
     kw_unknown_size,
     kw_unknown_address,
     kw_ir_block_address_taken,
     kw_machine_block_address_taken,
 
     // Metadata types.

llvm/lib/CodeGen/MIRParser/MILexer.cpp

return StringSwitch<MIToken::TokenKind>(Identifier)
       .Case("liveins", MIToken::kw_liveins)
       .Case("successors", MIToken::kw_successors)
       .Case("floatpred", MIToken::kw_floatpred)
       .Case("intpred", MIToken::kw_intpred)
       .Case("shufflemask", MIToken::kw_shufflemask)
       .Case("pre-instr-symbol", MIToken::kw_pre_instr_symbol)
       .Case("post-instr-symbol", MIToken::kw_post_instr_symbol)
       .Case("heap-alloc-marker", MIToken::kw_heap_alloc_marker)
+      .Case("pcsections", MIToken::kw_pcsections)
       .Case("cfi-type", MIToken::kw_cfi_type)
       .Case("bbsections", MIToken::kw_bbsections)
       .Case("unknown-size", MIToken::kw_unknown_size)
       .Case("unknown-address", MIToken::kw_unknown_address)
       .Case("distinct", MIToken::kw_distinct)
       .Case("ir-block-address-taken", MIToken::kw_ir_block_address_taken)
       .Case("machine-block-address-taken", MIToken::kw_machine_block_address_taken)
       .Default(MIToken::Identifier);

llvm/lib/CodeGen/MIRParser/MIParser.cpp

public:
   bool parseMemoryOperandFlag(MachineMemOperand::Flags &Flags);
   bool parseMemoryPseudoSourceValue(const PseudoSourceValue *&PSV);
   bool parseMachinePointerInfo(MachinePointerInfo &Dest);
   bool parseOptionalScope(LLVMContext &Context, SyncScope::ID &SSID);
   bool parseOptionalAtomicOrdering(AtomicOrdering &Order);
   bool parseMachineMemoryOperand(MachineMemOperand *&Dest);
   bool parsePreOrPostInstrSymbol(MCSymbol *&Symbol);
   bool parseHeapAllocMarker(MDNode *&Node);
+  bool parsePCSections(MDNode *&Node);
 
   bool parseTargetImmMnemonic(const unsigned OpCode, const unsigned OpIdx,
                               MachineOperand &Dest, const MIRFormatter &MF);
 
 private:
   /// Convert the integer literal in the current token into an unsigned integer.
   ///
   /// Return true if an error occurred.
bool MIParser::parse(MachineInstr *&MI) {
   unsigned OpCode, Flags = 0;
   if (Token.isError() || parseInstruction(OpCode, Flags))
     return true;
 
   // Parse the remaining machine operands.
   while (!Token.isNewlineOrEOF() && Token.isNot(MIToken::kw_pre_instr_symbol) &&
          Token.isNot(MIToken::kw_post_instr_symbol) &&
          Token.isNot(MIToken::kw_heap_alloc_marker) &&
+         Token.isNot(MIToken::kw_pcsections) &&
          Token.isNot(MIToken::kw_cfi_type) &&
          Token.isNot(MIToken::kw_debug_location) &&
          Token.isNot(MIToken::kw_debug_instr_number) &&
          Token.isNot(MIToken::coloncolon) && Token.isNot(MIToken::lbrace)) {
     auto Loc = Token.location();
     Optional<unsigned> TiedDefIdx;
     if (parseMachineOperandAndTargetFlags(OpCode, Operands.size(), MO, TiedDefIdx))
       return true;
bool MIParser::parse(MachineInstr *&MI) {
   MCSymbol *PostInstrSymbol = nullptr;
   if (Token.is(MIToken::kw_post_instr_symbol))
     if (parsePreOrPostInstrSymbol(PostInstrSymbol))
       return true;
   MDNode *HeapAllocMarker = nullptr;
   if (Token.is(MIToken::kw_heap_alloc_marker))
     if (parseHeapAllocMarker(HeapAllocMarker))
       return true;
+  MDNode *PCSections = nullptr;
+  if (Token.is(MIToken::kw_pcsections))
+    if (parsePCSections(PCSections))
+      return true;
 
   unsigned CFIType = 0;
   if (Token.is(MIToken::kw_cfi_type)) {
     lex();
     if (Token.isNot(MIToken::IntegerLiteral))
       return error("expected an integer literal after 'cfi-type'");
     // getUnsigned is sufficient for 32-bit integers.
     if (getUnsigned(CFIType))
bool MIParser::parse(MachineInstr *&MI) {
   if (assignRegisterTies(*MI, Operands))
     return true;
   if (PreInstrSymbol)
     MI->setPreInstrSymbol(MF, PreInstrSymbol);
   if (PostInstrSymbol)
     MI->setPostInstrSymbol(MF, PostInstrSymbol);
   if (HeapAllocMarker)
     MI->setHeapAllocMarker(MF, HeapAllocMarker);
+  if (PCSections)
+    MI->setPCSections(MF, PCSections);
   if (CFIType)
     MI->setCFIType(MF, CFIType);
   if (!MemOperands.empty())
     MI->setMemRefs(MF, MemOperands);
   if (InstrNum)
     MI->setDebugInstrNum(InstrNum);
   return false;
 }
if (Token.isNewlineOrEOF() || Token.is(MIToken::coloncolon) ||
       Token.is(MIToken::lbrace))
     return false;
   if (Token.isNot(MIToken::comma))
     return error("expected ',' before the next machine operand");
   lex();
   return false;
 }
 
+bool MIParser::parsePCSections(MDNode *&Node) {
+  assert(Token.is(MIToken::kw_pcsections) &&
+         "Invalid token for a PC sections!");
+  lex();
+  parseMDNode(Node);
+  if (!Node)
+    return error("expected a MDNode after 'pcsections'");
+  if (Token.isNewlineOrEOF() || Token.is(MIToken::coloncolon) ||
+      Token.is(MIToken::lbrace))
+    return false;
+  if (Token.isNot(MIToken::comma))
+    return error("expected ',' before the next machine operand");
+  lex();
+  return false;
+}
+
 static void initSlots2BasicBlocks(
     const Function &F,
     DenseMap<unsigned, const BasicBlock *> &Slots2BasicBlocks) {
   ModuleSlotTracker MST(F.getParent(), /*ShouldInitializeAllMetadata=*/false);
   MST.incorporateFunction(F);
   for (const auto &BB : F) {
     if (BB.hasName())
       continue;

llvm/lib/CodeGen/MIRPrinter.cpp

void MIPrinter::print(const MachineInstr &MI) {
   }
   if (MDNode *HeapAllocMarker = MI.getHeapAllocMarker()) {
     if (NeedComma)
       OS << ',';
     OS << " heap-alloc-marker ";
     HeapAllocMarker->printAsOperand(OS, MST);
     NeedComma = true;
   }
+  if (MDNode *PCSections = MI.getPCSections()) {
+    if (NeedComma)
+      OS << ',';
+    OS << " pcsections ";
+    PCSections->printAsOperand(OS, MST);
+    NeedComma = true;
+  }
   if (uint32_t CFIType = MI.getCFIType()) {
     if (NeedComma)
       OS << ',';
     OS << " cfi-type " << CFIType;
     NeedComma = true;
   }
 
   if (auto Num = MI.peekDebugInstrNum()) {

llvm/test/CodeGen/MIR/X86/instr-pcsections.mir

+# RUN: llc -march=x86-64 -run-pass none -o - %s | FileCheck %s
+# This test ensures that the MIR parser parses pcsections metadata correctly.
+
+--- |
+
+  define i8 @test(i8* %a) {
+  entry:
+    %0 = load i8, i8* %a, align 1, !pcsections !0
+    ret i8 %0
+  }
+
+  !0 = !{!"foo"}
+
+...
+---
+name:            test
+alignment:       16
+tracksRegLiveness: true
+liveins:
+  - { reg: '$rdi' }
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo: {}
+body:             |
+  bb.0.entry:
+    liveins: $rdi
+
+    ; CHECK-LABEL: name: test
+    ; CHECK: MOV{{.*}} pcsections !0
+    renamable $al = MOV8rm killed renamable $rdi, 1, $noreg, 0, $noreg, pcsections !0
+    RET64 implicit killed $al
+
+...