Skip to content

[LoopPeel] Use loop guards when checking if last iter can be peeled. #142605

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

fhahn
Copy link
Contributor

@fhahn fhahn commented Jun 3, 2025

Apply loop guards to BTC before checking if the last iteration should be peeled off. This also adds an assert to make sure applying the guards does not pessimize the results. I checked on a large test set and it did not trigger there, but it adds an additional guard to catch potential cases where loop-guards pessimize results.

Peels ~15% more loops.

@llvmbot
Copy link
Member

llvmbot commented Jun 3, 2025

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

Changes

Apply loop guards to BTC before checking if the last iteration should be peeled off. This also adds an assert to make sure applying the guards does not pessimize the results. I checked on a large test set and it did not trigger there, but it adds an additional guard to catch potential cases where loop-guards pessimize results.

Peels ~15% more loops.


Full diff: https://github.com/llvm/llvm-project/pull/142605.diff

2 Files Affected:

  • (modified) llvm/lib/Transforms/Utils/LoopPeel.cpp (+17-10)
  • (modified) llvm/test/Transforms/LoopUnroll/peel-last-iteration-with-guards.ll (+24-4)
diff --git a/llvm/lib/Transforms/Utils/LoopPeel.cpp b/llvm/lib/Transforms/Utils/LoopPeel.cpp
index bd025fddd0cf7..564c2298f6f2f 100644
--- a/llvm/lib/Transforms/Utils/LoopPeel.cpp
+++ b/llvm/lib/Transforms/Utils/LoopPeel.cpp
@@ -357,17 +357,14 @@ bool llvm::canPeelLastIteration(const Loop &L, ScalarEvolution &SE) {
                m_scev_AffineAddRec(m_SCEV(), m_scev_One(), m_SpecificLoop(&L)));
 }
 
-/// 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.
+/// Returns true if the last iteration should be peeled off, i.e. 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,
+                                    const SCEV *RightSCEV, const SCEV *BTC,
+                                    ScalarEvolution &SE,
                                     const TargetTransformInfo &TTI) {
-  if (!canPeelLastIteration(L, SE))
-    return false;
-
-  const SCEV *BTC = SE.getBackedgeTakenCount(&L);
   SCEVExpander Expander(SE, L.getHeader()->getDataLayout(), "loop-peel");
   if (!SE.isKnownNonZero(BTC) &&
       Expander.isHighCostExpansion(BTC, &L, SCEVCheapExpansionBudget, &TTI,
@@ -377,7 +374,6 @@ static bool shouldPeelLastIteration(Loop &L, CmpPredicate Pred,
   const SCEV *ValAtLastIter = LeftAR->evaluateAtIteration(BTC, SE);
   const SCEV *ValAtSecondToLastIter = LeftAR->evaluateAtIteration(
       SE.getMinusSCEV(BTC, SE.getOne(BTC->getType())), SE);
-
   return SE.isKnownPredicate(ICmpInst::getInversePredicate(Pred), ValAtLastIter,
                              RightSCEV) &&
          SE.isKnownPredicate(Pred, ValAtSecondToLastIter, RightSCEV);
@@ -484,8 +480,19 @@ countToEliminateCompares(Loop &L, unsigned MaxPeelCount, ScalarEvolution &SE,
     const SCEV *Step = LeftAR->getStepRecurrence(SE);
     if (!PeelWhilePredicateIsKnown(NewPeelCount, IterVal, RightSCEV, Step,
                                    Pred)) {
-      if (shouldPeelLastIteration(L, Pred, LeftAR, RightSCEV, SE, TTI))
+      if (!canPeelLastIteration(L, SE))
+        return;
+
+      const SCEV *BTC = SE.getBackedgeTakenCount(&L);
+      auto Guards = ScalarEvolution::LoopGuards::collect(&L, SE);
+      if (shouldPeelLastIteration(L, Pred, LeftAR,
+                                  SE.applyLoopGuards(RightSCEV, Guards),
+                                  SE.applyLoopGuards(BTC, Guards), SE, TTI))
         DesiredPeelCountLast = 1;
+      else
+        assert(!shouldPeelLastIteration(L, Pred, LeftAR, RightSCEV, BTC, SE,
+                                        TTI) &&
+               "loop guards pessimized result");
       return;
     }
 
diff --git a/llvm/test/Transforms/LoopUnroll/peel-last-iteration-with-guards.ll b/llvm/test/Transforms/LoopUnroll/peel-last-iteration-with-guards.ll
index af07a97131322..824e23fcf3e6e 100644
--- a/llvm/test/Transforms/LoopUnroll/peel-last-iteration-with-guards.ll
+++ b/llvm/test/Transforms/LoopUnroll/peel-last-iteration-with-guards.ll
@@ -13,14 +13,33 @@ define void @peel_with_guard_known_nonnegative_1(i32 %n) {
 ; CHECK-NEXT:    [[N_EXT:%.*]] = zext i32 [[N]] to i64
 ; CHECK-NEXT:    [[N_1:%.*]] = add i32 [[N]], 1
 ; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = zext i32 [[N_1]] to i64
+; CHECK-NEXT:    [[TMP0:%.*]] = add nsw i64 [[WIDE_TRIP_COUNT]], -1
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i64 [[TMP0]], 0
+; CHECK-NEXT:    br i1 [[TMP1]], label %[[PH_SPLIT:.*]], label %[[EXIT_LOOPEXIT_PEEL_BEGIN:.*]]
+; CHECK:       [[PH_SPLIT]]:
 ; CHECK-NEXT:    br label %[[LOOP:.*]]
 ; CHECK:       [[LOOP]]:
-; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, %[[PH]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[IV1:%.*]] = phi i64 [ 0, %[[PH_SPLIT]] ], [ [[IV_NEXT1:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT1]] = add nuw nsw i64 [[IV1]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[WIDE_TRIP_COUNT]], 1
+; CHECK-NEXT:    [[EC1:%.*]] = icmp eq i64 [[IV_NEXT1]], [[TMP2]]
+; CHECK-NEXT:    br i1 [[EC1]], label %[[EXIT_LOOPEXIT_PEEL_BEGIN_LOOPEXIT:.*]], label %[[LOOP]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       [[EXIT_LOOPEXIT_PEEL_BEGIN_LOOPEXIT]]:
+; CHECK-NEXT:    [[DOTPH:%.*]] = phi i64 [ [[IV_NEXT1]], %[[LOOP]] ]
+; CHECK-NEXT:    br label %[[EXIT_LOOPEXIT_PEEL_BEGIN]]
+; CHECK:       [[EXIT_LOOPEXIT_PEEL_BEGIN]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, %[[PH]] ], [ [[DOTPH]], %[[EXIT_LOOPEXIT_PEEL_BEGIN_LOOPEXIT]] ]
+; CHECK-NEXT:    br label %[[LOOP_PEEL:.*]]
+; CHECK:       [[LOOP_PEEL]]:
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i64 [[IV]], [[N_EXT]]
 ; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C]], i32 10, i32 20
-; CHECK-NEXT:    [[IV_NEXT]] = add i64 [[IV]], 1
+; CHECK-NEXT:    [[IV_NEXT:%.*]] = add i64 [[IV]], 1
 ; CHECK-NEXT:    [[EC:%.*]] = icmp eq i64 [[IV_NEXT]], [[WIDE_TRIP_COUNT]]
-; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT_LOOPEXIT:.*]], label %[[LOOP]]
+; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT_LOOPEXIT_PEEL_NEXT:.*]], label %[[EXIT_LOOPEXIT_PEEL_NEXT]]
+; CHECK:       [[EXIT_LOOPEXIT_PEEL_NEXT]]:
+; CHECK-NEXT:    br label %[[LOOP_PEEL_NEXT:.*]]
+; CHECK:       [[LOOP_PEEL_NEXT]]:
+; CHECK-NEXT:    br label %[[EXIT_LOOPEXIT:.*]]
 ; CHECK:       [[EXIT_LOOPEXIT]]:
 ; CHECK-NEXT:    br label %[[EXIT]]
 ; CHECK:       [[EXIT]]:
@@ -137,7 +156,7 @@ define void @peel_with_guard2(i32 %n) {
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw i32 [[IV]], 1
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i32 [[N]], 1
 ; CHECK-NEXT:    [[EC:%.*]] = icmp eq i32 [[IV_NEXT]], [[TMP2]]
-; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT_LOOPEXIT_PEEL_BEGIN_LOOPEXIT:.*]], label %[[LOOP_HEADER]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT_LOOPEXIT_PEEL_BEGIN_LOOPEXIT:.*]], label %[[LOOP_HEADER]], !llvm.loop [[LOOP2:![0-9]+]]
 ; CHECK:       [[EXIT_LOOPEXIT_PEEL_BEGIN_LOOPEXIT]]:
 ; CHECK-NEXT:    [[DOTPH:%.*]] = phi i32 [ [[IV_NEXT]], %[[LOOP_LATCH]] ]
 ; CHECK-NEXT:    br label %[[EXIT_LOOPEXIT_PEEL_BEGIN]]
