[SCEV] Leverage inferred no-self-wrap flags to refine trip counts

[reames]

Summary

This is a follow on to D114176. It adds the missing piece to actually build a new expression for the LHS using the newly discovered flag information.

I need to add a few more tests, before this lands. In the meantime, a question for the reviewers: is the approach of removing nodes from UniqueSCEVs reasonable? This has some serious downsides (breaking equality), but the other option requires a ton of duplicated code. Anyone see any other ideas?

[nikic]

I'm really uncomfortable with removing things from UniqueSCEVs. If I understood correctly, your motivation here is that getZeroExtendExpr will not try folds if it has already memoized that the fact that this expression doesn't fold. I think we'd be better off adding a parameter to skip that check for this use-site.

Though I'm generally starting to have doubts about the direction of the whole patch stack (starting D103991). I find this kind of after the fact fixup of SCEVs rather problematic. Just the additional nowrap flags inferred during zext/sext construction can already cause hard to understand "spooky action at a distance", and this takes things one step further. I didn't really appreciate that this is where this is going.

[reames]

In D114185#3143327, @nikic wrote:

I'm really uncomfortable with removing things from UniqueSCEVs. If I understood correctly, your motivation here is that getZeroExtendExpr will not try folds if it has already memoized that the fact that this expression doesn't fold. I think we'd be better off adding a parameter to skip that check for this use-site.

Fair enough, that's why I asked. Will restructure.

Though I'm generally starting to have doubts about the direction of the whole patch stack (starting D103991). I find this kind of after the fact fixup of SCEVs rather problematic. Just the additional nowrap flags inferred during zext/sext construction can already cause hard to understand "spooky action at a distance", and this takes things one step further. I didn't really appreciate that this is where this is going.

Yeah, the problem is that "spooky action at a distance" is pretty fundamental to the interaction of trip count logic and SCEV construction. I could move this into the SCEV construction, but then we have to add a really nasty "is this scev directly used by exit test?" piece.

I have to admit I'm not really a fan of this myself, but I *really* don't see any better options. I've tried - like *really* tried, this is what resulted in the whole misadventure in indvars - but so far, I have absolutely nothing better. If you want to brainstorm, I'm happy to see if we can come up with anything better, but after a lot of thought, I've concluded this is the "least bad" place.

[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

auto peekThroughInvertibleFunctions = [&](const SCEV *S) -> const SCEV * {
         }
         if (!NonInvariantOp)
           // The add itself is invariant.
           return Add;
         return NonInvariantOp;
       }
       return S;
     };
-    auto *InnerLHS = LHS;
+    SmallVector<const SCEV *> OpStack;
+    OpStack.push_back(LHS);
     while (true) {
-      auto *Temp = peekThroughInvertibleFunctions(InnerLHS);
-      if (Temp == InnerLHS)
+      auto *Last = OpStack.back();
+      auto *Temp = peekThroughInvertibleFunctions(Last);
+      if (Temp == Last)
         break;
-      InnerLHS = Temp;
+      OpStack.push_back(Temp);
     }
-    if (const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(InnerLHS)) {
+    if (const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(OpStack.back())) {
       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);
         setNoWrapFlags(const_cast<SCEVAddRecExpr *>(AR), Flags);
+
+        // Get a new SCEV (hopefully an addrec), if we know that replacing
+        // the loop varying operand of S with NewOp is likely to allow
+        // specialization.  Note that this has the side effect of removing
+        // old SCEVs from the set which breaks equality comparison between
+        // previously handed out pointers and newly queried ones.
+        auto getWithVaryingOperandReplaced =
+          [&](const SCEV *S, const SCEVAddRecExpr *NewOp) {
+          if (isa<SCEVZeroExtendExpr>(S) && NewOp->hasNoUnsignedWrap()) {
+            UniqueSCEVs.RemoveNode(const_cast<SCEV*>(S));
+            return getZeroExtendExpr(NewOp, S->getType());
+          } else if (isa<SCEVSignExtendExpr>(S) &&
+                     NewOp->hasNoSignedWrap()) {
+            UniqueSCEVs.RemoveNode(const_cast<SCEV*>(S));
+            return getSignExtendExpr(NewOp, S->getType());
+          } else if (auto *SAdd = dyn_cast<SCEVAddExpr>(S)) {
+            assert(isa<SCEVAddRecExpr>(NewOp));
+            UniqueSCEVs.RemoveNode(const_cast<SCEV*>(S));
+            SmallVector<const SCEV *> Operands;
+            for (auto *Op :  SAdd->operands())
+              Operands.push_back(isLoopInvariant(Op, L) ? Op : NewOp);
+            return getAddExpr(Operands, SAdd->getNoWrapFlags());
+          }
+          return S;
+        };
+
+        const SCEVAddRecExpr *Rebuilt = AR;
+        assert(OpStack.back() == AR);
+        OpStack.pop_back();
+        while (!OpStack.empty() && Rebuilt) {
+          auto *Next = OpStack.pop_back_val();
+          auto *S = getWithVaryingOperandReplaced(Next, Rebuilt);
+          Rebuilt = dyn_cast<SCEVAddRecExpr>(S);
+        }
+        if (Rebuilt)
+          LHS = Rebuilt;
       }
     }
   }
 
   switch (Pred) {
   case ICmpInst::ICMP_NE: {                     // while (X != Y)
     // Convert to: while (X-Y != 0)
     if (LHS->getType()->isPointerTy()) {

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

 
 for.cond.cleanup:
   ret void
 }
 
 define void @test_zext(i64 %N) mustprogress {
 ; CHECK-LABEL: 'test_zext'
 ; CHECK-NEXT:  Determining loop execution counts for: @test_zext
-; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.
-; CHECK-NEXT:  Loop %for.body: Unpredictable max backedge-taken count.
+; CHECK-NEXT:  Loop %for.body: backedge-taken count is (%N /u 2)
+; CHECK-NEXT:  Loop %for.body: max backedge-taken count is 9223372036854775807
 ; 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: <nusw>
+; CHECK:       Loop %for.body: Trip multiple is 1
 ;
 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

 
 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: backedge-taken count is ((-21 + %N) /u 2)
+; CHECK-NEXT:  Loop %for.body: max backedge-taken count is 9223372036854775807
 ; 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>
+; CHECK:       Loop %for.body: Trip multiple is 1
 ;
 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