[LoopUnrolling] Peel loops with invariant backedge Phi input

[mkazantsev]

Summary

If a loop contains a Phi node which has an invariant input from back
edge, it is profitable to peel such loops (rather than unroll them) to
use the advantage that this Phi is always invariant starting from 2nd
iteration. After the 1st iteration is peeled, other optimizations can
potentially simplify calculations with this invariant.

[anna]

Hi Max,

I'm not very sure about the change in priority of the loop peeling. There can be performance repercussions for that change alone, and might be better to get another opinion on the priority change.
Could you perhaps have this change as a strict improvement for loop peeling in itself?

lib/Transforms/Scalar/LoopUnrollPass.cpp

788 ↗

(On Diff #89091)

Can this patch be separated into 2 parts, where the second part is this change to the priority?

lib/Transforms/Utils/LoopUnrollPeel.cpp

145

There is a method for this getNumBackEdges. I think you can just reuse this and have a check against 1.

147

can change this to for (auto *Pred: predecessors(Header))

[mkazantsev]

lib/Transforms/Scalar/LoopUnrollPass.cpp

788 ↗

(On Diff #89091)

The problem here is that without this change, if UP.Partial is not set, line 803 prevents us from peeling (we return from method with false), and if UP.Partial is set, we may unroll the loop rather than peel it and have a worse result. So the test in this patch doesn't work without this re-prioritizing. I can split it into 2 patches, but priority change will be a parent and peeling change will depend on it.

[mkazantsev]

Addressed comments.

[mkuper]

lib/Transforms/Utils/LoopUnrollPeel.cpp

144

Can you just use L->getLoopLatch()? Or is that different from what you're looking for?

[reames]

lib/Transforms/Utils/LoopUnrollPeel.cpp

149

Out of curiosity, why the complexity about finding the backedge? Wouldn't *all* of the inputs to the phi be loop invariant in the case you're interested in?

[mkazantsev]

lib/Transforms/Utils/LoopUnrollPeel.cpp

149

Off course, all inputs from preheaders will be invariant. Finding the backedge for header with n predecessors takes O(n) (it needs traversal over all predecessors with "contains" check in set that takes O(1). Acquiring its input also takes O(n) for every Phi, so total complexity being O(n*m) for m Phis.

If we just check all inputs for being invariants, it will also take O(n*m), but we will have positive results for loops with multiple back edges. The current implementation of peeling expects the loop to have 1 back edge, otherwise it will bail and we do unneded work with such loops.

[mkuper]

This generally LGTM, except for some nits. But please wait for @reames as well.

lib/Transforms/Utils/LoopUnrollPeel.cpp

83

Why the parens around (Phi = ...)?

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

1

We generally prefer to test passes in isolation. Can you please make this a test for loop-unroll only?

[reames]

fine by me

[efriedma]

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

95 ↗

(On Diff #90508)

Do we need to do some sort of threshold check here? At first glance, it looks like this will peel a loop of any size.

lib/Transforms/Utils/LoopUnrollPeel.cpp

void llvm::computePeelCount(Loop *L, unsigned LoopSize,
   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.
+  if (auto *BackEdge = L->getLoopLatch()) {
+    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(); (Phi = dyn_cast<PHINode>(&*BI)); ++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;
+      return;
+    }
+  }
+
   // Bail if we know the statically calculated trip count.
   // In this case we rather prefer partial unrolling.
   if (TripCount)
     return;
 
   // If the user provided a peel count, use that.
   bool UserPeelCount = UnrollForcePeelCount.getNumOccurrences() > 0;
   if (UserPeelCount) {
if (*PeelCount) {
       DEBUG(dbgs() << "Max peel cost: " << UP.Threshold << "\n");
     }
   }
 
   return;
 }
 
 /// \brief Update the branch weights of the latch of a peeled-off loop
 /// iteration.
 /// This sets the branch weights for the latch of the recently peeled off loop
 /// iteration correctly. 
 /// Our goal is to make sure that:
 /// a) The total weight of all the copies of the loop body is preserved.
 /// b) The total weight of the loop exit is preserved.
 /// c) The body weight is reasonably distributed between the peeled iterations.
 ///
 /// \param Header The copy of the header block that belongs to next iteration.
 /// \param LatchBR The copy of the latch branch that belongs to this iteration.
 /// \param IterNumber The serial number of the iteration that was just
 /// peeled off.
 /// \param AvgIters The average number of iterations we expect the loop to have.
 /// \param[in,out] PeeledHeaderWeight The total number of dynamic loop

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

+; RUN: opt < %s -S -loop-unroll -indvars -loop-deletion -simplifycfg -instcombine | FileCheck %s
+
+define i32 @invariant_backedge_1(i32 %a, i32 %b) {
+; CHECK-LABEL: @invariant_backedge_1
+; CHECK:       entry:
+; CHECK-NEXT:    %0 = mul i32 %b, 999
+; CHECK-NEXT:    %1 = add i32 %0, %a
+; CHECK-NEXT:    ret i32 %1
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
+  %sum = phi i32 [ 0, %entry ], [ %incsum, %loop ]
+  %plus = phi i32 [ %a, %entry ], [ %b, %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
+}