[llvm] [VPlan] Record estimated branch probabilities on VPlan0 for cost modeling (PR #216172)

[Mircea Trofin via llvm-commits]

================
@@ -1055,6 +1056,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>
+vputils::computeBlockProbabilities(ArrayRef<VPBasicBlock *> Blocks) {
----------------
mtrofin wrote:

could this use the extensibility mechanism in BlockFrequencyInfo - we use the same way to derive BFI from BPI for either IR or Machine BBs, so probably VPBasicBlocks can work, too.

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