[LoopUnroll] Fix freqs for unconditional latches: N<=2

[jdenny-ornl]

As another step in issue #135812, this patch fixes block frequencies when LoopUnroll converts a conditional latch in an unrolled loop iteration to unconditional. It thus includes complete loop unrolling (the conditional backedge becomes an unconditional loop exit), which might be applied to the original loop or to its remainder loop.

As explained in detail in the header comments on the fixProbContradiction function that this patch introduces, these conversions mean LoopUnroll has proven that the original uniform latch probability is incorrect for the original loop iterations associated with the converted latches. However, LoopUnroll often is able to perform these corrections for only some iterations, leaving other iterations with the original latch probability, and thus corrupting the aggregate effect on the total frequency of the original loop body.

This patch ensures that the total frequency of the original loop body, summed across all its occurrences in the unrolled loop after the aforementioned conversions, is the same as in the original loop. Unlike other patches in this series, this patch cannot derive the required latch probabilities directly from the original uniform latch probability because it has been proven incorrect for some original loop iterations. Instead, this patch computes entirely new probabilities for the remaining N conditional latches in the unrolled loop.

This patch only handles N <= 2, for which it uses simple formulas to compute a single uniform probability across the latches. Future patches will handle N > 2.

This patch series does not consider the presence of non-latch loop exits, and I do not have a solid plan for that case. See fixme comments this patch introduces.

This patch depends on PR #182403 and PR #191008.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-support

Author: Joel E. Denny (jdenny-ornl)

Changes

As another step in issue #135812, this patch fixes block frequencies when LoopUnroll converts a conditional latch in an unrolled loop iteration to unconditional. It thus includes complete loop unrolling (the conditional backedge becomes an unconditional loop exit), which might be applied to the original loop or to its remainder loop.

As explained in detail in the header comments on the fixProbContradiction function that this patch introduces, these conversions mean LoopUnroll has proven that the original uniform latch probability is incorrect for the original loop iterations associated with the converted latches. However, LoopUnroll often is able to perform these corrections for only some iterations, leaving other iterations with the original latch probability, and thus corrupting the aggregate effect on the total frequency of the original loop body.

This patch ensures that the total frequency of the original loop body, summed across all its occurrences in the unrolled loop after the aforementioned conversions, is the same as in the original loop. Unlike other patches in this series, this patch cannot derive the required latch probabilities directly from the original uniform latch probability because it has been proven incorrect for some original loop iterations. Instead, this patch implements the following strategies to compute probabilities for the remaining N conditional latches in the unrolled loop:

• A. If N <= 2, use a simple formula to compute a single uniform probability across those latches.
• B. Otherwise, if -unroll-uniform-weights (a new option) is not specified, apply the original loop's probability to all N latches and then, as needed, adjust as few of them as possible.
• C. Otherwise, bisect the range [0,1] to find a single uniform probability across all N latches.

An issue with C is that it could impact compiler performance, so this patch makes it opt-in. Its appeal over B is that it treats all latches the same given that we have no evidence showing that any latch should have a higher or lower probability than any other. A has neither problem, but I do not know how to apply it for N > 2. More experience or feedback from others might determine that some strategies are not worthwhile to maintain.

This patch does not consider the presence of non-latch loop exits, and I do not have a solid plan for that case. See fixme comments this patch introduces.

---

Patch is 125.18 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/179520.diff

8 Files Affected:

• (modified) llvm/include/llvm/Support/BranchProbability.h (+1)
• (modified) llvm/lib/Transforms/Utils/LoopUnroll.cpp (+460-5)
• (modified) llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp (+2-2)
• (added) llvm/test/Transforms/LoopUnroll/branch-weights-freq/unroll-complete.ll (+1121)
• (modified) llvm/test/Transforms/LoopUnroll/branch-weights-freq/unroll-epilog.ll (+287-53)
• (added) llvm/test/Transforms/LoopUnroll/branch-weights-freq/unroll-partial-unconditional-latch.ll (+380)
• (modified) llvm/test/Transforms/LoopUnroll/branch-weights-freq/unroll-partial.ll (+2-1)
• (modified) llvm/test/Transforms/LoopUnroll/loop-probability-one.ll (+119-82)

