[llvm] [VPlan] Preserve branch weights from VPlan0 through to codegen. (PR #213143)

[Florian Hahn via llvm-commits]

================
@@ -1055,6 +1058,77 @@ SmallVector<VPUser *> vputils::collectUsersRecursively(VPValue *V) {
   return Users.takeVector();
 }
 
+/// Returns the probability of reaching each unique successor of \p VPBB, taken
+/// from the branch weights recorded on its terminator, or unknown if not
+/// available. See llvm::getBranchProbability in
+/// llvm/Transforms/Utils/LoopUtils.h for the IR version.
+static SmallVector<std::pair<const VPBasicBlock *, BranchProbability>, 2>
+getSuccessorProbabilities(const VPBasicBlock *VPBB) {
+  ArrayRef<VPBlockBase *> Successors = VPBB->getSuccessors();
+  // With a single successor the edge is always taken and needs no weights.
+  if (Successors.size() == 1)
+    return {{cast<VPBasicBlock>(Successors[0]), BranchProbability::getOne()}};
+
+  // Take the branch weights off the terminator. Without usable weights all
+  // successors have unknown probability; zero the weights, so the accumulation
+  // below still visits each of them.
+  SmallVector<uint32_t> Weights;
+  auto *Term = dyn_cast_if_present<VPInstruction>(VPBB->getTerminator());
+  if (!Term ||
+      !extractBranchWeights(Term->getMetadata(LLVMContext::MD_prof), Weights) ||
+      Weights.size() != Successors.size())
+    Weights.assign(Successors.size(), 0);
+  uint64_t Total = sum_of(Weights, uint64_t(0));
+
+  SmallMapVector<const VPBasicBlock *, uint64_t, 2> WeightPerSuccessor;
+  for (const auto &[Succ, Weight] : zip_equal(Successors, Weights))
+    WeightPerSuccessor[cast<VPBasicBlock>(Succ)] += Weight;
+
+  return map_to_vector<2>(WeightPerSuccessor, [Total](const auto &SuccWeight) {
+    auto [Succ, Weight] = SuccWeight;
+    return std::make_pair(
+        Succ, Total == 0
+                  ? BranchProbability::getUnknown()
+                  : BranchProbability::getBranchProbability(Weight, Total));
+  });
+}
+
+DenseMap<const VPBasicBlock *, VPExecutionProbability>
+vputils::computeExecutionProbabilities(ArrayRef<VPBasicBlock *> Blocks) {
+  assert(!Blocks.empty() && "expected at least the header block");
+  // Push each block's probability along its outgoing edges, accumulating it in
+  // the successors. Blocks is in reverse post-order and the blocks form a DAG
+  // (the backedge of a loop region is implicit), so all incoming edges of a
+  // block have contributed by the time it is visited and its probability is
+  // final.
+  DenseMap<const VPBasicBlock *, VPExecutionProbability> Probabilities;
+  Probabilities.reserve(Blocks.size());
+  // The header (first block) always executes.
+  Probabilities[Blocks.front()] = VPExecutionProbability::getOne();
+  for (VPBasicBlock *VPBB : Blocks.drop_front())
+    Probabilities[VPBB] = VPExecutionProbability::getZero();
+
+  for (VPBasicBlock *VPBB : Blocks) {
+    VPExecutionProbability SrcProb = Probabilities.at(VPBB);
+    for (auto [Succ, EdgeProb] : getSuccessorProbabilities(VPBB)) {
+      VPExecutionProbability &SuccProb = Probabilities.at(Succ);
+      // An unknown edge or predecessor poisons the successor: its probability
+      // is only known if all edges on paths reaching it carry branch weights.
+      if (SrcProb.isUnknown() || EdgeProb.isUnknown() || SuccProb.isUnknown()) {
+        SuccProb = VPExecutionProbability::getUnknown();
+        continue;
+      }
+      VPExecutionProbability Contribution = SrcProb * EdgeProb;
----------------
fhahn wrote:

Fair enough, at the very low range of probabilties there are small-ish differences; with how the info is now managed (own kind), it is fairly straight-forward to change it without moving things around too much. updated thanks.

https://github.com/llvm/llvm-project/pull/213143