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

[Luke Lau via llvm-commits]

================
@@ -1050,6 +1051,79 @@ SmallVector<VPUser *> vputils::collectUsersRecursively(VPValue *V) {
   return Users.takeVector();
 }
 
+/// Returns the probability of the edge from \p Src to \p Dst, taken from the
+/// branch weights recorded on Src's terminator, or unknown if not available.
+/// See llvm::getBranchProbability in llvm/Transforms/Utils/LoopUtils.h for the
+/// IR version.
+static BranchProbability getEdgeProbability(const VPBasicBlock *Src,
+                                            const VPBasicBlock *Dst) {
+  // With a single successor the edge is always taken.
+  ArrayRef<VPBlockBase *> Successors = Src->getSuccessors();
+  if (Successors.size() == 1)
+    return BranchProbability::getOne();
+
+  auto *Term = dyn_cast_if_present<VPInstruction>(Src->getTerminator());
+  SmallVector<uint32_t> Weights;
+  if (!Term ||
+      !extractBranchWeights(Term->getMetadata(LLVMContext::MD_prof), Weights) ||
+      Weights.size() != Successors.size())
+    return BranchProbability::getUnknown();
+
+  uint64_t Total = sum_of(Weights, uint64_t(0));
+  if (Total == 0)
+    return BranchProbability::getUnknown();
+
+  // Sum the weights of all edges from Src to Dst; the same block may be the
+  // destination of multiple successors, e.g. for switches.
+  uint64_t ToDst = 0;
+  for (const auto &[Succ, Weight] : zip(Successors, Weights))
+    if (Succ == Dst)
+      ToDst += Weight;
+  return BranchProbability::getBranchProbability(ToDst, Total);
+}
+
+DenseMap<const VPBasicBlock *, BranchProbability>
----------------
lukel97 wrote:

Is there a particular reason why you're using BranchProbability for blocks as opposed to BlockFrequency? The former stores both a numerator + denominator whereas the latter would always be relative to the entry block. Avoiding the mixed denominators might be better for accuracy?

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