diff --git a/llvm/include/llvm/Support/BranchProbability.h b/llvm/include/llvm/Support/BranchProbability.h
index 0b0d4343a9fcb..2f2224891ab30 100644
--- a/llvm/include/llvm/Support/BranchProbability.h
+++ b/llvm/include/llvm/Support/BranchProbability.h
@@ -46,6 +46,7 @@ class BranchProbability {
   LLVM_ABI BranchProbability(uint32_t Numerator, uint32_t Denominator);
 
   bool isZero() const { return N == 0; }
+  bool isOne() const { return N == D; }
   bool isUnknown() const { return N == UnknownN; }
 
   static BranchProbability getZero() { return BranchProbability(0); }
diff --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
index d9422afe5e82a..b33d5ce770553 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
@@ -65,6 +65,7 @@
 #include "llvm/Transforms/Utils/UnrollLoop.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <assert.h>
+#include <cmath>
 #include <numeric>
 #include <vector>
 
@@ -88,6 +89,11 @@ UnrollRuntimeEpilog("unroll-runtime-epilog", cl::init(false), cl::Hidden,
                     cl::desc("Allow runtime unrolled loops to be unrolled "
                              "with epilog instead of prolog."));
 
+static cl::opt<bool> UnrollUniformWeights(
+    "unroll-uniform-weights", cl::init(false), cl::Hidden,
+    cl::desc("If new branch weights must be found, work harder to keep them "
+             "uniform."));
+
 static cl::opt<bool>
 UnrollVerifyDomtree("unroll-verify-domtree", cl::Hidden,
                     cl::desc("Verify domtree after unrolling"),
@@ -438,6 +444,407 @@ static bool canHaveUnrollRemainder(const Loop *L) {
   return true;
 }
 
+// If LoopUnroll has proven OriginalLoopProb is incorrect for some iterations
+// of the original loop, adjust latch probabilities in the unrolled loop to
+// maintain the original total frequency of the original loop body.
+//
+// OriginalLoopProb is practical but imprecise
+// -------------------------------------------
+//
+// The latch branch weights that LLVM originally adds to a loop encode one latch
+// probability, OriginalLoopProb, applied uniformly across the loop's infinite
+// set of theoretically possible iterations.  While this uniform latch
+// probability serves as a practical statistic summarizing the trip counts
+// observed during profiling, it is imprecise.  Specifically, unless it is zero,
+// it is impossible for it to be the actual probability observed at every
+// individual iteration.  To see why, consider that the only way to actually
+// observe at run time that the latch probability remains non-zero is to profile
+// at least one loop execution that has an infinite number of iterations.  I do
+// not know how to profile an infinite number of loop iterations, and most loops
+// I work with are always finite.
+//
+// LoopUnroll proves OriginalLoopProb is incorrect
+// ------------------------------------------------
+//
+// LoopUnroll reorganizes the original loop so that loop iterations are no
+// longer all implemented by the same code, and then it analyzes some of those
+// loop iteration implementations independently of others.  In particular, it
+// converts some of their conditional latches to unconditional.  That is, by
+// examining code structure without any profile data, LoopUnroll proves that the
+// actual latch probability at the end of such an iteration is either 1 or 0.
+// When an individual iteration's actual latch probability is 1 or 0, that means
+// it always behaves the same, so it is impossible to observe it as having any
+// other probability.  The original uniform latch probability is rarely 1 or 0
+// because, when applied to all possible iterations, that would yield an
+// estimated trip count of infinity or 1, respectively.
+//
+// Thus, the new probabilities of 1 or 0 are proven corrections to
+// OriginalLoopProb for individual iterations in the original loop.  However,
+// LoopUnroll often is able to perform these corrections for only some
+// iterations, leaving other iterations with OriginalLoopProb, and thus
+// corrupting the aggregate effect on the total frequency of the original loop
+// body.
+//
+// Adjusting latch probabilities
+// -----------------------------
+//
+// This function ensures that the total frequency of the original loop body,
+// summed across all its occurrences in the unrolled loop after the
+// aforementioned latch conversions, is the same as in the original loop.  To do
+// so, it adjusts probabilities on the remaining conditional latches.  However,
+// it cannot derive the new probabilities directly from the original uniform
+// latch probability because the latter has been proven incorrect for some
+// original loop iterations.
+//
+// There are often many sets of latch probabilities that can produce the
+// original total loop body frequency.  If there are many remaining conditional
+// latches and !UnrollUniformWeights, this function just quickly hacks a few of
+// their probabilities to restore the original total loop body frequency.
+// Otherwise, it tries harder to determine less arbitrary probabilities.
+static void fixProbContradiction(UnrollLoopOptions ULO,
+                                 BranchProbability OriginalLoopProb,
+                                 bool CompletelyUnroll,
+                                 std::vector<unsigned> &IterCounts,
+                                 const std::vector<BasicBlock *> &CondLatches,
+                                 std::vector<BasicBlock *> &CondLatchNexts) {
+  // Runtime unrolling is handled later in LoopUnroll not here.
+  //
+  // There are two scenarios in which LoopUnroll sets ProbUpdateRequired to true
+  // because it needs to update probabilities that were originally
+  // OriginalLoopProb, but only in one scenario has LoopUnroll proven
+  // OriginalLoopProb incorrect for iterations within the original loop:
+  // - If ULO.Runtime, LoopUnroll adds new guards that enforce new reaching
+  //   conditions for new loop iteration implementations (e.g., one unrolled
+  //   loop iteration executes only if at least ULO.Count original loop
+  //   iterations remain).  Those reaching conditions dictate how conditional
+  //   latches can be converted to unconditional (e.g., within an unrolled loop
+  //   iteration, there is no need to recheck the number of remaining original
+  //   loop iterations).  None of this reorganization alters the set of possible
+  //   original loop iteration counts or proves OriginalLoopProb incorrect for
+  //   any of the original loop iterations.  Thus, LoopUnroll derives
+  //   probabilities for the new guards and latches directly from
+  //   OriginalLoopProb based on the probabilities that their reaching
+  //   conditions would occur in the original loop.  Doing so maintains the
+  //   total frequency of the original loop body.
+  // - If !ULO.Runtime, LoopUnroll initially adds new loop iteration
+  //   implementations, which have the same latch probabilities as in the
+  //   original loop because there are no new guards that change their reaching
+  //   conditions.  Sometimes, LoopUnroll is then done, and so does not set
+  //   ProbUpdateRequired to true.  Other times, LoopUnroll then proves that
+  //   some latches are unconditional, directly contradicting OriginalLoopProb
+  //   for the corresponding original loop iterations.  That reduces the set of
+  //   possible original loop iteration counts, possibly producing a finite set
+  //   if it manages to eliminate the backedge.  LoopUnroll has to choose a new
+  //   set of latch probabilities that produce the same total loop body
+  //   frequency.
+  //
+  // This function addresses the second scenario only.
+  if (ULO.Runtime)
+    return;
+
+  // If CondLatches.empty(), there are no latch branches with probabilities we
+  // can adjust.  That should mean that the actual trip count is always exactly
+  // the number of remaining unrolled iterations, and so OriginalLoopProb should
+  // have yielded that trip count as the original loop body frequency.  Of
+  // course, OriginalLoopProb could be based on bad profile data, but there is
+  // nothing we can do about that here.
+  if (CondLatches.empty())
+    return;
+
+  // If the original latch probability is 1, the original frequency is infinity.
+  // Leaving all remaining probabilities set to 1 might or might not get us
+  // there (e.g., a completely unrolled loop cannot be infinite), but it is the
+  // closest we can come.
+  assert(!OriginalLoopProb.isUnknown() &&
+         "Expected to have loop probability to fix");
+  if (OriginalLoopProb.isOne())
+    return;
+
+  // FreqDesired is the frequency implied by the original loop probability.
+  double FreqDesired = 1 / (1 - OriginalLoopProb.toDouble());
+
+  // Get the probability at CondLatches[I].
+  auto GetProb = [&](unsigned I) {
+    BranchInst *B = cast<BranchInst>(CondLatches[I]->getTerminator());
+    bool FirstTargetIsNext = B->getSuccessor(0) == CondLatchNexts[I];
+    return getBranchProbability(B, FirstTargetIsNext).toDouble();
+  };
+
+  // Set the probability at CondLatches[I] to Prob.
+  auto SetProb = [&](unsigned I, double Prob) {
+    BranchInst *B = cast<BranchInst>(CondLatches[I]->getTerminator());
+    bool FirstTargetIsNext = B->getSuccessor(0) == CondLatchNexts[I];
+    bool Success = setBranchProbability(
+        B, BranchProbability::getBranchProbability(Prob), FirstTargetIsNext);
+    assert(Success && "Expected to be able to set branch probability");
+  };
+
+  // Set all probabilities in CondLatches to Prob.
+  auto SetAllProbs = [&](double Prob) {
+    for (unsigned I = 0, E = CondLatches.size(); I < E; ++I)
+      SetProb(I, Prob);
+  };
+
+  // If UnrollUniformWeights or n <= 2, we choose the simplest probability model
+  // we can think of: every remaining conditional branch instruction has the
+  // same probability, Prob, of continuing to the next iteration.  This model
+  // has several helpful properties:
+  // - There is only one search parameter, Prob.
+  // - We have no reason to think one latch branch's probability should be
+  //   higher or lower than another, and so this model makes them all the same.
+  //   In the worst cases, we thus avoid setting just some probabilities to 0 or
+  //   1, which can unrealistically make some code appear unreachable.  There
+  //   are cases where they *all* must become 0 or 1 to achieve the total
+  //   frequency of original loop body, and our model does permit that.
+  // - The frequency, FreqOne, of the original loop body in a single iteration
+  //   of the unrolled loop is computed by a simple polynomial, where p=Prob,
+  //   n=CondLatches.size(), and c_i=IterCounts[i]:
+  //
+  //     FreqOne = Sum(i=0..n)(c_i * p^i)
+  //
+  // - If the backedge has been eliminated:
+  //   - FreqOne is the total frequency of the original loop body in the
+  //     unrolled loop.
+  //   - If Prob == 1, the total frequency of the original loop body is exactly
+  //     the number of remaining loop iterations, as expected because every
+  //     remaining loop iteration always then executes.
+  // - If the backedge remains:
+  //   - Sum(i=0..inf)(FreqOne * p^(n*i)) = FreqOne / (1 - p^n) is the total
+  //     frequency of the original loop body in the unrolled loop, regardless of
+  //     whether the backedge is conditional or unconditional.
+  //   - As Prob approaches 1, the total frequency of the original loop body
+  //     approaches infinity, as expected because the loop approaches never
+  //     exiting.
+  // - For n <= 2, we can use simple formulas to solve the above polynomial
+  //   equation exactly for p without performing a search.   For n == 2, we use
+  //   ComputeProbForQuadratic below.  For n == 1, we use ComputeProb below.
+  // - For n > 2, evaluating each point in the search space, using ComputeFreq
+  //   below, requires about as few instructions as we could hope for.  That is,
+  //   the probability is constant across the conditional branches, so the only
+  //   computation is across conditional branches and any backedge, as required
+  //   for any model for Prob.
+  // - Prob == 1 produces the maximum possible total frequency for the original
+  //   loop body, as described above.  Prob == 0 produces the minimum, 0.
+  //   Increasing or decreasing Prob monotonically increases or decreases the
+  //   frequency, respectively.  Thus, for every possible frequency, there
+  //   exists some Prob that can produce it, and we can easily use bisection to
+  //   search the problem space.
+
+  // When iterating for a solution, we stop early if we find probabilities
+  // that produce a Freq whose difference from FreqDesired is small
+  // (FreqPrec).  Otherwise, we expect to compute a solution at least that
+  // accurate (but surely far more accurate).
+  const double FreqPrec = 1e-6;
+
+  // Compute the new frequency produced by using Prob throughout CondLatches.
+  auto ComputeFreq = [&](double Prob) {
+    double ProbReaching = 1;        // p^0
+    double FreqOne = IterCounts[0]; // c_0*p^0
+    for (unsigned I = 0, E = CondLatches.size(); I < E; ++I) {
+      ProbReaching *= Prob;                        // p^(I+1)
+      FreqOne += IterCounts[I + 1] * ProbReaching; // c_(I+1)*p^(I+1)
+    }
+    double ProbReachingBackedge = CompletelyUnroll ? 0 : ProbReaching;
+    assert(FreqOne > 0 && "Expected at least one iteration before first latch");
+    if (ProbReachingBackedge == 1)
+      return std::numeric_limits<double>::infinity();
+    return FreqOne / (1 - ProbReachingBackedge);
+  };
+
+  // Compute the probability that, used throughout CondLatches where
+  // CondLatches.size() == 2, gets as close as possible to FreqDesired.
+  auto ComputeProbForQuadratic = [&]() {
+    // The polynomial is quadratic, so just solve it.
+    double A = IterCounts[2] + (CompletelyUnroll ? 0 : FreqDesired);
+    double B = IterCounts[1];
+    double C = IterCounts[0] - FreqDesired;
+    assert(A > 0 && "Expected iterations after last conditional latch");
+    double Prob = (-B + sqrt(B * B - 4 * A * C)) / (2 * A);
+    // If it computes an invalid Prob, FreqDesired is impossibly low or high.
+    // Otherwise, Prob should produce nearly FreqDesired.
+    assert((Prob < 0 || Prob > 1 ||
+            fabs(ComputeFreq(Prob) - FreqDesired) < FreqPrec) &&
+           "Expected accurate frequency when quadratic case is possible");
+    Prob = std::max(Prob, 0.);
+    Prob = std::min(Prob, 1.);
+    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
+      // that analytically the probability has to be so near 1 that, in floating
+      // point, the frequency is computed as infinite.
+      //
+      // TODO: Cover this case in the test suite if you can.
+      Freq = std::numeric_limits<double>::infinity();
+    } else {
+      assert(FreqBeforeCompute > 0 &&
+             "Expected at least one iteration before first latch");
+      // In this equation, if we replace the left-hand side with FreqDesired and
+      // then solve for ProbCompute, we get the ProbCompute formula above.
+      Freq = (FreqBeforeCompute + Freq * ProbCompute) /
+             (1 - ProbReachingBackedge * ProbCompute);
+    }
+    return {ProbCompute, Freq};
+  };
+
+  // Determine and set branch weights.
+  //
+  // Prob < 0 and Prob > 1 cannot be represented as branch weights.  We might
+  // compute such a Prob if FreqDesired is impossible (e.g., due to bad profile
+  // data) for the maximum trip count we have determined when completely
+  // unrolling.  In that case, so just go with whichever is closest.
+  if (CondLatches.size() == 2) {
+    // The polynomial is quadratic, so just solve it.
+    SetAllProbs(ComputeProbForQuadratic());
+  } else if (CondLatches.size() == 1 || !UnrollUniformWeights) {
+    // Either:
+    // - There's just one conditional latch, so just compute the probability
+    //   it requires to produce the original total frequency.
+    // - The polynomial is too complex for a simple formula and the quick and
+    //   dirty fix has been selected.  Adjust probabilities starting from the
+    //   first latch, which has the most influence on the total frequency, so
+    //   starting there sh...
[truncated]

[jdenny-ornl]

Would it be better to move some of the new LoopUnroll.cpp code into a new module and possibly extract some of the lambdas into separate functions?

[jdenny-ornl]

I plan to break this patch into smaller patches to facilitate the review. If anyone has already started reviewing and that would interfere, please say so now.

[jdenny-ornl]

I plan to break this patch into smaller patches to facilitate the review.

Done.

[github-actions]

🐧 Linux x64 Test Results

• 193106 tests passed
• 4996 tests skipped

✅ The build succeeded and all tests passed.

[jdenny-ornl]

ping

[jdoerfert]

Is there a test we can see the before/after effect?
Or could we precommit one?

[jdenny-ornl]

Is there a test we can see the before/after effect? Or could we precommit one?

I've extracted all tests to PR #191008 so the tests diff in this PR shows just the improvements.

[jdoerfert]

I see the problem this is solving and I see how it fixes the tests.
That said, I don't fully grasp how we know the IterCounts.
Given that they are in a different commit and I can assume they are accurate, I am fine with this.
As I wrote this, I realized something and left a comment below.
Otherwise, I think this looks good.

[jdoerfert]

I am not sure about this. We have their current probabilities, no?
So if latch 1 has prob p and latch 2 has prob q, couldn't we preserve the ratio p/q for the new probabilities we pick?
For the overall trip count, it doesn't matter, but for probabilities in the loop, it does, right?
This also only affects the N >= 2 case.
The equation system we get for N == 2 would then be

A * p' * q' + B * q' + C = 0
p' / q' = p / q

or

A * p' * p' * q / p + B * p' * q / p + C = 0

[jdenny-ornl]

Here, we know the original loop has only one latch. If it did not, the getLoopProbability call in UnrollLoop would return an unknown probability, and so fixProbContradiction would not be called.

With only one original latch, there is only one original probability, OriginalLoopProb.

Multiple latches show up due to unrolling, but I know of nothing suggesting they should have different probabilities (except that LoopUnroll proves some are unconditional).

So this patch actually does as you suggest: tries to preserve the original ratio (1:1) between latch probabilities.

[jdenny-ornl]

I don't fully grasp how we know the IterCounts. Given that they are in a different commit and I can assume they are accurate, I am fine with this.

Feel free to add questions to that PR. That might lead to improved source comments.

Otherwise, I think this looks good.

Thanks for the review!