[llvm] [LoopUnroll] Fix freqs for unconditional latches: N>2, fast (PR #182404)

[Joel E. Denny via llvm-commits]

================
@@ -625,13 +639,120 @@ static void fixProbContradiction(UnrollLoopOptions ULO,
     return Prob;
   };
 
+  // Compute the probability required at CondLatches[ComputeIdx] to get as close
+  // as possible to FreqDesired without replacing probabilities elsewhere in
+  // CondLatches.  Return {Prob, Freq} where 0 <= Prob <= 1 and Freq is the new
+  // frequency.
+  auto ComputeProb = [&](unsigned ComputeIdx) -> std::pair<double, double> {
+    assert(ComputeIdx < CondLatches.size());
+
+    // Accumulate the frequency from before ComputeIdx into FreqBeforeCompute,
+    // and accumulate the rest in Freq without yet multiplying the latter by any
+    // probability for ComputeIdx (i.e., treat it as 1 for now).
+    double ProbReaching = 1;     // p^0
+    double Freq = IterCounts[0]; // c_0*p^0
+    double FreqBeforeCompute;
+    for (unsigned I = 0, E = CondLatches.size(); I < E; ++I) {
+      // Get the branch probability for CondLatches[I].
+      double Prob;
+      if (I == ComputeIdx) {
+        FreqBeforeCompute = Freq;
+        Freq = 0;
+        Prob = 1;
+      } else {
+        Prob = GetProb(I);
+      }
+      ProbReaching *= Prob;                     // p^(I+1)
+      Freq += IterCounts[I + 1] * ProbReaching; // c_(I+1)*p^(I+1)
+    }
+
+    // Compute the required probability, and limit it to a valid probability (0
+    // <= p <= 1).  See the Freq formula below for how to derive the ProbCompute
+    // formula.
+    double ProbReachingBackedge = CompletelyUnroll ? 0 : ProbReaching;
+    double ProbComputeNumerator = FreqDesired - FreqBeforeCompute;
+    double ProbComputeDenominator = Freq + FreqDesired * ProbReachingBackedge;
+    double ProbCompute;
+    if (ProbComputeNumerator <= 0) {
+      // FreqBeforeCompute has already reached or surpassed FreqDesired, so add
+      // no more frequency.  It is possible that ProbComputeDenominator == 0
+      // here because some latch probability (maybe the original) was set to
+      // zero, so this check avoids setting ProbCompute=1 (in the else if below)
+      // and division by zero where the numerator <= 0 (in the else below).
+      ProbCompute = 0;
+    } else if (ProbComputeDenominator == 0) {
+      // Analytically, this case seems impossible.  It would occur if either:
+      // - Both Freq and FreqDesired are zero.  But the latter would cause
+      //   ProbComputeNumerator < 0, which we catch above, and FreqDesired
+      //   should always be >= 1 anyway.
+      // - There are no iterations after CondLatches[ComputeIdx], not even via
+      //   a backedge, so that both Freq and ProbReachingBackedge are zero.
+      //   But iterations should exist after even the last conditional latch.
+      // - Some latch probability (maybe the original) was set to zero so that
+      //   both Freq and ProbReachingBackedge are zero.  But that should not
+      //   have happened because, according to the above ProbComputeNumerator
+      //   check, we have not yet reached FreqDesired (which, if the original
+      //   latch probability is zero, is just 1 and thus always reached or
+      //   surpassed).
+      //
+      // Numerically, perhaps this case is possible.  We interpret it to mean we
+      // need more frequency (ProbComputeNumerator > 0) but have no way to get
+      // any (ProbComputeDenominator is analytically too small to distinguish it
+      // from 0 in floating point), suggesting infinite probability is needed,
+      // but 1 is the maximum valid probability and thus the best we can do.
+      //
+      // TODO: Cover this case in the test suite if you can.
+      ProbCompute = 1;
+    } else {
+      ProbCompute = ProbComputeNumerator / ProbComputeDenominator;
+      ProbCompute = std::max(ProbCompute, 0.);
+      ProbCompute = std::min(ProbCompute, 1.);
+    }
+
+    // Compute the resulting total frequency.
+    if (ProbReachingBackedge * ProbCompute == 1) {
+      // Analytically, this case seems impossible.  It requires that there is a
+      // backedge and that FreqDesired == infinity so that every conditional
+      // latch's probability had to be set to 1.  But FreqDesired == infinity
+      // means OriginalLoopProb.isOne(), which we guarded against earlier.
+      //
+      // Numerically, perhaps this case is possible.  We interpret it to mean
----------------
jdenny-ornl wrote:

I'm not sure if it can happen.  I never came up with a test to produce it.

I added a remark.  Is that reasonable?

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