[llvm-branch-commits] [llvm] [LoopUnroll] Fix block frequencies for epilogue (PR #159163)

[Joel E. Denny via llvm-branch-commits]

================
@@ -1159,10 +1160,17 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
     // the unrolled loop as a whole without considering the branch weights for
     // each unrolled iteration's latch within it, we store the new trip count as
     // separate metadata.
-    unsigned NewTripCount = *OriginalTripCount / ULO.Count;
-    if (!ULO.Runtime && *OriginalTripCount % ULO.Count)
-      NewTripCount += 1;
-    setLoopEstimatedTripCount(L, NewTripCount);
+    if (!OriginalLoopProb.isUnknown() && ULO.Runtime && EpilogProfitability) {
+      // Where p is always the probability of executing at least 1 more
+      // iteration, the probability for at least n more iterations is p^n.
+      setLoopProbability(L, OriginalLoopProb.pow(ULO.Count));
+    }
+    if (OriginalTripCount) {
+      unsigned NewTripCount = *OriginalTripCount / ULO.Count;
+      if (!ULO.Runtime && *OriginalTripCount % ULO.Count)
+        NewTripCount += 1;
----------------
jdenny-ornl wrote:

Thanks.  Done.  I was thinking in Python, I guess.

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