[LoopUnroll] Peel iterations based on select conditions

[caojoshua]

Summary

This also allows us to peel loops with a select:

for (int i = 0; i <= N; ++i);
  f3(i == 0 ? a : b); // select instruction

into:

f3(a); // peel one iteration
for (int i = 1; i <= N; ++i)
  f3(b);

[caojoshua]

More details examples in commit msg

[caojoshua]

update commit msg

[nikic]

Please restrict this patch to handle select only. Let's handle and/or in a followup.

llvm/lib/Transforms/Utils/LoopPeel.cpp

352

Why is this code needed now?

452

You probably meant m_And/m_Or here. m_CombineAnd/m_CombineOr will put &I into both LHS and RHS, which is certainly not what you intended.

[caojoshua]

llvm/lib/Transforms/Utils/LoopPeel.cpp

352

Latch branch conditions can peel the entire loop.

for (i = 0; i < 4; ++i) {
  foo();
  // The backedge branch, i < 4, is known for all iterations, and could make us peel the entire loop
}

When peeling only branch instructions, we avoid this by ignoring the branch condition of the latch block. But with this patch, we don't ignore and/or's that are used by the latch branch.

I will take out and/or's from this patch as recommended, but I would still like to keep this code block here. It's not breaking and it avoids unnecessary transformations.

[caojoshua]

Restrict patch to only handle selects. Handle and/or in another patch.

[nikic]

llvm/lib/Transforms/Utils/LoopPeel.cpp

352

The getZExtValue() here can assert if we the BECount is larger than 64 bits. You can use getLimitedValue() instead.

When peeling only branch instructions, we avoid this by ignoring the branch condition of the latch block. But with this patch, we don't ignore and/or's that are used by the latch branch.

FWIW, I think that's a problem in your and/or handling: You shouldn't pick up random and/ors, but only those rooted at a br or select.

But anyway, I don't mind adding the check.

[caojoshua]

use getLimitedValue() instead of getZExtValue()

[caojoshua]

llvm/lib/Transforms/Utils/LoopPeel.cpp

352

Thanks, using getLimitedValue() in newest changes. I added a test for peeling large loops BECounts (>2^64) to precommit tests locally.

FWIW, I think that's a problem in your and/or handling: You shouldn't pick up random and/ors, but only those rooted at a br or select.

Benefit of checking random and/ors is that allows peeling nested and/ors easily. We can also peel for and/ors not used in branches/selects, for example, we might pass a boolean to a function (probably rare in practice though).

[nikic]

LGTM

llvm/lib/Transforms/Utils/LoopPeel.cpp

352

Benefit of checking random and/ors is that allows peeling nested and/ors easily. We can also peel for and/ors not used in branches/selects, for example, we might pass a boolean to a function (probably rare in practice though).

If we wanted to peel such cases, we would just look at all icmps, regardless of where they occur. But we specifically don't want to peel cases like that, because it's not useful. (It does not enable any meaningful follow-up simplification.)

