!fixup address latest comments, thanks

Author: fhahn

llvm/include/llvm/Transforms/Utils/LoopPeel.h

@@ -21,11 +21,15 @@ namespace llvm {
 
 bool canPeel(const Loop *L);
 
+/// Returns true if the last iteration of \p L can be peeled off. It makes sure
+/// the loop exit condition can be adjusted when peeling and that the loop
+/// executes at least 2 iterations.
+bool canPeelLastIteration(const Loop &L, ScalarEvolution &SE);
+
 /// VMap is the value-map that maps instructions from the original loop to
 /// instructions in the last peeled-off iteration. If \p PeelLast is true, peel
-/// off the last \p PeelCount iterations from \p L. In that case, the caller has
-/// to make sure that the exit condition can be adjusted when peeling and that
-/// the loop executes at least 2 iterations.
+/// off the last \p PeelCount iterations from \p L (canPeelLastIteration must be
+/// true for \p L), otherwise peel off the first \p PeelCount iterations.
 bool peelLoop(Loop *L, unsigned PeelCount, bool PeelLast, LoopInfo *LI,
               ScalarEvolution *SE, DominatorTree &DT, AssumptionCache *AC,
               bool PreserveLCSSA, ValueToValueMapTy &VMap);

llvm/lib/Transforms/Utils/LoopPeel.cpp

@@ -326,13 +326,7 @@ static unsigned peelToTurnInvariantLoadsDerefencebale(Loop &L,
   return 0;
 }
 
-/// Returns true if the last iteration can be peeled off and the condition (Pred
-/// LeftAR, RightSCEV) is known at the last iteration and the inverse condition
-/// is known at the second-to-last. This function also has to make sure the loop
-/// exit condition can be adjusted when peeling and that the loop executes at
-/// least 2 iterations.
-static bool canPeelLastIteration(Loop &L, const SCEVAddRecExpr *LeftAR,
-                                 const SCEV *RightSCEV, ScalarEvolution &SE) {
+bool llvm::canPeelLastIteration(const Loop &L, ScalarEvolution &SE) {
   const SCEV *BTC = SE.getBackedgeTakenCount(&L);
   Value *Inc;
   CmpPredicate Pred;
@@ -351,16 +345,27 @@ static bool canPeelLastIteration(Loop &L, const SCEVAddRecExpr *LeftAR,
   // * the exit condition must be a NE/EQ compare of an induction with step
   // of 1.
   BasicBlock *Latch = L.getLoopLatch();
-  if (Latch != L.getExitingBlock() ||
-      !match(Latch->getTerminator(),
-             m_Br(m_ICmp(Pred, m_Value(Inc), m_Value()), m_BasicBlock(Succ1),
-                  m_BasicBlock(Succ2))) ||
-      !((Pred == CmpInst::ICMP_EQ && Succ2 == L.getHeader()) ||
-        (Pred == CmpInst::ICMP_NE && Succ1 == L.getHeader())) ||
-      !isa<SCEVAddRecExpr>(SE.getSCEV(Inc)) ||
-      !cast<SCEVAddRecExpr>(SE.getSCEV(Inc))->getStepRecurrence(SE)->isOne())
+  return Latch == L.getExitingBlock() &&
+         match(Latch->getTerminator(),
+               m_Br(m_ICmp(Pred, m_Value(Inc), m_Value()), m_BasicBlock(Succ1),
+                    m_BasicBlock(Succ2))) &&
+         ((Pred == CmpInst::ICMP_EQ && Succ2 == L.getHeader()) ||
+          (Pred == CmpInst::ICMP_NE && Succ1 == L.getHeader())) &&
+         isa<SCEVAddRecExpr>(SE.getSCEV(Inc)) &&
+         cast<SCEVAddRecExpr>(SE.getSCEV(Inc))->getStepRecurrence(SE)->isOne();
+}
+
+/// Returns true if the last iteration can be peeled off and the condition (Pred
+/// LeftAR, RightSCEV) is known at the last iteration and the inverse condition
+/// is known at the second-to-last.
+static bool shouldPeelLastIteration(Loop &L, CmpPredicate Pred,
+                                    const SCEVAddRecExpr *LeftAR,
+                                    const SCEV *RightSCEV,
+                                    ScalarEvolution &SE) {
+  if (!canPeelLastIteration(L, SE))
     return false;
 
+  const SCEV *BTC = SE.getBackedgeTakenCount(&L);
   const SCEV *ValAtLastIter =
       SE.applyLoopGuards(LeftAR->evaluateAtIteration(BTC, SE), &L);
   const SCEV *ValAtSecondToLastIter = LeftAR->evaluateAtIteration(
@@ -470,7 +475,7 @@ static int countToEliminateCompares(Loop &L, unsigned MaxPeelCount,
     const SCEV *Step = LeftAR->getStepRecurrence(SE);
     if (!PeelWhilePredicateIsKnown(NewPeelCount, IterVal, RightSCEV, Step,
                                    Pred)) {
-      if (canPeelLastIteration(L, LeftAR, RightSCEV, SE))
+      if (shouldPeelLastIteration(L, Pred, LeftAR, RightSCEV, SE))
         DesiredPeelCount = -1;
       return;
     }
@@ -987,8 +992,9 @@ bool llvm::peelLoop(Loop *L, unsigned PeelCount, bool PeelLast, LoopInfo *LI,
                     bool PreserveLCSSA, ValueToValueMapTy &LVMap) {
   assert(PeelCount > 0 && "Attempt to peel out zero iterations?");
   assert(canPeel(L) && "Attempt to peel a loop which is not peelable?");
-  assert((!PeelLast || PeelCount == 1) &&
-         "can only peel off a single iteration from the end for now");
+  assert((!PeelLast || (canPeelLastIteration(*L, *SE) && PeelCount == 1)) &&
+         "when peeling the last iteration, the loop must be supported and can "
+         "only peel a single iteration");
 
   LoopBlocksDFS LoopBlocks(L);
   LoopBlocks.perform(LI);

llvm/test/Transforms/LoopUnroll/peel-last-iteration.ll

@@ -85,6 +85,7 @@ define i64 @peel_single_block_loop_iv_step_1_eq_pred() {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
 ; CHECK-NEXT:    br label %[[LOOP:.*]]
 ; CHECK:       [[LOOP]]:
+<<<<<<< HEAD
 ; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[IV]], 63
 ; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i32 10, i32 20
@@ -93,7 +94,38 @@ define i64 @peel_single_block_loop_iv_step_1_eq_pred() {
 ; CHECK-NEXT:    [[EC:%.*]] = icmp eq i64 [[IV_NEXT]], 64
 ; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT:.*]], label %[[LOOP]]
 ; CHECK:       [[EXIT]]:
+||||||| parent of 4da8bc5afc5b (!fixup address latest comments, thanks)
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[CMP18_NOT:%.*]] = icmp eq i64 [[IV]], 63
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP18_NOT]], i32 10, i32 20
+; CHECK-NEXT:    call void @foo(i32 [[COND]])
+; CHECK-NEXT:    [[IV_NEXT]] = add i64 [[IV]], 1
+; CHECK-NEXT:    [[EC:%.*]] = icmp eq i64 [[IV_NEXT]], 64
+; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT:.*]], label %[[LOOP]]
+; CHECK:       [[EXIT]]:
+=======
+; CHECK-NEXT:    [[IV1:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[IV:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    call void @foo(i32 20)
+; CHECK-NEXT:    [[IV]] = add nuw nsw i64 [[IV1]], 1
+; CHECK-NEXT:    [[CMP18_NOT:%.*]] = icmp eq i64 [[IV]], 63
+; CHECK-NEXT:    br i1 [[CMP18_NOT]], label %[[EXIT_PEEL_BEGIN:.*]], label %[[LOOP]], !llvm.loop [[LOOP2:![0-9]+]]
+; CHECK:       [[EXIT_PEEL_BEGIN]]:
+>>>>>>> 4da8bc5afc5b (!fixup address latest comments, thanks)
 ; CHECK-NEXT:    [[IV_LCSSA:%.*]] = phi i64 [ [[IV]], %[[LOOP]] ]
+; CHECK-NEXT:    [[IV_LCSSA1:%.*]] = phi i64 [ [[IV1]], %[[LOOP]] ]
+; CHECK-NEXT:    br label %[[LOOP_PEEL:.*]]
+; CHECK:       [[LOOP_PEEL]]:
+; CHECK-NEXT:    [[CMP_PEEL:%.*]] = icmp eq i64 [[IV_LCSSA]], 63
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP_PEEL]], i32 10, i32 20
+; CHECK-NEXT:    call void @foo(i32 [[COND]])
+; CHECK-NEXT:    [[IV_NEXT:%.*]] = add i64 [[IV_LCSSA]], 1
+; CHECK-NEXT:    [[EC:%.*]] = icmp eq i64 [[IV_NEXT]], 64
+; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT_PEEL_NEXT:.*]], label %[[EXIT_PEEL_NEXT]]
+; CHECK:       [[EXIT_PEEL_NEXT]]:
+; CHECK-NEXT:    br label %[[LOOP_PEEL_NEXT:.*]]
+; CHECK:       [[LOOP_PEEL_NEXT]]:
+; CHECK-NEXT:    br label %[[EXIT:.*]]
+; CHECK:       [[EXIT]]:
 ; CHECK-NEXT:    ret i64 [[IV_LCSSA]]
 ;
 entry:
