[SCEV] Look through invertible functions when infering no-self-wrap from mustprogres

[reames]

Summary

This is an extension to the logic introduced with D103991. The basic idea is that any invertible function of a repeating sequence is itself a repeating sequence, and that since invertible functions are 1-to-1, we can conclude that a wrapped IV can't cycle.

Note that this only shows benefit for predicated SCEV results. The missing piece to directly influence trip counts is rebuilding the invertible chain using the newly inferred flags. That'll be a separate change. :)

[reames]

Will reopen when I get time to return to this. Closing in the meantime to get off review queues.

llvm/lib/Analysis/ScalarEvolution.cpp

ScalarEvolution::computeExitLimitFromICmp(const Loop *L,
   // If this loop must exit based on this condition (or execute undefined
   // behaviour), and we can prove the test sequence produced must repeat
   // the same values on self-wrap of the IV, then we can infer that IV
   // doesn't self wrap because if it did, we'd have an infinite (undefined)
   // loop.
   if (ControlsExit && isLoopInvariant(RHS, L) && loopHasNoAbnormalExits(L) &&
       loopIsFiniteByAssumption(L)) {
 
-    // TODO: We can peel off any functions which are invertible *in L*.  Loop
-    // invariant terms are effectively constants for our purposes here.
+    // Peel off any functions which are invertible *in L*.  Loop invariant
+    // terms are effectively constants for our purposes here.
+    auto peekThroughInvertibleFunctions = [&](const SCEV *S) -> const SCEV * {
+      if (auto *ZExt = dyn_cast<SCEVZeroExtendExpr>(S))
+        return ZExt->getOperand();
+      else if (auto *SExt = dyn_cast<SCEVSignExtendExpr>(S))
+        return SExt->getOperand();
+      else if (auto *Add = dyn_cast<SCEVAddExpr>(S)) {
+        const SCEV *NonInvariantOp = nullptr;
+        for (auto *Op : Add->operands()) {
+          if (isLoopInvariant(Op, L))
+            continue;
+          if (!NonInvariantOp) {
+            NonInvariantOp = Op;
+            continue;
+          }
+          if (Op != NonInvariantOp)
+            // Can't peak through this
+            return Add;
+        }
+        if (!NonInvariantOp)
+          // The add itself is invariant.
+          return Add;
+        return NonInvariantOp;
+      }
+      return S;
+    };
     auto *InnerLHS = LHS;
-    if (auto *ZExt = dyn_cast<SCEVZeroExtendExpr>(LHS))
-      InnerLHS = ZExt->getOperand();
+    while (true) {
+      auto *Temp = peekThroughInvertibleFunctions(InnerLHS);
+      if (Temp == InnerLHS)
+        break;
+      InnerLHS = Temp;
+    }
     if (const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(InnerLHS)) {
       auto *StrideC = dyn_cast<SCEVConstant>(AR->getStepRecurrence(*this));
       if (!AR->hasNoSelfWrap() && AR->getLoop() == L && AR->isAffine() && 
           StrideC && StrideC->getAPInt().isPowerOf2()) {
         auto Flags = AR->getNoWrapFlags();
         Flags = setFlags(Flags, SCEV::FlagNW);
         SmallVector<const SCEV*> Operands{AR->operands()};
         Flags = StrengthenNoWrapFlags(this, scAddRecExpr, Operands, Flags);

llvm/test/Analysis/ScalarEvolution/ne-overflow.ll

for.cond.cleanup:
   ret void
 }
 
 define void @test_sext(i64 %N) mustprogress {
 ; CHECK-LABEL: 'test_sext'
 ; CHECK-NEXT:  Determining loop execution counts for: @test_sext
 ; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.
 ; CHECK-NEXT:  Loop %for.body: Unpredictable max backedge-taken count.
-; CHECK-NEXT:  Loop %for.body: Unpredictable predicated backedge-taken count.
+; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (%N /u 2)
+; CHECK-NEXT:   Predicates:
+; CHECK-NEXT:    {0,+,2}<nuw><%for.body> Added Flags: <nssw>
 ;
 entry:
   br label %for.body
 
 for.body:
   %iv = phi i32 [ %iv.next, %for.body ], [ 0, %entry ]
   %iv.next = add i32 %iv, 2
   %sext = sext i32 %iv to i64
   %cmp = icmp ne i64 %sext, %N
   br i1 %cmp, label %for.body, label %for.cond.cleanup
 
 for.cond.cleanup:
   ret void
 }
 
 define void @test_zext_of_sext(i64 %N) mustprogress {
 ; CHECK-LABEL: 'test_zext_of_sext'
 ; CHECK-NEXT:  Determining loop execution counts for: @test_zext_of_sext
 ; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.
 ; CHECK-NEXT:  Loop %for.body: Unpredictable max backedge-taken count.
-; CHECK-NEXT:  Loop %for.body: Unpredictable predicated backedge-taken count.
+; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (%N /u 2)
+; CHECK-NEXT:   Predicates:
+; CHECK-NEXT:    {0,+,2}<nuw><%for.body> Added Flags: <nssw>
+; CHECK-NEXT:    {0,+,2}<nuw><%for.body> Added Flags: <nusw>
 ;
 entry:
   br label %for.body
 
 for.body:
   %iv = phi i32 [ %iv.next, %for.body ], [ 0, %entry ]
   %iv.next = add i32 %iv, 2
   %sext = sext i32 %iv to i48
   %zext = zext i48 %sext to i64
   %cmp = icmp ne i64 %zext, %N
   br i1 %cmp, label %for.body, label %for.cond.cleanup
 
 for.cond.cleanup:
   ret void
 }
 
 define void @test_zext_offset(i64 %N) mustprogress {
 ; CHECK-LABEL: 'test_zext_offset'
 ; CHECK-NEXT:  Determining loop execution counts for: @test_zext_offset
 ; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.
 ; CHECK-NEXT:  Loop %for.body: Unpredictable max backedge-taken count.
-; CHECK-NEXT:  Loop %for.body: Unpredictable predicated backedge-taken count.
+; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is ((-21 + %N) /u 2)
+; CHECK-NEXT:   Predicates:
+; CHECK-NEXT:    {0,+,2}<nuw><%for.body> Added Flags: <nusw>
 ;
 entry:
   br label %for.body
 
 for.body:
   %iv = phi i32 [ %iv.next, %for.body ], [ 0, %entry ]
   %iv.next = add i32 %iv, 2
   %zext = zext i32 %iv to i64
   %offset = add i64 %zext, 21
   %cmp = icmp ne i64 %offset, %N
   br i1 %cmp, label %for.body, label %for.cond.cleanup
 
 for.cond.cleanup:
   ret void
 }
 
 define void @test_sext_offset(i64 %N) mustprogress {
 ; CHECK-LABEL: 'test_sext_offset'
 ; CHECK-NEXT:  Determining loop execution counts for: @test_sext_offset
 ; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.
 ; CHECK-NEXT:  Loop %for.body: Unpredictable max backedge-taken count.
-; CHECK-NEXT:  Loop %for.body: Unpredictable predicated backedge-taken count.
+; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is ((-21 + %N) /u 2)
+; CHECK-NEXT:   Predicates:
+; CHECK-NEXT:    {0,+,2}<nuw><%for.body> Added Flags: <nssw>
 ;
 entry:
   br label %for.body
 
 for.body:
   %iv = phi i32 [ %iv.next, %for.body ], [ 0, %entry ]
   %iv.next = add i32 %iv, 2
   %sext = sext i32 %iv to i64
   %offset = add i64 %sext, 21
   %cmp = icmp ne i64 %offset, %N
   br i1 %cmp, label %for.body, label %for.cond.cleanup
 
 for.cond.cleanup:
   ret void
 }

llvm/test/Analysis/ScalarEvolution/trip-count-implied-addrec.ll

 
 define void @nw_implies_nsw(i16 %n) mustprogress {
 ; CHECK-LABEL: 'nw_implies_nsw'
 ; CHECK-NEXT:  Determining loop execution counts for: @nw_implies_nsw
 ; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.
 ; CHECK-NEXT:  Loop %for.body: Unpredictable max backedge-taken count.
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (128 + (-128 smax %n))
 ; CHECK-NEXT:   Predicates:
-; CHECK-NEXT:    {-128,+,1}<%for.body> Added Flags: <nssw>
+; CHECK-NEXT:    {-128,+,1}<nw><%for.body> Added Flags: <nssw>
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.body
   %iv = phi i8 [ %iv.next, %for.body ], [ -128, %entry ]
   %iv.next = add i8 %iv, 1
   %zext = sext i8 %iv to i16