MIR Serialization: Serialize the machine basic block's successor weights.

Reviewers: Duncan P. N. Exon Smith
llvm-svn: 243659

include/llvm/CodeGen/MIRYamlMapping.h

@@ -73,6 +73,37 @@ template <> struct ScalarTraits<FlowStringValue> {
   static bool mustQuote(StringRef Scalar) { return needsQuotes(Scalar); }
 };
 
+/// A wrapper around unsigned which contains a source range that's being set
+/// during parsing.
+struct UnsignedValue {
+  unsigned Value;
+  SMRange SourceRange;
+
+  UnsignedValue() : Value(0) {}
+  UnsignedValue(unsigned Value) : Value(Value) {}
+
+  bool operator==(const UnsignedValue &Other) const {
+    return Value == Other.Value;
+  }
+};
+
+template <> struct ScalarTraits<UnsignedValue> {
+  static void output(const UnsignedValue &Value, void *Ctx, raw_ostream &OS) {
+    return ScalarTraits<unsigned>::output(Value.Value, Ctx, OS);
+  }
+
+  static StringRef input(StringRef Scalar, void *Ctx, UnsignedValue &Value) {
+    if (const auto *Node =
+            reinterpret_cast<yaml::Input *>(Ctx)->getCurrentNode())
+      Value.SourceRange = Node->getSourceRange();
+    return ScalarTraits<unsigned>::input(Scalar, Ctx, Value.Value);
+  }
+
+  static bool mustQuote(StringRef Scalar) {
+    return ScalarTraits<unsigned>::mustQuote(Scalar);
+  }
+};
+
 template <> struct ScalarEnumerationTraits<MachineJumpTableInfo::JTEntryKind> {
   static void enumeration(yaml::IO &IO,
                           MachineJumpTableInfo::JTEntryKind &EntryKind) {
@@ -94,6 +125,7 @@ template <> struct ScalarEnumerationTraits<MachineJumpTableInfo::JTEntryKind> {
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::StringValue)
 LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(llvm::yaml::FlowStringValue)
+LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(llvm::yaml::UnsignedValue)
 
 namespace llvm {
 namespace yaml {
@@ -139,8 +171,8 @@ struct MachineBasicBlock {
   unsigned Alignment = 0;
   bool IsLandingPad = false;
   bool AddressTaken = false;
-  // TODO: Serialize the successor weights.
   std::vector<FlowStringValue> Successors;
+  std::vector<UnsignedValue> SuccessorWeights;
   std::vector<FlowStringValue> LiveIns;
   std::vector<StringValue> Instructions;
 };
@@ -156,6 +188,7 @@ template <> struct MappingTraits<MachineBasicBlock> {
     YamlIO.mapOptional("isLandingPad", MBB.IsLandingPad);
     YamlIO.mapOptional("addressTaken", MBB.AddressTaken);
     YamlIO.mapOptional("successors", MBB.Successors);
+    YamlIO.mapOptional("weights", MBB.SuccessorWeights);
     YamlIO.mapOptional("liveins", MBB.LiveIns);
     YamlIO.mapOptional("instructions", MBB.Instructions);
   }

include/llvm/CodeGen/MachineBasicBlock.h

@@ -25,6 +25,7 @@ class Pass;
 class BasicBlock;
 class MachineFunction;
 class MCSymbol;
+class MIRPrinter;
 class SlotIndexes;
 class StringRef;
 class raw_ostream;
@@ -424,6 +425,9 @@ public:
   /// which refer to fromMBB to refer to this.
   void transferSuccessorsAndUpdatePHIs(MachineBasicBlock *fromMBB);
 
+  /// Return true if any of the successors have weights attached to them.
+  bool hasSuccessorWeights() const { return !Weights.empty(); }
+
   /// isPredecessor - Return true if the specified MBB is a predecessor of this
   /// block.
   bool isPredecessor(const MachineBasicBlock *MBB) const;
@@ -685,6 +689,7 @@ private:
   const_weight_iterator getWeightIterator(const_succ_iterator I) const;
 
   friend class MachineBranchProbabilityInfo;
+  friend class MIRPrinter;
 
   /// getSuccWeight - Return weight of the edge from this block to MBB. This
   /// method should NOT be called directly, but by using getEdgeWeight method

lib/CodeGen/MIRParser/MIRParser.cpp

@@ -356,12 +356,24 @@ bool MIRParserImpl::initializeMachineBasicBlock(
   MBB.setIsLandingPad(YamlMBB.IsLandingPad);
   SMDiagnostic Error;
   // Parse the successors.
+  const auto &Weights = YamlMBB.SuccessorWeights;
+  bool HasWeights = !Weights.empty();
+  if (HasWeights && Weights.size() != YamlMBB.Successors.size()) {
+    bool IsFew = Weights.size() < YamlMBB.Successors.size();
+    return error(IsFew ? Weights.back().SourceRange.End
+                       : Weights[YamlMBB.Successors.size()].SourceRange.Start,
+                 Twine("too ") + (IsFew ? "few" : "many") +
+                     " successor weights, expected " +
+                     Twine(YamlMBB.Successors.size()) + ", have " +
+                     Twine(Weights.size()));
+  }
+  size_t SuccessorIndex = 0;
   for (const auto &MBBSource : YamlMBB.Successors) {
     MachineBasicBlock *SuccMBB = nullptr;
     if (parseMBBReference(SuccMBB, MBBSource, MF, PFS))
       return true;
     // TODO: Report an error when adding the same successor more than once.
-    MBB.addSuccessor(SuccMBB);
+    MBB.addSuccessor(SuccMBB, HasWeights ? Weights[SuccessorIndex++].Value : 0);
   }
   // Parse the liveins.
   for (const auto &LiveInSource : YamlMBB.LiveIns) {

lib/CodeGen/MIRPrinter.cpp

@@ -56,6 +56,10 @@ struct FrameIndexOperand {
   }
 };
 
+} // end anonymous namespace
+
+namespace llvm {
+
 /// This class prints out the machine functions using the MIR serialization
 /// format.
 class MIRPrinter {
@@ -88,6 +92,10 @@ private:
   void initRegisterMaskIds(const MachineFunction &MF);
 };
 
+} // end namespace llvm
+
+namespace {
+
 /// This class prints out the machine instructions using the MIR serialization
 /// format.
 class MIPrinter {
@@ -363,6 +371,11 @@ void MIRPrinter::convert(ModuleSlotTracker &MST,
         .printMBBReference(*SuccMBB);
     YamlMBB.Successors.push_back(StrOS.str());
   }
+  if (MBB.hasSuccessorWeights()) {
+    for (auto I = MBB.succ_begin(), E = MBB.succ_end(); I != E; ++I)
+      YamlMBB.SuccessorWeights.push_back(
+          yaml::UnsignedValue(MBB.getSuccWeight(I)));
+  }
   // Print the live in registers.
   const auto *TRI = MBB.getParent()->getSubtarget().getRegisterInfo();
   assert(TRI && "Expected target register info");

test/CodeGen/MIR/X86/successor-basic-blocks-few-weights.mir

@@ -0,0 +1,41 @@
+# RUN: not llc -march=x86-64 -start-after branch-folder -stop-after branch-folder -o /dev/null %s 2>&1 | FileCheck %s
+
+--- |
+
+  define i32 @foo(i32 %a) {
+  entry:
+    %0 = icmp sle i32 %a, 10
+    br i1 %0, label %less, label %exit
+
+  less:
+    ret i32 0
+
+  exit:
+    ret i32 %a
+  }
+
+...
+---
+name:            foo
+body:
+  - id:          0
+    name:        entry
+    successors:  [ '%bb.1.less', '%bb.2.exit' ]
+# CHECK: [[@LINE+1]]:23: too few successor weights, expected 2, have 1
+    weights:     [ 16 ]
+    liveins:     [ '%edi' ]
+    instructions:
+      - 'CMP32ri8 %edi, 10, implicit-def %eflags'
+      - 'JG_1 %bb.2.exit, implicit killed %eflags'
+  - id:          1
+    name:        less
+    instructions:
+      - '%eax = MOV32r0 implicit-def dead %eflags'
+      - 'RETQ killed %eax'
+  - id:          2
+    name:        exit
+    liveins:     [ '%edi' ]
+    instructions:
+      - '%eax = COPY killed %edi'
+      - 'RETQ killed %eax'
+...

test/CodeGen/MIR/X86/successor-basic-blocks-many-weights.mir

@@ -0,0 +1,41 @@
+# RUN: not llc -march=x86-64 -start-after branch-folder -stop-after branch-folder -o /dev/null %s 2>&1 | FileCheck %s
+
+--- |
+
+  define i32 @foo(i32 %a) {
+  entry:
+    %0 = icmp sle i32 %a, 10
+    br i1 %0, label %less, label %exit
+
+  less:
+    ret i32 0
+
+  exit:
+    ret i32 %a
+  }
+
+...
+---
+name:            foo
+body:
+  - id:          0
+    name:        entry
+    successors:  [ '%bb.1.less', '%bb.2.exit' ]
+# CHECK: [[@LINE+1]]:28: too many successor weights, expected 2, have 3
+    weights:     [ 16, 16, 16 ]
+    liveins:     [ '%edi' ]
+    instructions:
+      - 'CMP32ri8 %edi, 10, implicit-def %eflags'
+      - 'JG_1 %bb.2.exit, implicit killed %eflags'
+  - id:          1
+    name:        less
+    instructions:
+      - '%eax = MOV32r0 implicit-def dead %eflags'
+      - 'RETQ killed %eax'
+  - id:          2
+    name:        exit
+    liveins:     [ '%edi' ]
+    instructions:
+      - '%eax = COPY killed %edi'
+      - 'RETQ killed %eax'
+...

test/CodeGen/MIR/X86/successor-basic-blocks-weights.mir

@@ -0,0 +1,46 @@
+# RUN: llc -march=x86-64 -start-after branch-folder -stop-after branch-folder -o /dev/null %s | FileCheck %s
+# This test ensures that the MIR parser parses basic block successors and
+# weights correctly.
+
+--- |
+
+  define i32 @foo(i32 %a) {
+  entry:
+    %0 = icmp sle i32 %a, 10
+    br i1 %0, label %less, label %exit
+
+  less:
+    ret i32 0
+
+  exit:
+    ret i32 %a
+  }
+
+...
+---
+name:            foo
+body:
+  # CHECK:      name: entry
+  # CHECK:      successors: [ '%bb.1.less', '%bb.2.exit' ]
+  # CHECK-NEXT: weights: [ 16, 32 ]
+  # CHECK:      name: less
+  - id:          0
+    name:        entry
+    successors:  [ '%bb.1.less', '%bb.2.exit' ]
+    weights:     [ 16, 32 ]
+    liveins:     [ '%edi' ]
+    instructions:
+      - 'CMP32ri8 %edi, 10, implicit-def %eflags'
+      - 'JG_1 %bb.2.exit, implicit killed %eflags'
+  - id:          1
+    name:        less
+    instructions:
+      - '%eax = MOV32r0 implicit-def dead %eflags'
+      - 'RETQ killed %eax'
+  - id:          2
+    name:        exit
+    liveins:     [ '%edi' ]
+    instructions:
+      - '%eax = COPY killed %edi'
+      - 'RETQ killed %eax'
+...