[llvm] [LoopPeel] Implement initial peeling off the last loop iteration. (PR #139551)

[Florian Hahn via llvm-commits]

================
@@ -325,19 +326,66 @@ 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) {
+  const SCEV *BTC = SE.getBackedgeTakenCount(&L);
+  Value *Inc;
+  CmpPredicate Pred;
+  BasicBlock *Succ1;
+  BasicBlock *Succ2;
+  // The loop must execute at least 2 iterations to guarantee that peeled
+  // iteration executes.
+  // TODO: Add checks during codegen.
+  if (isa<SCEVCouldNotCompute>(BTC) ||
+      !SE.isKnownPredicate(CmpInst::ICMP_UGT, BTC, SE.getOne(BTC->getType())))
+    return false;
+
+  // Check if the exit condition of the loop can be adjusted by the peeling
+  // codegen. For now, it must
+  // * exit via the latch,
+  // * 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 false;
+
+  const SCEV *ValAtLastIter =
+      SE.applyLoopGuards(LeftAR->evaluateAtIteration(BTC, SE), &L);
+  const SCEV *ValAtSecondToLastIter = LeftAR->evaluateAtIteration(
+      SE.getMinusSCEV(BTC, SE.getOne(BTC->getType())), SE);
+
+  return SE.isKnownPredicate(ICmpInst::getInversePredicate(Pred), ValAtLastIter,
+                             SE.applyLoopGuards(RightSCEV, &L)) &&
+         SE.isKnownPredicate(Pred, ValAtSecondToLastIter, RightSCEV);
+}
+
 // Return the number of iterations to peel off that make conditions in the
-// body true/false. For example, if we peel 2 iterations off the loop below,
-// the condition i < 2 can be evaluated at compile time.
+// body true/false. Positive return values indicate the iterations to peel of
+// from the front and negative return values indicate the number of iterations
+// from the back after removing the sign. For example, if we peel 2 iterations
+// off the loop below, the condition i < 2 can be evaluated at compile time.
 //  for (i = 0; i < n; i++)
 //    if (i < 2)
 //      ..
 //    else
 //      ..
 //   }
-static unsigned countToEliminateCompares(Loop &L, unsigned MaxPeelCount,
-                                         ScalarEvolution &SE) {
+static int countToEliminateCompares(Loop &L, unsigned MaxPeelCount,
----------------
fhahn wrote:

Split up thanks!

https://github.com/llvm/llvm-project/pull/139551