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

Generalize countToEliminateCompares to also consider peeling off the
last iteration if it eliminates a compare.

At the moment, codegen for peeling off the last iteration is quite
restrictive and callers have to make sure that the exit condition can be
adjusted when peeling and that the loop executes at least 2 iterations.

Both will be relaxed in follow-ups.

PR: #139551

Author: fhahn

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -680,6 +680,9 @@ class TargetTransformInfo {
     /// If the value is true the peeling cost model can decide to peel only
     /// some iterations and in this case it will set this to false.
     bool PeelProfiledIterations;
+
+    /// Peel off the last PeelCount loop iterations.
+    bool PeelLast;
   };
 
   /// Get target-customized preferences for the generic loop peeling

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

@@ -21,11 +21,18 @@ 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.
-bool peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI, ScalarEvolution *SE,
-              DominatorTree &DT, AssumptionCache *AC, bool PreserveLCSSA,
-              ValueToValueMapTy &VMap);
+/// instructions in the last peeled-off iteration. If \p PeelLast is true, peel
+/// 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);
 
 TargetTransformInfo::PeelingPreferences
 gatherPeelingPreferences(Loop *L, ScalarEvolution &SE,

llvm/lib/Transforms/Scalar/LoopFuse.cpp

@@ -790,7 +790,8 @@ struct LoopFuser {
                       << " iterations of the first loop. \n");
 
     ValueToValueMapTy VMap;
-    FC0.Peeled = peelLoop(FC0.L, PeelCount, &LI, &SE, DT, &AC, true, VMap);
+    FC0.Peeled =
+        peelLoop(FC0.L, PeelCount, false, &LI, &SE, DT, &AC, true, VMap);
     if (FC0.Peeled) {
       LLVM_DEBUG(dbgs() << "Done Peeling\n");
 

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

@@ -1314,7 +1314,8 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
     });
 
     ValueToValueMapTy VMap;
-    if (peelLoop(L, PP.PeelCount, LI, &SE, DT, &AC, PreserveLCSSA, VMap)) {
+    if (peelLoop(L, PP.PeelCount, PP.PeelLast, LI, &SE, DT, &AC, PreserveLCSSA,
+                 VMap)) {
       simplifyLoopAfterUnroll(L, true, LI, &SE, &DT, &AC, &TTI, nullptr);
       // If the loop was peeled, we already "used up" the profile information
       // we had, so we don't want to unroll or peel again.