[LoopPeeling] Get rid of Phis that become invariant after N steps

[mkazantsev]

Summary

This patch is a generalization of the improvement introduced in rL296898.
Previously, we were able to peel one iteration of a loop to get rid of a Phi that becomes
an invariant on the 2nd iteration. In more general case, if a Phi becomes invariant after
N iterations, we can peel N times and turn it into invariant.
In order to do this, we for every Phi in loop's header we define the Invariant Depth value
which is calculated as follows:

Given %x = phi <Inputs from above the loop>, ..., [%y, %back.edge].

If %y is a loop invariant, then Depth(%x) = 1.
If %y is a Phi from the loop header, Depth(%x) = Depth(%y) + 1.
Otherwise, Depth(%x) is infinite.

Notice that if we peel a loop, all Phis with Depth = 1 become invariants,
and all other Phis with finite depth decrease the depth by 1.
Thus, peeling N first iterations allows us to turn all Phis with Depth <= N
into invariants.

[anna]

Hi Max,
Couple of comments inline. The logic as such looks correct to me.

lib/Transforms/Utils/LoopUnrollPeel.cpp

101 ↗

(On Diff #93904)

You can use a SmallDenseMap here since number of phis within a single basic block (i.e. header) is generally a small number.

153 ↗

(On Diff #93904)

Nit: iterations.

test/Transforms/LoopUnroll/peel-loop-not-forced.ll

90 ↗

(On Diff #93904)

Could you please add a test case that exercises the logic of choosing the depth based on the code size? i.e. say we have 2 phi's each of which have maxDepth of 3 and 4, but due to the code size restriction on peeling, you need to peel by only 3.

[mkazantsev]

Added a new test, fixed typo, used SmallDenseMap for storage of Phis depth.

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

134 ↗

(On Diff #94044)

Actually this one is not needed, because if the value is present in map, we will return it early in CalculateInvariantDepth.

[mkazantsev]

Removed redundant piece of code.

[sanjoy]

Minor comments inline.

lib/Transforms/Utils/LoopUnrollPeel.cpp

109 ↗

(On Diff #94050)

I think the logic may be more understandable if instead of returning 0 for a loop-varying-non-header-phi %x, you return DepthInfinity where DepthInfinity is an alias for UINT_MAX. Returning 0 in these cases seems a bit odd, since given your scheme I'd expect a dept of 0 meaning that the value is loop invariant, not that it is "hopelessly loop invariant".

111 ↗

(On Diff #94050)

I think this is too large to be a lambda -- can you please extract it out into a static helper?

126 ↗

(On Diff #94050)

Use else if (PHINode *IncPhi = dyn_cast<PHINode>(Input)) {

137 ↗

(On Diff #94050)

One TODO here would be to consider some minor binary operations as well, like:

A = PHI(0, C)
B = PHI(0, 5)
C = B + 1

Ideally the depth of A should be 2, not 0.

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

137 ↗

(On Diff #94050)

Good catch!

[sanjoy]

Comments inline.

lib/Transforms/Utils/LoopUnrollPeel.cpp

50 ↗

(On Diff #94201)

s/InfiniteDepth/InfiniteInvariantDepth/

84 ↗

(On Diff #94201)

I should have mentioned this in the earlier revision, but do you think computeIterationsToInvariance is a better name for this function? Depth does not clarify what we're computing here very much.

If you do this, please also rename InfiniteDepth (this suggestion also invalidates the "expand Inv to Invariant" suggestion).

85 ↗

(On Diff #94201)

I'd expand Inv to Invariant.

94 ↗

(On Diff #94201)

This is a suggestion, feel free to ignore:

getLoopLatch isn't super-cheap (it re-discovers the latch every time). How about passing the latch explicitly to this function?

109 ↗

(On Diff #94201)

s/estimate/compute an upper bound on/

138 ↗

(On Diff #94201)

I don't think there is a need to explain the algorithm here (you should do that on the implementation, as you've already done). Instead renaming the helper to computeIterationsToInvariance will make the logic here more obvious.

149 ↗

(On Diff #94201)

Why did you change the condition to 2 * LoopSize <= UP.Threshold?

165 ↗

(On Diff #94201)

Can UP.Threshold be 0 (say it was directly set by the user by UnrollThreshold)?

168 ↗

(On Diff #94201)

I don't know the rest of the code well enough to be conclusive on this, but should this be UP.PeelCount = std::max(UP.PeelCount, DesiredPeelCount);? That is, if we've earlier decided, due to some other reasons, that peeling for 5 iterations is a good idea, and here we decide that peeling for only 3 iterations is a good idea, perhaps we should still peel for 5 iterations?

Actually, even if UP.PeelCount > DesiredPeelCount is impossible, I think UP.PeelCount = std::max(UP.PeelCount, DesiredPeelCount); seems cleaner.

test/Transforms/LoopUnroll/peel-loop-not-forced.ll

60 ↗

(On Diff #94201)

s/trice/thrice/

92 ↗

(On Diff #94201)

s/trice/thrice/

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

149 ↗

(On Diff #94201)

Because initially it was wrong. The threshold applies to the total number of iterations after peeling, not to loop's size, see line 202. If we peel once, the cost would be 2 * LoopSize.

[sanjoy]

lib/Transforms/Utils/LoopUnrollPeel.cpp

149 ↗

(On Diff #94201)

Can you please fix that separately with a test?

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

165 ↗

(On Diff #94201)

Only if loop size is also 0, we enter here witn condition 2 * LoopSize <= UP.Threshold :) So I Don't think it's real.

168 ↗

(On Diff #94201)

Line 126

UP.PeelCount = 0;

It is never changed after that. The puspose of this method is to only set it once.

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

149 ↗

(On Diff #94201)

https://reviews.llvm.org/D31753

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

165 ↗

(On Diff #94201)

Line 131 rejects empty loops:

if (!L->empty())
  return;

[mkazantsev]

Rebase & some renames.

[sanjoy]

lgtm with minor comments

lib/Transforms/Utils/LoopUnrollPeel.cpp

94 ↗

(On Diff #94201)

Please add an assert that BackEdge is, in fact, the latch.

165 ↗

(On Diff #94201)

I don't think line 131 has anything to do with LoopSize, it is just checking if there are *subloops* or not.

But I think your point on LoopSize being non-zero is correct. Can you please add an assert?

168 ↗

(On Diff #94201)

SGTM.

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

165 ↗

(On Diff #94201)

Ok. Zero loop size is not allowed where it is calculated, like

// Don't allow an estimate of size zero.  This would allows unrolling of loops
// with huge iteration counts, which is a compile time problem even if it's
// not a problem for code quality. Also, the code using this size may assume
// that each loop has at least three instructions (likely a conditional
// branch, a comparison feeding that branch, and some kind of loop increment
// feeding that comparison instruction).
LoopSize = std::max(LoopSize, BEInsns + 1);

So I'll add this assert in the very beginning of the method.

[lebedev.ri]

@mkazantsev hi.
I have encountered a missing peeling issue, and here is the most simplified version of the code: https://godbolt.org/g/xZgb4T
If i manually use opt, and specify -loop-unroll -unroll-force-peel-count=2, the peeling happens as i would expect.
Which i suppose means, only the analysis (calculateIterationsToInvariance()) does not support that pattern.
If i change the line 117 of this diff of that function with something like:

-  }
+  } else ToInvariance = 2u;

it also peels that case correctly, so i do think it is just an analysis problem.

Do you plan on working on the loop peeling again anytime soon?
If yes, then i guess i should not worry.
If not, any pointers? :)

[mkazantsev]

Hi @lebedev.ri

This is an interesting case, however it has nothing to do with this patch. This patch is only for pattern "some Phi becomes invariant after N steps", and in your situation is "some condition becomes invariant after few steps". It is a completely different situation. In case if your condition was if (x < 60), using peeling here would be complete bizzare.

I personally don't plan to do anything in peeling in the observable future.

In LLVM, we have a pass InductiveRangeChecksElimination which is not included into clang pipeline, but it handles exactly this kind of cases. Try running opt with -irce option and see what happens.

• Max

llvm/trunk/lib/Transforms/Utils/LoopUnrollPeel.cpp

 static cl::opt<unsigned> UnrollPeelMaxCount(
     "unroll-peel-max-count", cl::init(7), cl::Hidden,
     cl::desc("Max average trip count which will cause loop peeling."));
 
 static cl::opt<unsigned> UnrollForcePeelCount(
     "unroll-force-peel-count", cl::init(0), cl::Hidden,
     cl::desc("Force a peel count regardless of profiling information."));
 
+// Designates that a Phi is estimated to become invariant after an "infinite"
+// number of loop iterations (i.e. only may become an invariant if the loop is
+// fully unrolled).
+static const unsigned InfiniteIterationsToInvariance = UINT_MAX;
+
 // Check whether we are capable of peeling this loop.
 static bool canPeel(Loop *L) {
   // Make sure the loop is in simplified form
   if (!L->isLoopSimplifyForm())
     return false;
 
   // Only peel loops that contain a single exit
   if (!L->getExitingBlock() || !L->getUniqueExitBlock())
     return false;
 
   // Don't try to peel loops where the latch is not the exiting block.
   // This can be an indication of two different things:
   // 1) The loop is not rotated.
   // 2) The loop contains irreducible control flow that involves the latch.
   if (L->getLoopLatch() != L->getExitingBlock())
     return false;
 
   return true;
 }
 
+// This function calculates the number of iterations after which the given Phi
+// becomes an invariant. The pre-calculated values are memorized in the map. The
+// function (shortcut is I) is calculated according to the following definition:
+// Given %x = phi <Inputs from above the loop>, ..., [%y, %back.edge].
+//   If %y is a loop invariant, then I(%x) = 1.
+//   If %y is a Phi from the loop header, I(%x) = I(%y) + 1.
+//   Otherwise, I(%x) is infinite.
+// TODO: Actually if %y is an expression that depends only on Phi %z and some
+//       loop invariants, we can estimate I(%x) = I(%z) + 1. The example
+//       looks like:
+//         %x = phi(0, %a),  <-- becomes invariant starting from 3rd iteration.
+//         %y = phi(0, 5),
+//         %a = %y + 1.
+static unsigned calculateIterationsToInvariance(
+    PHINode *Phi, Loop *L, BasicBlock *BackEdge,
+    SmallDenseMap<PHINode *, unsigned> &IterationsToInvariance) {
+  assert(Phi->getParent() == L->getHeader() &&
+         "Non-loop Phi should not be checked for turning into invariant.");
+  assert(BackEdge == L->getLoopLatch() && "Wrong latch?");
+  // If we already know the answer, take it from the map.
+  auto I = IterationsToInvariance.find(Phi);
+  if (I != IterationsToInvariance.end())
+    return I->second;
+
+  // Otherwise we need to analyze the input from the back edge.
+  Value *Input = Phi->getIncomingValueForBlock(BackEdge);
+  // Place infinity to map to avoid infinite recursion for cycled Phis. Such
+  // cycles can never stop on an invariant.
+  IterationsToInvariance[Phi] = InfiniteIterationsToInvariance;
+  unsigned ToInvariance = InfiniteIterationsToInvariance;
+
+  if (L->isLoopInvariant(Input))
+    ToInvariance = 1u;
+  else if (PHINode *IncPhi = dyn_cast<PHINode>(Input)) {
+    // Only consider Phis in header block.
+    if (IncPhi->getParent() != L->getHeader())
+      return InfiniteIterationsToInvariance;
+    // If the input becomes an invariant after X iterations, then our Phi
+    // becomes an invariant after X + 1 iterations.
+    unsigned InputToInvariance = calculateIterationsToInvariance(
+        IncPhi, L, BackEdge, IterationsToInvariance);
+    if (InputToInvariance != InfiniteIterationsToInvariance)
+      ToInvariance = InputToInvariance + 1u;
+  }
+
+  // If we found that this Phi lies in an invariant chain, update the map.
+  if (ToInvariance != InfiniteIterationsToInvariance)
+    IterationsToInvariance[Phi] = ToInvariance;
+  return ToInvariance;
+}
+
 // Return the number of iterations we want to peel off.
 void llvm::computePeelCount(Loop *L, unsigned LoopSize,
                             TargetTransformInfo::UnrollingPreferences &UP,
                             unsigned &TripCount) {
+  assert(LoopSize > 0 && "Zero loop size is not allowed!");
   UP.PeelCount = 0;
   if (!canPeel(L))
     return;
 
   // Only try to peel innermost loops.
   if (!L->empty())
     return;
 
-  // Try to find a Phi node that has the same loop invariant as an input from
-  // its only back edge. If there is such Phi, peeling 1 iteration from the
-  // loop is profitable, because starting from 2nd iteration we will have an
-  // invariant instead of this Phi.
+  // Here we try to get rid of Phis which become invariants after 1, 2, ..., N
+  // iterations of the loop. For this we compute the number for iterations after
+  // which every Phi is guaranteed to become an invariant, and try to peel the
+  // maximum number of iterations among these values, thus turning all those
+  // Phis into invariants.
   // First, check that we can peel at least one iteration.
   if (2 * LoopSize <= UP.Threshold && UnrollPeelMaxCount > 0) {
+    // Store the pre-calculated values here.
+    SmallDenseMap<PHINode *, unsigned> IterationsToInvariance;
+    // Now go through all Phis to calculate their the number of iterations they
+    // need to become invariants.
+    unsigned DesiredPeelCount = 0;
     BasicBlock *BackEdge = L->getLoopLatch();
     assert(BackEdge && "Loop is not in simplified form?");
-    BasicBlock *Header = L->getHeader();
-    // Iterate over Phis to find one with invariant input on back edge.
-    bool FoundCandidate = false;
-    PHINode *Phi;
-    for (auto BI = Header->begin(); isa<PHINode>(&*BI); ++BI) {
-      Phi = cast<PHINode>(&*BI);
-      Value *Input = Phi->getIncomingValueForBlock(BackEdge);
-      if (L->isLoopInvariant(Input)) {
-        FoundCandidate = true;
-        break;
-      }
-    }
-    if (FoundCandidate) {
-      DEBUG(dbgs() << "Peel one iteration to get rid of " << *Phi
-                   << " because starting from 2nd iteration it is always"
-                   << " an invariant\n");
-      UP.PeelCount = 1;
+    for (auto BI = L->getHeader()->begin(); isa<PHINode>(&*BI); ++BI) {
+      PHINode *Phi = cast<PHINode>(&*BI);
+      unsigned ToInvariance = calculateIterationsToInvariance(
+          Phi, L, BackEdge, IterationsToInvariance);
+      if (ToInvariance != InfiniteIterationsToInvariance)
+        DesiredPeelCount = std::max(DesiredPeelCount, ToInvariance);
+    }
+    if (DesiredPeelCount > 0) {
+      // Pay respect to limitations implied by loop size and the max peel count.
+      unsigned MaxPeelCount = UnrollPeelMaxCount;
+      MaxPeelCount = std::min(MaxPeelCount, UP.Threshold / LoopSize - 1);
+      DesiredPeelCount = std::min(DesiredPeelCount, MaxPeelCount);
+      // Consider max peel count limitation.
+      assert(DesiredPeelCount > 0 && "Wrong loop size estimation?");
+      DEBUG(dbgs() << "Peel " << DesiredPeelCount << " iteration(s) to turn"
+                   << " some Phis into invariants.\n");
+      UP.PeelCount = DesiredPeelCount;
       return;
     }
   }
 
   // Bail if we know the statically calculated trip count.
   // In this case we rather prefer partial unrolling.
   if (TripCount)
     return;

llvm/trunk/test/Transforms/LoopUnroll/peel-loop-not-forced.ll

-; RUN: opt < %s -S -loop-unroll -unroll-threshold=8 | FileCheck %s
+; RUN: opt < %s -S -loop-unroll -unroll-threshold=30 | FileCheck %s
 
 define i32 @invariant_backedge_1(i32 %a, i32 %b) {
 ; CHECK-LABEL: @invariant_backedge_1
 ; CHECK-NOT:     %plus = phi
 ; CHECK:       loop.peel:
 ; CHECK:       loop:
 ; CHECK:         %i = phi
 ; CHECK:         %sum = phi
loop:
   %cmp = icmp slt i32 %i, 1000
 
   br i1 %cmp, label %loop, label %exit
 
 exit:
   ret i32 %sum
 }
 
-; Peeling should fail due to method size.
 define i32 @invariant_backedge_2(i32 %a, i32 %b) {
+; This loop should be peeled twice because it has a Phi which becomes invariant
+; starting from 3rd iteration.
 ; CHECK-LABEL: @invariant_backedge_2
-; CHECK-NOT:   loop.peel:
+; CHECK:       loop.peel{{.*}}:
+; CHECK:       loop.peel{{.*}}:
+; CHECK:         %i = phi
+; CHECK:         %sum = phi
+; CHECK-NOT:     %half.inv = phi
+; CHECK-NOT:     %plus = phi
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
+  %sum = phi i32 [ 0, %entry ], [ %incsum, %loop ]
+  %half.inv = phi i32 [ %a, %entry ], [ %b, %loop ]
+  %plus = phi i32 [ %a, %entry ], [ %half.inv, %loop ]
+
+  %incsum = add i32 %sum, %plus
+  %inc = add i32 %i, 1
+  %cmp = icmp slt i32 %i, 1000
+
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret i32 %sum
+}
+
+define i32 @invariant_backedge_3(i32 %a, i32 %b) {
+; This loop should be peeled thrice because it has a Phi which becomes invariant
+; starting from 4th iteration.
+; CHECK-LABEL: @invariant_backedge_3
+; CHECK:       loop.peel{{.*}}:
+; CHECK:       loop.peel{{.*}}:
+; CHECK:       loop.peel{{.*}}:
+; CHECK:         %i = phi
+; CHECK:         %sum = phi
+; CHECK-NOT:     %half.inv = phi
+; CHECK-NOT:     %half.inv.2 = phi
+; CHECK-NOT:     %plus = phi
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
+  %sum = phi i32 [ 0, %entry ], [ %incsum, %loop ]
+  %half.inv = phi i32 [ %a, %entry ], [ %b, %loop ]
+  %half.inv.2 = phi i32 [ %a, %entry ], [ %half.inv, %loop ]
+  %plus = phi i32 [ %a, %entry ], [ %half.inv.2, %loop ]
+
+  %incsum = add i32 %sum, %plus
+  %inc = add i32 %i, 1
+  %cmp = icmp slt i32 %i, 1000
+
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret i32 %sum
+}
+
+define i32 @invariant_backedge_limited_by_size(i32 %a, i32 %b) {
+; This loop should normally be peeled thrice because it has a Phi which becomes
+; invariant starting from 4th iteration, but the size of the loop only allows
+; us to peel twice because we are restricted to 30 instructions in resulting
+; code. Thus, %plus Phi node should stay in loop even despite its backedge
+; input is an invariant.
+; CHECK-LABEL: @invariant_backedge_limited_by_size
+; CHECK:       loop.peel{{.*}}:
+; CHECK:       loop.peel{{.*}}:
+; CHECK:         %i = phi
+; CHECK:         %sum = phi
+; CHECK:         %plus = phi i32 [ %a, {{.*}} ], [ %b, %loop ]
+; CHECK-NOT:     %half.inv = phi
+; CHECK-NOT:     %half.inv.2 = phi
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
+  %sum = phi i32 [ 0, %entry ], [ %incsum, %loop ]
+  %half.inv = phi i32 [ %a, %entry ], [ %b, %loop ]
+  %half.inv.2 = phi i32 [ %a, %entry ], [ %half.inv, %loop ]
+  %plus = phi i32 [ %a, %entry ], [ %half.inv.2, %loop ]
+
+  %incsum = add i32 %sum, %plus
+  %inc = add i32 %i, 1
+  %cmp = icmp slt i32 %i, 1000
+
+  %incsum2 = add i32 %incsum, %plus
+  %incsum3 = add i32 %incsum, %plus
+  %incsum4 = add i32 %incsum, %plus
+  %incsum5 = add i32 %incsum, %plus
+  %incsum6 = add i32 %incsum, %plus
+  %incsum7 = add i32 %incsum, %plus
+
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret i32 %sum
+}
+
+; Peeling should fail due to method size.
+define i32 @invariant_backedge_negative(i32 %a, i32 %b) {
+; CHECK-LABEL: @invariant_backedge_negative
+; CHECK-NOT:   loop.peel{{.*}}:
 ; CHECK:       loop:
 ; CHECK:         %i = phi
 ; CHECK:         %sum = phi
 ; CHECK:         %plus = phi
 entry:
   br label %loop
 
 loop:
   %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
   %sum = phi i32 [ 0, %entry ], [ %incsum2, %loop ]
   %plus = phi i32 [ %a, %entry ], [ %b, %loop ]
 
   %incsum = add i32 %sum, %plus
   %incsum2 = add i32 %incsum, %plus
+  %incsum3 = add i32 %incsum, %plus
+  %incsum4 = add i32 %incsum, %plus
+  %incsum5 = add i32 %incsum, %plus
+  %incsum6 = add i32 %incsum, %plus
+  %incsum7 = add i32 %incsum, %plus
+  %incsum8 = add i32 %incsum, %plus
+  %incsum9 = add i32 %incsum, %plus
+  %incsum10 = add i32 %incsum, %plus
+  %incsum11 = add i32 %incsum, %plus
+  %incsum12 = add i32 %incsum, %plus
+  %incsum13 = add i32 %incsum, %plus
+  %incsum14 = add i32 %incsum, %plus
+  %incsum15 = add i32 %incsum, %plus
+  %inc = add i32 %i, 1
+  %cmp = icmp slt i32 %i, 1000
+
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  ret i32 %sum
+}
+
+define i32 @cycled_phis(i32 %a, i32 %b) {
+; Make sure that we do not crash working with cycled Phis and don't peel it.
+; TODO: Actually this loop should be partially unrolled with factor 2.
+; CHECK-LABEL: @cycled_phis
+; CHECK-NOT:   loop.peel{{.*}}:
+; CHECK:       loop:
+; CHECK:         %i = phi
+; CHECK:         %phi.a = phi
+; CHECK:         %phi.b = phi
+; CHECK:         %sum = phi
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
+  %phi.a = phi i32 [ %a, %entry ], [ %phi.b, %loop ]
+  %phi.b = phi i32 [ %b, %entry ], [ %phi.a, %loop ]
+  %sum = phi i32 [ 0, %entry], [ %incsum, %loop ]
+  %incsum = add i32 %sum, %phi.a
   %inc = add i32 %i, 1
   %cmp = icmp slt i32 %i, 1000
 
   br i1 %cmp, label %loop, label %exit
 
 exit:
   ret i32 %sum
 }