[MIR] Normalize explicit successor probabilities

2 files changed, 37 insertions, 2 deletions

diff --git a/backend/lib/CodeGen/MIRParser/MIParser.cpp b/backend/lib/CodeGen/MIRParser/MIParser.cpp
index b21e713..62f1f14 100644
--- a/backend/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/backend/lib/CodeGen/MIRParser/MIParser.cpp
@@ -2500,6 +2500,9 @@ bool MIParser::Impl::parseBlockSuccessors(MachineBasicBlock &Block) {
   ExplicitSuccessorBlocks.insert(&Block);
   if (token().isNewlineOrEnd())
     return false;
+
+  std::vector<MachineBasicBlock *> Successors;
+  std::vector<BranchProbability> Probabilities;
   do {
     if (!at(MITokenKind::MachineBasicBlock))
       return fail("expected a machine basic block reference");
@@ -2517,10 +2520,17 @@ bool MIParser::Impl::parseBlockSuccessors(MachineBasicBlock &Block) {
       if (expect(MITokenKind::RightParen, "')' after successor weight"))
         return true;
     }
-    auto Added = Function.cfgUpdater().addSuccessor(Block, *Successor, Probability);
+    Successors.push_back(Successor);
+    Probabilities.push_back(Probability);
+  } while (consume(MITokenKind::Comma));
+
+  BranchProbability::normalizeProbabilities(Probabilities);
+  for (std::size_t Index = 0; Index < Successors.size(); ++Index) {
+    auto Added =
+        Function.cfgUpdater().addSuccessor(Block, *Successors[Index], Probabilities[Index]);
     if (!Added)
       return importError(token(), Added.error());
-  } while (consume(MITokenKind::Comma));
+  }
   return false;
 }
 
diff --git a/backend/lib/CodeGen/MachineIRTypes.cpp b/backend/lib/CodeGen/MachineIRTypes.cpp
index b5348b3..2810c92 100644
--- a/backend/lib/CodeGen/MachineIRTypes.cpp
+++ b/backend/lib/CodeGen/MachineIRTypes.cpp
@@ -2,7 +2,9 @@
 
 #include "rcc/CodeGen/MachineIRTypes.h"
 
+#include <algorithm>
 #include <bit>
+#include <cstdint>
 
 namespace rcc {
 
@@ -30,6 +32,29 @@ Expected<BranchProbability> BranchProbability::fromFraction(std::uint32_t Numera
   return BranchProbability(static_cast<std::uint32_t>(Rounded));
 }
 
+void BranchProbability::normalizeProbabilities(
+    std::span<BranchProbability> Probabilities) noexcept {
+  if (Probabilities.empty())
+    return;
+
+  std::uint64_t Sum = 0;
+  for (const BranchProbability Probability : Probabilities)
+    Sum += Probability.rawValue();
+
+  if (Sum == 0) {
+    const std::uint64_t Count = Probabilities.size();
+    const auto Equal =
+        static_cast<std::uint32_t>((static_cast<std::uint64_t>(MaxRawValue) + Count / 2U) / Count);
+    std::ranges::fill(Probabilities, BranchProbability(Equal));
+    return;
+  }
+
+  for (BranchProbability &Probability : Probabilities) {
+    const std::uint64_t Scaled = static_cast<std::uint64_t>(Probability.rawValue()) * MaxRawValue;
+    Probability = BranchProbability(static_cast<std::uint32_t>((Scaled + Sum / 2U) / Sum));
+  }
+}
+
 MachineInstrFlags::MachineInstrFlags(std::initializer_list<MachineInstrFlag> Flags) noexcept {
   for (const MachineInstrFlag Flag : Flags)
     set(Flag);