llvm/lib/Transforms/Utils/LoopPeel.cpp

 //    else
 //      ..
 //   }
 static unsigned countToEliminateCompares(Loop &L, unsigned MaxPeelCount,
                                          ScalarEvolution &SE) {
   assert(L.isLoopSimplifyForm() && "Loop needs to be in loop simplify form");
   unsigned DesiredPeelCount = 0;
 
-  for (auto *BB : L.blocks()) {
-    auto *BI = dyn_cast<BranchInst>(BB->getTerminator());
-    if (!BI || BI->isUnconditional())
-      continue;
+  // Do not peel the entire loop.
+  const SCEV *BE = SE.getConstantMaxBackedgeTakenCount(&L);
+  if (const SCEVConstant *SC = dyn_cast<SCEVConstant>(BE))
+    MaxPeelCount =
+        std::min((unsigned)SC->getAPInt().getLimitedValue() - 1, MaxPeelCount);
 
-    // Ignore loop exit condition.
-    if (L.getLoopLatch() == BB)
-      continue;
+  auto ComputePeelCount = [&](Value *Condition) -> void {
+    if (!Condition->getType()->isIntegerTy())
+      return;
 
-    Value *Condition = BI->getCondition();
     Value *LeftVal, *RightVal;
     CmpInst::Predicate Pred;
     if (!match(Condition, m_ICmp(Pred, m_Value(LeftVal), m_Value(RightVal))))
-      continue;
+      return;
 
     const SCEV *LeftSCEV = SE.getSCEV(LeftVal);
     const SCEV *RightSCEV = SE.getSCEV(RightVal);
 
     // Do not consider predicates that are known to be true or false
     // independently of the loop iteration.
     if (SE.evaluatePredicate(Pred, LeftSCEV, RightSCEV))
-      continue;
+      return;
 
     // Check if we have a condition with one AddRec and one non AddRec
     // expression. Normalize LeftSCEV to be the AddRec.
     if (!isa<SCEVAddRecExpr>(LeftSCEV)) {
       if (isa<SCEVAddRecExpr>(RightSCEV)) {
         std::swap(LeftSCEV, RightSCEV);
         Pred = ICmpInst::getSwappedPredicate(Pred);
       } else
-        continue;
+        return;
     }
 
     const SCEVAddRecExpr *LeftAR = cast<SCEVAddRecExpr>(LeftSCEV);
 
     // Avoid huge SCEV computations in the loop below, make sure we only
     // consider AddRecs of the loop we are trying to peel.
     if (!LeftAR->isAffine() || LeftAR->getLoop() != &L)
-      continue;
+      return;
     if (!(ICmpInst::isEquality(Pred) && LeftAR->hasNoSelfWrap()) &&
         !SE.getMonotonicPredicateType(LeftAR, Pred))
-      continue;
+      return;
 
     // Check if extending the current DesiredPeelCount lets us evaluate Pred
     // or !Pred in the loop body statically.
     unsigned NewPeelCount = DesiredPeelCount;
 
     const SCEV *IterVal = LeftAR->evaluateAtIteration(
         SE.getConstant(LeftSCEV->getType(), NewPeelCount), SE);
 
auto ComputePeelCount = [&](Value *Condition) -> void {
     while (CanPeelOneMoreIteration() &&
            SE.isKnownPredicate(Pred, IterVal, RightSCEV))
       PeelOneMoreIteration();
 
     // With *that* peel count, does the predicate !Pred become known in the
     // first iteration of the loop body after peeling?
     if (!SE.isKnownPredicate(ICmpInst::getInversePredicate(Pred), IterVal,
                              RightSCEV))
-      continue; // If not, give up.
+      return; // If not, give up.
 
     // However, for equality comparisons, that isn't always sufficient to
     // eliminate the comparsion in loop body, we may need to peel one more
     // iteration. See if that makes !Pred become unknown again.
     if (ICmpInst::isEquality(Pred) &&
         !SE.isKnownPredicate(ICmpInst::getInversePredicate(Pred), NextIterVal,
                              RightSCEV) &&
         !SE.isKnownPredicate(Pred, IterVal, RightSCEV) &&
         SE.isKnownPredicate(Pred, NextIterVal, RightSCEV)) {
       if (!CanPeelOneMoreIteration())
-        continue; // Need to peel one more iteration, but can't. Give up.
+        return; // Need to peel one more iteration, but can't. Give up.
       PeelOneMoreIteration(); // Great!
     }
 
     DesiredPeelCount = std::max(DesiredPeelCount, NewPeelCount);
+  };
+
+  for (BasicBlock *BB : L.blocks()) {
+    for (Instruction &I : *BB) {
+      if (SelectInst *SI = dyn_cast<SelectInst>(&I))
+        ComputePeelCount(SI->getCondition());
+    }
+
+    auto *BI = dyn_cast<BranchInst>(BB->getTerminator());
+    if (!BI || BI->isUnconditional())
+      continue;
+
+    // Ignore loop exit condition.
+    if (L.getLoopLatch() == BB)
+      continue;
+
+    ComputePeelCount(BI->getCondition());
   }
 
   return DesiredPeelCount;
 }
 
 /// This "heuristic" exactly matches implicit behavior which used to exist
 /// inside getLoopEstimatedTripCount.  It was added here to keep an
 /// improvement inside that API from causing peeling to become more aggressive.

llvm/test/Transforms/LoopUnroll/peel-loop-conditions.ll

 
 ; Test that we can peel based on a select's condition
 define void @test19(i32 %num, i32 %a, i32 %b) {
 ; CHECK-LABEL: @test19(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP5:%.*]] = icmp sgt i32 [[NUM:%.*]], 0
 ; CHECK-NEXT:    br i1 [[CMP5]], label [[FOR_BODY_PREHEADER:%.*]], label [[FOR_COND_CLEANUP:%.*]]
 ; CHECK:       for.body.preheader:
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL_BEGIN:%.*]]
+; CHECK:       for.body.peel.begin:
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL:%.*]]
+; CHECK:       for.body.peel:
+; CHECK-NEXT:    [[CMP1_PEEL:%.*]] = icmp eq i32 0, 0
+; CHECK-NEXT:    [[COND_PEEL:%.*]] = select i1 [[CMP1_PEEL]], i32 [[A:%.*]], i32 [[B:%.*]]
+; CHECK-NEXT:    tail call void @f3(i32 noundef [[COND_PEEL]])
+; CHECK-NEXT:    [[INC_PEEL:%.*]] = add nuw nsw i32 0, 1
+; CHECK-NEXT:    [[EXITCOND_NOT_PEEL:%.*]] = icmp eq i32 [[INC_PEEL]], [[NUM]]
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT_PEEL]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], label [[FOR_BODY_PEEL_NEXT:%.*]]
+; CHECK:       for.body.peel.next:
+; CHECK-NEXT:    br label [[FOR_BODY_PEEL_NEXT1:%.*]]
+; CHECK:       for.body.peel.next1:
+; CHECK-NEXT:    br label [[FOR_BODY_PREHEADER_PEEL_NEWPH:%.*]]
+; CHECK:       for.body.preheader.peel.newph:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup.loopexit.loopexit:
+; CHECK-NEXT:    br label [[FOR_COND_CLEANUP_LOOPEXIT]]
 ; CHECK:       for.cond.cleanup.loopexit:
 ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
 ; CHECK:       for.cond.cleanup:
 ; CHECK-NEXT:    ret void
 ; CHECK:       for.body:
-; CHECK-NEXT:    [[I_06:%.*]] = phi i32 [ [[INC:%.*]], [[FOR_BODY]] ], [ 0, [[FOR_BODY_PREHEADER]] ]
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i32 [[I_06]], 0
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP1]], i32 [[A:%.*]], i32 [[B:%.*]]
-; CHECK-NEXT:    tail call void @f3(i32 noundef [[COND]])
+; CHECK-NEXT:    [[I_06:%.*]] = phi i32 [ [[INC:%.*]], [[FOR_BODY]] ], [ [[INC_PEEL]], [[FOR_BODY_PREHEADER_PEEL_NEWPH]] ]
+; CHECK-NEXT:    tail call void @f3(i32 noundef [[B]])
 ; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I_06]], 1
 ; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i32 [[INC]], [[NUM]]
-; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], label [[FOR_BODY]]
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP_LOOPEXIT_LOOPEXIT:%.*]], label [[FOR_BODY]], !llvm.loop [[LOOP16:![0-9]+]]
 ;
 entry:
   %cmp5 = icmp sgt i32 %num, 0
   br i1 %cmp5, label %for.body, label %for.cond.cleanup
 
 for.cond.cleanup:                                 ; preds = %for.body, %entry
   ret void