[MIR] Normalize explicit successor probabilities

1 files changed, 25 insertions, 0 deletions

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);