From 9430569e0817c61d6fdaa37cf5cc35445f2f6289 Mon Sep 17 00:00:00 2001
From: Richard-Qin-X <richard.qin.001@gmail.com>
Date: Fri, 14 Aug 2026 21:39:49 +0800
Subject: [MIR] Normalize explicit successor probabilities

---
 backend/include/rcc/CodeGen/MachineIRTypes.h |  3 +++
 backend/lib/CodeGen/MIRParser/MIParser.cpp   | 14 +++++++++++--
 backend/lib/CodeGen/MachineIRTypes.cpp       | 25 ++++++++++++++++++++++
 backend/unittests/MIR/MIParserTest.cpp       | 31 ++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 2 deletions(-)

(limited to 'backend')

diff --git a/backend/include/rcc/CodeGen/MachineIRTypes.h b/backend/include/rcc/CodeGen/MachineIRTypes.h
index 922adef..66369e0 100644
--- a/backend/include/rcc/CodeGen/MachineIRTypes.h
+++ b/backend/include/rcc/CodeGen/MachineIRTypes.h
@@ -11,6 +11,7 @@
 #include <cstdint>
 #include <functional>
 #include <initializer_list>
+#include <span>
 #include <string>
 #include <type_traits>
 
@@ -207,6 +208,8 @@ public:
   [[nodiscard]] static Expected<BranchProbability> fromFraction(std::uint32_t Numerator,
                                                                 std::uint32_t Denominator);
 
+  static void normalizeProbabilities(std::span<BranchProbability> Probabilities) noexcept;
+
   [[nodiscard]] constexpr std::uint32_t rawValue() const noexcept {
     return RawValue;
   }
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);
diff --git a/backend/unittests/MIR/MIParserTest.cpp b/backend/unittests/MIR/MIParserTest.cpp
index 4ed39ae..61292ac 100644
--- a/backend/unittests/MIR/MIParserTest.cpp
+++ b/backend/unittests/MIR/MIParserTest.cpp
@@ -553,6 +553,36 @@ void testLLVMCompareExchangeMemorySyntax(TestContext &Context) {
                  "canonical LLVM MMOs must omit an absent zero-offset pointer base");
 }
 
+void testUnweightedSuccessorNormalization(TestContext &Context) {
+  auto Target = createTarget();
+  auto Function = createFunction(Context, 9);
+  Context.expect(Target.has_value(), "successor probability target description must be valid");
+  if (!Target || !Function)
+    return;
+
+  const std::string Body = R"(bb.0:
+  successors: %bb.1, %bb.2
+  TEST_NOP
+
+bb.1:
+  TEST_NOP
+
+bb.2:
+  TEST_NOP
+)";
+  rcc::mir::MIParser Parser(*Function, *Target, "successors.mir", Body);
+  const auto Parsed = Parser.parse();
+  Context.expect(Parsed.has_value(), "unweighted MIR successors must parse");
+  if (!Parsed)
+    return;
+
+  const rcc::MachineBasicBlock *Entry = Function->firstBlock();
+  Context.expect(Entry != nullptr && Entry->successors().size() == 2 &&
+                     Entry->successors()[0].Probability.rawValue() == 0x40000000U &&
+                     Entry->successors()[1].Probability.rawValue() == 0x40000000U,
+                 "unweighted successors must receive equal normalized probabilities");
+}
+
 void expectParseError(TestContext &Context, const rcc::TargetDescription &Target, std::string Body,
                       std::string_view Message) {
   auto Function = createFunction(Context, 10);
@@ -604,6 +634,7 @@ int main() {
   testTypedIntegerTruncation(Context);
   testLLVMInlineAsmAndConstantMemorySyntax(Context);
   testLLVMCompareExchangeMemorySyntax(Context);
+  testUnweightedSuccessorNormalization(Context);
   testDiagnostics(Context);
   return Context.result();
 }
-- 
cgit v1.2.3