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

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

2 files changed, 26 insertions, 1 deletions

diff --git a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
index d49085b6fe5..a179c593380 100644
--- a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
+++ b/llvm/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) {
diff --git a/llvm/lib/CodeGen/MIRPrinter.cpp b/llvm/lib/CodeGen/MIRPrinter.cpp
index 84cf09f904c..f34cef7bcb4 100644
--- a/llvm/lib/CodeGen/MIRPrinter.cpp
+++ b/llvm/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");