[LoopPeeling] Add flag to disable support for peeling loops with non-latch exits

[asbirlea]

Summary

Add a flag to allow disabling the changes in
https://reviews.llvm.org/D134803.

[tra]

Thank you, Alina!

llvm/lib/Transforms/Utils/LoopPeel.cpp

905

Style nit: no need for {} around if/else bodies here.

[bgraur]

Thanks a lot Alina!

[nikic]

llvm/lib/Transforms/Utils/LoopPeel.cpp

112

Is it possible to reduce these conditions to the ones relevant to the affected test case? Like, I doubt that the check for BranchInst is relevant, and it would be good to know whether for the problematic case having a non-latch exit is the issue, or the structure of other exits.

516

Why is all this code needed? I would have expected the change to be just in canPeel(). For the cases where canPeel() previously returned true, the new implementation should behave the same apart from minor differences in branch weights.

Does that mean that you see issues if you limit peeling, but use the new branch weight updating code? That would be quite surprising, and would point towards something there being broken.

[nikic]

Generally fine with a temporary option, as long as the code duplication introduced here is not actually necessary.

[asbirlea]

Simplified to only the changes in canPeel.

[asbirlea]

Simplify conditions in canPeel.

llvm/lib/Transforms/Utils/LoopPeel.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> DisableAdvancedPeeling(
+    "disable-advanced-peeling", cl::init(false), cl::Hidden,
+    cl::desc(
+        "Disable advance peeling. Issues for convergent targets (D134803)."));
+
 static const char *PeeledCountMetaData = "llvm.loop.peeled.count";
 
 // Check whether we are capable of peeling this loop.
 bool llvm::canPeel(Loop *L) {
   // Make sure the loop is in simplified form
-  return L->isLoopSimplifyForm();
+  if (!L->isLoopSimplifyForm())
+    return false;
+  if (!DisableAdvancedPeeling)
+    return true;
+
+  SmallVector<BasicBlock *, 4> Exits;
+  L->getUniqueNonLatchExitBlocks(Exits);
+  // The latch must either be the only exiting block or all non-latch exit
+  // blocks have either a deopt or unreachable terminator or compose a chain of
+  // blocks where the last one is either deopt or unreachable terminated. Both
+  // deopt and unreachable terminators are a strong indication they are not
+  // taken. Note that this is a profitability check, not a legality check. Also
+  // note that LoopPeeling currently can only update the branch weights of latch
+  // blocks and branch weights to blocks with deopt or unreachable do not need
+  // updating.
+  return llvm::all_of(Exits, IsBlockFollowedByDeoptOrUnreachable);
 }
 
 // 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 Optional<unsigned> calculateIterationsToInvariance(
     PHINode *Phi, Loop *L, BasicBlock *BackEdge,
     SmallDenseMap<PHINode *, Optional<unsigned> > &IterationsToInvariance) {
   assert(Phi->getParent() == L->getHeader() &&
          "Non-loop Phi should not be checked for turning into invariant.");
Term->setMetadata(LLVMContext::MD_prof,
                     MDB.createBranchWeights(Info.Weights));
   for (auto [Idx, SubWeight] : enumerate(Info.SubWeights))
     if (SubWeight != 0)
       Info.Weights[Idx] = Info.Weights[Idx] > SubWeight
                               ? Info.Weights[Idx] - SubWeight
                               : 1;
 }
 
 /// Initialize the weights for all exiting blocks.
 static void initBranchWeights(DenseMap<Instruction *, WeightInfo> &WeightInfos,
                               Loop *L) {
   SmallVector<BasicBlock *> ExitingBlocks;
   L->getExitingBlocks(ExitingBlocks);
   for (BasicBlock *ExitingBlock : ExitingBlocks) {
     Instruction *Term = ExitingBlock->getTerminator();
     SmallVector<uint32_t> Weights;
     if (!extractBranchWeights(*Term, Weights))
for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     Instruction *LatchInst = dyn_cast<Instruction>(NewVal);
     if (LatchInst && L->contains(LatchInst))
       NewVal = LVMap[LatchInst];
 
     PHI->setIncomingValueForBlock(NewPreHeader, NewVal);
   }
 
   for (const auto &[Term, Info] : Weights)
     fixupBranchWeights(Term, Info);
 
   // Update Metadata for count of peeled off iterations.
   unsigned AlreadyPeeled = 0;
   if (auto Peeled = getOptionalIntLoopAttribute(L, PeeledCountMetaData))
     AlreadyPeeled = *Peeled;
   addStringMetadataToLoop(L, PeeledCountMetaData, AlreadyPeeled + PeelCount);
 
   if (Loop *ParentLoop = L->getParentLoop())

llvm/test/Transforms/LoopUnroll/peel-branch-weights.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
 ; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=2 2>&1 | FileCheck %s
+; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=2 -disable-advanced-peeling 2>&1 | FileCheck %s --check-prefix=DISABLEADV
 
 declare i32 @get.x()
 
 ; Test branch weight update for terminator with multiple fallthrough and
 ; multiple exit edges.
 define void @test() {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:  entry:

 ; CHECK-NEXT:    i32 2, label [[LOOP_EXIT_LOOPEXIT]]
 ; CHECK-NEXT:    ], !prof [[PROF2:![0-9]+]]
 ; CHECK:       loop.latch:
 ; CHECK-NEXT:    br label [[LOOP]], !llvm.loop [[LOOP3:![0-9]+]]
 ; CHECK:       loop.exit.loopexit:
 ; CHECK-NEXT:    br label [[LOOP_EXIT]]
 ; CHECK:       loop.exit:
 ; CHECK-NEXT:    ret void
-;
+
+; DISABLEADV-LABEL: @test()
+; DISABLEADV-NEXT: entry:
+; DISABLEADV-NEXT:  br label %loop
+; DISABLEADV: loop
+; DISABLEADV-NEXT:  %x = call i32 @get.x()
+; DISABLEADV-NEXT:  switch i32 %x, label %loop.latch [
+; DISABLEADV-NEXT:    i32 0, label %loop.latch
+; DISABLEADV-NEXT:    i32 1, label %loop.exit
+; DISABLEADV-NEXT:    i32 2, label %loop.exit
+; DISABLEADV-NEXT:  ], !prof !0
+; DISABLEADV: loop.latch:
+; DISABLEADV-NEXT:  br label %loop
+; DISABLEADV: loop.exit:
+; DISABLEADV-NEXT:  ret void
+
 entry:
   br label %loop
 
 loop:
   %x = call i32 @get.x()
   switch i32 %x, label %loop.latch [
   i32 0, label %loop.latch
   i32 1, label %loop.exit