@@ -152,7 +184,7 @@ define i64 @peel_single_block_loop_iv_step_1_nested_loop() {
 ; CHECK-NEXT:    [[OUTER_IV:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[OUTER_IV_NEXT:%.*]], %[[OUTER_LATCH:.*]] ]
 ; CHECK-NEXT:    br label %[[LOOP:.*]]
 ; CHECK:       [[LOOP]]:
-;
+<<<;
 entry:
   br label %outer.header
 
@@ -435,7 +467,7 @@ define i32 @peel_loop_with_branch_and_phi_uses(ptr %x, i1 %c) {
 ; CHECK-NEXT:    [[ADD1]] = add nsw i32 [[L1]], [[RED1]]
 ; CHECK-NEXT:    [[IV_NEXT1]] = add nuw nsw i32 [[IV1]], 1
 ; CHECK-NEXT:    [[EC1:%.*]] = icmp ne i32 [[IV_NEXT1]], 99
-; CHECK-NEXT:    br i1 [[EC1]], label %[[LOOP_HEADER]], label %[[LOOPEXIT_PEEL_BEGIN:.*]], !llvm.loop [[LOOP4:![0-9]+]]
+; CHECK-NEXT:    br i1 [[EC1]], label %[[LOOP_HEADER]], label %[[LOOPEXIT_PEEL_BEGIN:.*]], !llvm.loop [[LOOP5:![0-9]+]]
 ; CHECK:       [[LOOPEXIT_PEEL_BEGIN]]:
 ; CHECK-NEXT:    [[RED:%.*]] = phi i32 [ [[ADD1]], %[[LOOP_LATCH]] ]
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[IV_NEXT1]], %[[LOOP_LATCH]] ]