[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

84

Nit: if (!Exists.empty())

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();
 
 // 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<BasicBlock *, 4> Exits;
+    L->getUniqueNonLatchExitBlocks(Exits);
+
+    if (Exits.size()) {
+      // 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) &&
+             all_of(Exits, [](const BasicBlock *BB) {
+               return BB->getTerminatingDeoptimizeCall();
+             });
+    }
+  }
+
   // 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.

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}
+