[MIR] Add MIR parsing for heap alloc site instruction markers

[akhuang]

Summary

This patch adds MIR parsing and printing for heap alloc markers, which were
added in D69136. They are printed as an operand similar to pre-/post-instr
symbols, with a heap-alloc-marker token and a metadata node.

[akhuang]

formatting

[rnk]

lgtm, thanks!

llvm/lib/CodeGen/MIRParser/MILexer.h

enum TokenKind {
     kw_landing_pad,
     kw_liveins,
     kw_successors,
     kw_floatpred,
     kw_intpred,
     kw_shufflemask,
     kw_pre_instr_symbol,
     kw_post_instr_symbol,
+    kw_heap_alloc_marker,
     kw_unknown_size,
 
     // Named metadata keywords
     md_tbaa,
     md_alias_scope,
     md_noalias,
     md_range,
     md_diexpr,

llvm/lib/CodeGen/MIRParser/MILexer.cpp

return StringSwitch<MIToken::TokenKind>(Identifier)
       .Case("landing-pad", MIToken::kw_landing_pad)
       .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("unknown-size", MIToken::kw_unknown_size)
       .Default(MIToken::Identifier);
 }
 
 static Cursor maybeLexIdentifier(Cursor C, MIToken &Token) {
   if (!isalpha(C.peek()) && C.peek() != '_')
     return None;
   auto Range = C;
return StringSwitch<MIToken::TokenKind>(Identifier)
       .Case("!alias.scope", MIToken::md_alias_scope)
       .Case("!noalias", MIToken::md_noalias)
       .Case("!range", MIToken::md_range)
       .Case("!DIExpression", MIToken::md_diexpr)
       .Case("!DILocation", MIToken::md_dilocation)
       .Default(MIToken::Error);
 }
 