@@ -188,4 +207,5 @@ exit:
 ;.
 ; CHECK: [[LOOP0]] = distinct !{[[LOOP0]], [[META1:![0-9]+]]}
 ; CHECK: [[META1]] = !{!"llvm.loop.peeled.count", i32 1}
+; CHECK: [[LOOP2]] = distinct !{[[LOOP2]], [[META1]]}
 ;.

else
assert(!shouldPeelLastIteration(L, Pred, LeftAR, RightSCEV, BTC, SE,
TTI) &&
"loop guards pessimized result");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this assertion should exist. Based on past experience, you can always construct edge cases where this happens.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah fair enough, I removed it. Still wondering if there's a nice way to have some kind of verification as opt-in?

if (!canPeelLastIteration(L, SE))
return false;

const SCEV *BTC = SE.getBackedgeTakenCount(&L);
SCEVExpander Expander(SE, L.getHeader()->getDataLayout(), "loop-peel");
if (!SE.isKnownNonZero(BTC) &&
Expander.isHighCostExpansion(BTC, &L, SCEVCheapExpansionBudget, &TTI,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The high cost expansion check should probably not use the BTC with loop guards? We're not going to actually expand it with the loop guards, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep loop-guards aren't used during expansion, removed for now. Some of the info may be used by later passes to simplifiy the expansion.

Apply loop guards to BTC before checking if the last iteration should be
peeled off. This also adds an assert to make sure applying the guards
does not pessimize the results. I checked on a large test set and it did
not trigger there, but it adds an additional guard to catch potential
cases where loop-guards pessimize results.

Peels ~15% more loops.
@fhahn fhahn force-pushed the peel-last-iter-with-guards branch from 8e0bea9 to 4b4a394 Compare June 3, 2025 20:50
@fhahn fhahn force-pushed the peel-last-iter-with-guards branch from 4b4a394 to c7deb11 Compare June 3, 2025 20:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants