[Loop Peeling] Enable peeling for loops with multiple exits

[skatkov]

Summary

This CL enables peeling of the loop with multiple exits where
one exit should be from latch and others are basic blocks with
call to deopt.

The peeling is enabled under the flag which is false by default.

[reames]

lib/Transforms/Utils/LoopUnrollPeel.cpp

65 ↗

(On Diff #207021)

Submit disabled, then enable in future commit?

And maybe you want a three way choice: none, deopt-only, all?

84 ↗

(On Diff #207021)

This looks like a useful helper function which could be factored out and reused, for example by canProfitablyUnrollMultiExitLoop.

Maybe: bool allNonLatchExitsDeoptimize or allNonLatchExitsPredicatedNeverTaken?

lib/Transforms/Utils/LoopUtils.cpp

626 ↗

(On Diff #207021)

A suggestion here. Extend the existing function to allow multiple exits, but only if those exits have 0 taken probability. Return None for any exit with non-zero exit.

(I'm trying to avoid the need for the new function.)

[skatkov]

lib/Transforms/Utils/LoopUnrollPeel.cpp

65 ↗

(On Diff #207021)

Sure, will change to false.

I do not support the non deopt-only now in terms of update of branch weights. So introducing the explicit difference between deopt-only and all is not what I think make sense.

lib/Transforms/Utils/LoopUtils.cpp

626 ↗

(On Diff #207021)

Unfortunately, zero probability is not what is used in llvm. Even BranchProbability analysis will set smallest but greater than zero probability for edge coming to block with unreachable terminator.

So if want to add a complexity to this method we should introduce some threshold below which we consider the exit has '0' probability.

Do you still see the value in it?
Please note that getLoopEstimatedTripCount is widely used in code base.

[skatkov]

Rebase

[fhahn]

lib/Transforms/Utils/LoopUnrollPeel.cpp

65 ↗

(On Diff #207021)

I think it would be good for the naming to somehow reflect that we do not support any multi-exit loop, but just a very restricted set. The current naming/description sounds more general than the actual implementation.

[reames]

LGTM w/two requested changes before or after commit, no further review needed.

1. allNonLatchExitsDeopt should be separated into a utility file and reused to implement the same (existing) check in the LoopUnroll code to ensure there's no accidental divergence.

2. getEstimatedTripCount should return a answer for all loops it knows how to. After this change sequence, that will be single exit loops, and multiple exit loops where all non-latch exits are deopt calls. Existing callers should be updated to either a) add the single exit check or b) call a new shim function (getETCForSingleExitLoop) which contains the check.

[skatkov]

Handled Philip's comments.

I decided to avoid introduction of utility function for check that all side exits
end up with deopt. The current code with all_of looks short enough and clear
to understand. At the same time, the case in LoopUnrollruntime is more narrow
due to it supports only one side exits, so updating that case to use the same utility
function does not improve readability from my view.

Please take a look.

[reames]

still LGTM

[xbolva00]

llvm/lib/Transforms/Utils/LoopUnrollPeel.cpp

96

Nit: if (!Exists.empty())

llvm/include/llvm/Transforms/Utils/LoopUtils.h

 void addStringMetadataToLoop(Loop *TheLoop, const char *MDString,
                              unsigned V = 0);
 
 /// Get a loop's estimated trip count based on branch weight metadata.
 /// Returns 0 when the count is estimated to be 0, or None when a meaningful
 /// estimate can not be made.
 Optional<unsigned> getLoopEstimatedTripCount(Loop *L);
 
+/// Get a loop's estimated trip count based on branch weight metadata.
+/// Only Latch branch is considered.
+/// Returns 0 when the count is estimated to be 0, or None when a meaningful
+/// estimate can not be made.
+Optional<unsigned> getLoopEstimatedTripCountBasingOnLatchOnly(Loop *L);
+
 /// Check inner loop (L) backedge count is known to be invariant on all
 /// iterations of its outer loop. If the loop has no parent, this is trivially
 /// true.
 bool hasIterationCountInvariantInParent(Loop *L, ScalarEvolution &SE);
 
 /// Helper to consistently add the set of standard passes to a loop pass's \c
 /// AnalysisUsage.
 ///

llvm/lib/Transforms/Utils/LoopUnroll.cpp

LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   bool Peeled = false;
   if (ULO.PeelCount) {
     Peeled = peelLoop(L, ULO.PeelCount, LI, SE, DT, AC, PreserveLCSSA);
 
     // Successful peeling may result in a change in the loop preheader/trip
     // counts. If we later unroll the loop, we want these to be updated.
     if (Peeled) {
       // According to our guards and profitability checks the only
-      // meaningful exit should be latch block.
+      // meaningful exit should be latch block. Other exits go to deopt,
+      // so we do not worry about them.
       BasicBlock *ExitingBlock = L->getLoopLatch();
       assert(ExitingBlock && "Loop without exiting block?");
       assert(L->isLoopExiting(ExitingBlock) && "Latch is not exiting?");
       Preheader = L->getLoopPreheader();
       ULO.TripCount = SE->getSmallConstantTripCount(L, ExitingBlock);
       ULO.TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock);
     }
   }

llvm/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."));
 
+static cl::opt<bool> UnrollPeelMultiDeoptExit(
+    "unroll-peel-multi-deopt-exit", cl::init(false), cl::Hidden,
+    cl::desc("Allow peeling of loops with multiple deopt exits."));
+
 // 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 =
     std::numeric_limits<unsigned>::max();
 
+// Returns true if all non latch exits ends up with a deopt.
+static bool allNonLatchExitsDeoptimize(
+    const BasicBlock *Latch,
+    SmallVectorImpl<std::pair<const BasicBlock *, const BasicBlock *> > &
+        ExitEdges) {
+  return all_of(ExitEdges,
+                [Latch](
+                    std::pair<const BasicBlock *, const BasicBlock *> Edge) {
+    return Latch == Edge.first || Edge.second->getTerminatingDeoptimizeCall();
+  });
+}
+
 // Check whether we are capable of peeling this loop.
 bool llvm::canPeel(Loop *L) {
   // Make sure the loop is in simplified form
   if (!L->isLoopSimplifyForm())
     return false;
 
+  if (UnrollPeelMultiDeoptExit) {
+    SmallVector<std::pair<const BasicBlock *, const BasicBlock *>, 4> ExitEdges;
+    L->getExitEdges(ExitEdges);
+
+    if (ExitEdges.size() > 1) {
+      // Latch's terminator is a conditional branch, Latch is exiting and
+      // all non Latch exits ends up with deoptimize.
+      const BasicBlock *Latch = L->getLoopLatch();
+      const BranchInst *T = dyn_cast<BranchInst>(Latch->getTerminator());
+      return T && T->isConditional() && L->isLoopExiting(Latch) &&
+             allNonLatchExitsDeoptimize(Latch, ExitEdges);
+    }
+  }
+
   // 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 (TripCount)
     return;
 
   // If we don't know the trip count, but have reason to believe the average
   // trip count is low, peeling should be beneficial, since we will usually
   // hit the peeled section.
   // We only do this in the presence of profile information, since otherwise
   // our estimates of the trip count are not reliable enough.
   if (L->getHeader()->getParent()->hasProfileData()) {
-    Optional<unsigned> PeelCount = getLoopEstimatedTripCount(L);
+    Optional<unsigned> PeelCount =
+        getLoopEstimatedTripCountBasingOnLatchOnly(L);
     if (!PeelCount)
       return;
 
     LLVM_DEBUG(dbgs() << "Profile-based estimated trip count is " << *PeelCount
                       << "\n");
 
     if (*PeelCount) {
       if ((*PeelCount <= UnrollPeelMaxCount) &&

llvm/lib/Transforms/Utils/LoopUtils.cpp

if (LI) {
     LI->erase(L);
   }
 }
 
 Optional<unsigned> llvm::getLoopEstimatedTripCount(Loop *L) {
   // Only support loops with a unique exiting block, and a latch.
   if (!L->getExitingBlock())
     return None;
+  return getLoopEstimatedTripCountBasingOnLatchOnly(L);
+}
+
+Optional<unsigned> llvm::getLoopEstimatedTripCountBasingOnLatchOnly(Loop *L) {
 
   // Get the branch weights for the loop's backedge.
   BranchInst *LatchBR =
       dyn_cast<BranchInst>(L->getLoopLatch()->getTerminator());
   if (!LatchBR || LatchBR->getNumSuccessors() != 2)
     return None;
 
   assert((LatchBR->getSuccessor(0) == L->getHeader() ||

llvm/test/Transforms/LoopUnroll/peel-loop-pgo-deopt.ll

+; RUN: opt < %s -S -debug-only=loop-unroll -loop-unroll -unroll-runtime -unroll-peel-multi-deopt-exit 2>&1 | FileCheck %s
+; RUN: opt < %s -S -debug-only=loop-unroll -unroll-peel-multi-deopt-exit -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s
+
+; Make sure we use the profile information correctly to peel-off 3 iterations
+; from the loop, and update the branch weights for the peeled loop properly.
+
+; CHECK: Loop Unroll: F[basic]
+; CHECK: PEELING loop %for.body with iteration count 3!
+
+; CHECK-LABEL: @basic
+; CHECK: br i1 %{{.*}}, label %[[NEXT0:.*]], label %for.cond.for.end_crit_edge, !prof !16
+; CHECK: [[NEXT0]]:
+; CHECK: br i1 %{{.*}}, label %[[NEXT1:.*]], label %for.cond.for.end_crit_edge, !prof !17
+; CHECK: [[NEXT1]]:
+; CHECK: br i1 %{{.*}}, label %[[NEXT2:.*]], label %for.cond.for.end_crit_edge, !prof !18
+; CHECK: [[NEXT2]]:
+; CHECK: br i1 %{{.*}}, label %for.body, label %{{.*}}, !prof !19
+
+define i32 @basic(i32* %p, i32 %k, i1 %c) #0 !prof !15 {
+entry:
+  %cmp3 = icmp slt i32 0, %k
+  br i1 %cmp3, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.body
+  %i.05 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %continue ]
+  %p.addr.04 = phi i32* [ %p, %for.body.lr.ph ], [ %incdec.ptr, %continue ]
+  %incdec.ptr = getelementptr inbounds i32, i32* %p.addr.04, i32 1
+  store i32 %i.05, i32* %p.addr.04, align 4
+  %inc = add nsw i32 %i.05, 1
+  %cmp = icmp slt i32 %inc, %k
+  br i1 %c, label %continue, label %side_exit, !prof !17
+
+continue:
+  br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge, !prof !16
+
+for.cond.for.end_crit_edge:                       ; preds = %for.body
+  br label %for.end
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry
+  %res = phi i32 [ 0, %entry ], [ %inc, %for.cond.for.end_crit_edge ]
+  ret i32 %res
+
+side_exit:
+  %rval = call i32(...) @llvm.experimental.deoptimize.i32() [ "deopt"(i32 %inc) ]
+  ret i32 %rval
+}
+
+declare i32 @llvm.experimental.deoptimize.i32(...)
+
+attributes #0 = { nounwind }
+attributes #1 = { nounwind optsize }
+
+!llvm.module.flags = !{!1}
+
+!1 = !{i32 1, !"ProfileSummary", !2}
+!2 = !{!3, !4, !5, !6, !7, !8, !9, !10}
+!3 = !{!"ProfileFormat", !"InstrProf"}
+!4 = !{!"TotalCount", i64 10}
+!5 = !{!"MaxCount", i64 3}
+!6 = !{!"MaxInternalCount", i64 1}
+!7 = !{!"MaxFunctionCount", i64 3}
+!8 = !{!"NumCounts", i64 2}
+!9 = !{!"NumFunctions", i64 2}
+!10 = !{!"DetailedSummary", !11}
+!11 = !{!12, !13, !14}
+!12 = !{i32 10000, i64 3, i32 2}
+!13 = !{i32 999000, i64 1, i32 10}
+!14 = !{i32 999999, i64 1, i32 10}
+!15 = !{!"function_entry_count", i64 1}
+!16 = !{!"branch_weights", i32 3001, i32 1001}
+!17 = !{!"branch_weights", i32 1, i32 0}
+
+;CHECK: !16 = !{!"branch_weights", i32 900, i32 101}
+;CHECK: !17 = !{!"branch_weights", i32 540, i32 360}
+;CHECK: !18 = !{!"branch_weights", i32 162, i32 378}
+;CHECK: !19 = !{!"branch_weights", i32 1399, i32 162}
+