-static Cursor maybeLexExlaim(Cursor C, MIToken &Token,
+static Cursor maybeLexExclaim(Cursor C, MIToken &Token,
                               ErrorCallbackType ErrorCallback) {
   if (C.peek() != '!')
     return None;
   auto Range = C;
   C.advance(1);
   if (isdigit(C.peek()) || !isIdentifierChar(C.peek())) {
     Token.reset(MIToken::exclaim, Range.upto(C));
     return C;
   }
StringRef llvm::lexMIToken(StringRef Source, MIToken &Token,
   if (Cursor R = maybeLexExternalSymbol(C, Token, ErrorCallback))
     return R.remaining();
   if (Cursor R = maybeLexMCSymbol(C, Token, ErrorCallback))
     return R.remaining();
   if (Cursor R = maybeLexHexadecimalLiteral(C, Token))
     return R.remaining();
   if (Cursor R = maybeLexNumericalLiteral(C, Token))
     return R.remaining();
-  if (Cursor R = maybeLexExlaim(C, Token, ErrorCallback))
+  if (Cursor R = maybeLexExclaim(C, Token, ErrorCallback))
     return R.remaining();
   if (Cursor R = maybeLexSymbol(C, Token))
     return R.remaining();
   if (Cursor R = maybeLexNewline(C, Token))
     return R.remaining();
   if (Cursor R = maybeLexEscapedIRValue(C, Token, ErrorCallback))
     return R.remaining();
   if (Cursor R = maybeLexStringConstant(C, Token, ErrorCallback))
     return R.remaining();
 
   Token.reset(MIToken::Error, C.remaining());
   ErrorCallback(C.location(),
                 Twine("unexpected character '") + Twine(C.peek()) + "'");
   return C.remaining();
 }

llvm/lib/CodeGen/MIRParser/MIParser.cpp

public:
   bool parseIRValue(const Value *&V);
   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);
 
 private:
   /// Convert the integer literal in the current token into an unsigned integer.
   ///
   /// Return true if an error occurred.
   bool getUnsigned(unsigned &Result);
 
   /// Convert the integer literal in the current token into an uint64.
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_debug_location) &&
          Token.isNot(MIToken::coloncolon) && Token.isNot(MIToken::lbrace)) {
     auto Loc = Token.location();
     Optional<unsigned> TiedDefIdx;
     if (parseMachineOperandAndTargetFlags(MO, TiedDefIdx))
       return true;
     if (OpCode == TargetOpcode::DBG_VALUE && MO.isReg())
       MO.setIsDebug();
bool MIParser::parse(MachineInstr *&MI) {
   MCSymbol *PreInstrSymbol = nullptr;
   if (Token.is(MIToken::kw_pre_instr_symbol))
     if (parsePreOrPostInstrSymbol(PreInstrSymbol))
       return true;
   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;
 
   DebugLoc DebugLocation;
   if (Token.is(MIToken::kw_debug_location)) {
     lex();
     MDNode *Node = nullptr;
     if (Token.is(MIToken::exclaim)) {
       if (parseMDNode(Node))
         return true;
bool MIParser::parse(MachineInstr *&MI) {
   for (const auto &Operand : Operands)
     MI->addOperand(MF, Operand.Operand);
   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 (!MemOperands.empty())
     MI->setMemRefs(MF, MemOperands);
   return false;
 }
 
 bool MIParser::parseStandaloneMBB(MachineBasicBlock *&MBB) {
   lex();
   if (Token.isNot(MIToken::MachineBasicBlock))
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::parseHeapAllocMarker(MDNode *&Node) {
+  assert(Token.is(MIToken::kw_heap_alloc_marker) &&
+         "Invalid token for a heap alloc marker!");
+  lex();
+  parseMDNode(Node);
+  if (!Node)
+    return error("expected a MDNode after 'heap-alloc-marker'");
+  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 (auto &BB : F) {
     if (BB.hasName())
       continue;

llvm/lib/CodeGen/MIRPrinter.cpp

void MIPrinter::print(const MachineInstr &MI) {
   }
   if (MCSymbol *PostInstrSymbol = MI.getPostInstrSymbol()) {
     if (NeedComma)
       OS << ',';
     OS << " post-instr-symbol ";
     MachineOperand::printSymbol(OS, *PostInstrSymbol);
     NeedComma = true;
   }
+  if (MDNode *HeapAllocMarker = MI.getHeapAllocMarker()) {
+    if (NeedComma)
+      OS << ',';
+    OS << " heap-alloc-marker ";
+    HeapAllocMarker->printAsOperand(OS, MST);
+    NeedComma = true;
+  }
 
   if (const DebugLoc &DL = MI.getDebugLoc()) {
     if (NeedComma)
       OS << ',';
     OS << " debug-location ";
     DL->printAsOperand(OS, MST);
   }
 

llvm/lib/CodeGen/MachineInstr.cpp

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 (MDNode *HeapAllocMarker = getHeapAllocMarker()) {
     if (!FirstOp) {
       FirstOp = false;
       OS << ',';
     }
-    OS << " heap-alloc-marker";
+    OS << " heap-alloc-marker ";
+    HeapAllocMarker->printAsOperand(OS, MST);
   }
 
   if (!SkipDebugLoc) {
     if (const DebugLoc &DL = getDebugLoc()) {
       if (!FirstOp)
         OS << ',';
       OS << " debug-location ";
       DL->printAsOperand(OS, MST);

llvm/test/CodeGen/MIR/X86/instr-heap-alloc-operands.mir

+# RUN: llc -march=x86-64 -run-pass none -o - %s | FileCheck %s
+# This test ensures that the MIR parser parses heap alloc markers correctly.
+
+--- |
+  declare i8* @f(i32) nounwind
+
+  define i32 @test(i32 %x) nounwind {
+  entry:
+    call i8* @f(i32 %x), !heapallocsite !2
+    ret i32 undef
+  }
+
+  !llvm.dbg.cu = !{!0}
+  !llvm.module.flags = !{!3, !4}
+
+  !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, nameTableKind: None)
+  !1 = !DIFile(filename: "test.ll", directory: "")
+  !2 = !{}
+  !3 = !{i32 2, !"Debug Info Version", i32 3}
+  !4 = !{i32 1, !"wchar_size", i32 2}
+
+...
+---
+name: test
+# CHECK-LABEL: name: test
+tracksRegLiveness: true
+frameInfo:
+  hasCalls: true
+body: |
+  bb.0.entry:
+    liveins: $ecx
+
+    %0:gr32 = COPY $ecx
+    ADJCALLSTACKDOWN64 32, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    $ecx = COPY %0
+    CALL64pcrel32 @f, csr_win64, implicit $rsp, implicit $ssp, implicit $ecx, implicit-def $rsp, implicit-def $ssp, implicit-def dead $rax, heap-alloc-marker !2
+  ; CHECK: CALL64pcrel32 @f, {{.*}} heap-alloc-marker !2
+    ADJCALLSTACKUP64 32, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    $eax = MOV32r0 implicit-def dead $eflags
+    RET 0, killed $eax
+
+...

llvm/test/CodeGen/X86/heap-alloc-markers.mir

+# RUN: llc -mtriple=x86_64-unknown-windows-msvc -start-before=x86-flags-copy-lowering -o - %s | FileCheck %s
+#
+# Test the emission of heap alloc site instruction labels.
+
+--- |
+  declare i8* @alloc(i32) nounwind
+
+  define i32 @test(i32 %x) nounwind !dbg !6 {
+  entry:
+    call i8* @alloc(i32 %x), !dbg !11, !heapallocsite !2
+    ret i32 0, !dbg !12
+  }
+
+  !llvm.dbg.cu = !{!0}
+  !llvm.module.flags = !{!3, !4, !5}
+
+  !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, nameTableKind: None)
+  !1 = !DIFile(filename: "<stdin>", directory: "")
+  !2 = !{}
+  !3 = !{i32 2, !"CodeView", i32 1}
+  !4 = !{i32 2, !"Debug Info Version", i32 3}
+  !5 = !{i32 1, !"wchar_size", i32 2}
+  !6 = distinct !DISubprogram(name: "test", scope: !7, file: !7, line: 2, type: !8, scopeLine: 2, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !2)
+  !7 = !DIFile(filename: "test.c", directory: "")
+  !8 = !DISubroutineType(types: !9)
+  !9 = !{!10, !10}
+  !10 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+  !11 = !DILocation(line: 3, scope: !6)
+  !12 = !DILocation(line: 4, scope: !6)
+
+...
+---
+name: test
+# CHECK-LABEL: {{^}}test:
+tracksRegLiveness: true
+frameInfo:
+  hasCalls: true
+body: |
+  bb.0.entry:
+    liveins: $ecx
+
+    ADJCALLSTACKDOWN64 32, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp, debug-location !11
+    CALL64pcrel32 @alloc, csr_win64, implicit $rsp, implicit $ssp, implicit $ecx, implicit-def $rsp, implicit-def $ssp, implicit-def dead $rax, heap-alloc-marker !2, debug-location !11
+  ; CHECK:      callq alloc
+  ; CHECK-NEXT: .Ltmp{{.*}}:
+    ADJCALLSTACKUP64 32, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp, debug-location !11
+    $eax = MOV32r0 implicit-def dead $eflags, debug-location !12
+    RET 0, killed $eax, debug-location !12
+
+...