[SCEV] Distinguish between full and wrapping AddRec in proveNoWrapViaCR.

Try to widen integer AddRecs to a type one bit wider to distinguish
between the AddRec wrapping or just hitting all possible values.

Alternative to llvm#131538.

Note that now we can end up in the awkward situation that we fail to
compute an unpredicated BTC on the first try, but succeed on the second
try, because we now have a accurate max BTC.

For now, I updated getPredicatedBackedgeTakencCount to remove cached BTC
if that happens, but perhaps there's a better solution?

Author: fhahn

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -5092,9 +5092,29 @@ ScalarEvolution::proveNoWrapViaConstantRanges(const SCEVAddRecExpr *AR) {
   }
 
   if (!AR->hasNoUnsignedWrap()) {
-    ConstantRange AddRecRange = getUnsignedRange(AR);
-    ConstantRange IncRange = getUnsignedRange(AR->getStepRecurrence(*this));
-
+    const SCEVAddRecExpr *NewAR = AR;
+    unsigned BitWidth = getTypeSizeInBits(AR->getType());
+    // For integer AddRecs, try to evaluate the AddRec in a type one bit wider
+    // than the original type, to be able to differentiate between the AddRec
+    // hitting the full range or wrapping.
+    const SCEV *Step = AR->getStepRecurrence(*this);
+    if (AR->getType()->isIntegerTy() && isKnownNonNegative(Step)) {
+      Type *WiderTy = IntegerType::get(getContext(), BitWidth + 1);
+      NewAR = cast<SCEVAddRecExpr>(
+          getAddRecExpr(getSignExtendExpr(AR->getStart(), WiderTy),
+                        getZeroExtendExpr(Step, WiderTy), AR->getLoop(),
+                        AR->getNoWrapFlags()));
+      ConstantRange AddRecRange = getUnsignedRange(NewAR);
+      // If the wider AddRec range matches the original range after stripping
+      // the top bit, the original AddRec does not self-wrap.
+      if (AddRecRange !=
+          AddRecRange.truncate(BitWidth).zeroExtend(BitWidth + 1))
+        NewAR = AR;
+    }
+    ConstantRange AddRecRange = getUnsignedRange(NewAR);
+    ConstantRange IncRange = getUnsignedRange(Step);
+    if (NewAR != AR)
+      IncRange = IncRange.zeroExtend(getTypeSizeInBits(NewAR->getType()));
     auto NUWRegion = ConstantRange::makeGuaranteedNoWrapRegion(
         Instruction::Add, IncRange, OBO::NoUnsignedWrap);
     if (NUWRegion.contains(AddRecRange))
@@ -8336,7 +8356,13 @@ const SCEV *ScalarEvolution::getPredicatedExitCount(
 
 const SCEV *ScalarEvolution::getPredicatedBackedgeTakenCount(
     const Loop *L, SmallVectorImpl<const SCEVPredicate *> &Preds) {
-  return getPredicatedBackedgeTakenInfo(L).getExact(L, this, &Preds);
+  auto *Res = getPredicatedBackedgeTakenInfo(L).getExact(L, this, &Preds);
+  if (!isa<SCEVCouldNotCompute>(Res) && Preds.empty()) {
+    auto I = BackedgeTakenCounts.find(L);
+    if (I != BackedgeTakenCounts.end() && !I->second.isComplete())
+      BackedgeTakenCounts.erase(I);
+  }
+  return Res;
 }
 
 const SCEV *ScalarEvolution::getBackedgeTakenCount(const Loop *L,
@@ -13858,7 +13884,8 @@ static void PrintLoopInfo(raw_ostream &OS, ScalarEvolution *SE,
   SmallVector<const SCEVPredicate *, 4> Preds;
   auto *PBT = SE->getPredicatedBackedgeTakenCount(L, Preds);
   if (PBT != BTC) {
-    assert(!Preds.empty() && "Different predicated BTC, but no predicates");
+    assert((!Preds.empty() || PBT == SE->getBackedgeTakenCount(L)) &&
+           "Different predicated BTC, but no predicates");
     OS << "Loop ";
     L->getHeader()->printAsOperand(OS, /*PrintType=*/false);
     OS << ": ";
@@ -13877,7 +13904,8 @@ static void PrintLoopInfo(raw_ostream &OS, ScalarEvolution *SE,
   auto *PredConstantMax =
       SE->getPredicatedConstantMaxBackedgeTakenCount(L, Preds);
   if (PredConstantMax != ConstantBTC) {
-    assert(!Preds.empty() &&
+    assert((!Preds.empty() ||
+            PredConstantMax == SE->getConstantMaxBackedgeTakenCount(L)) &&
            "different predicated constant max BTC but no predicates");
     OS << "Loop ";
     L->getHeader()->printAsOperand(OS, /*PrintType=*/false);
@@ -13897,7 +13925,8 @@ static void PrintLoopInfo(raw_ostream &OS, ScalarEvolution *SE,
   auto *PredSymbolicMax =
       SE->getPredicatedSymbolicMaxBackedgeTakenCount(L, Preds);
   if (SymbolicBTC != PredSymbolicMax) {
-    assert(!Preds.empty() &&
+    assert((!Preds.empty() ||
+            PredSymbolicMax == SE->getSymbolicMaxBackedgeTakenCount(L)) &&
            "Different predicated symbolic max BTC, but no predicates");
     OS << "Loop ";
     L->getHeader()->printAsOperand(OS, /*PrintType=*/false);

llvm/test/Analysis/ScalarEvolution/exit-count-non-strict.ll

@@ -153,26 +153,22 @@ define void @ule_from_zero_no_nuw(i32 %M, i32 %N) {
 ; CHECK-NEXT:    exit count for loop: ***COULDNOTCOMPUTE***
 ; CHECK-NEXT:    predicated exit count for loop: (1 + (zext i32 %M to i64))<nuw><nsw>
 ; CHECK-NEXT:     Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-EMPTY:
 ; CHECK-NEXT:    exit count for latch: %N
 ; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is i32 -1
 ; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is %N
 ; CHECK-NEXT:    symbolic max exit count for loop: ***COULDNOTCOMPUTE***
 ; CHECK-NEXT:    predicated symbolic max exit count for loop: (1 + (zext i32 %M to i64))<nuw><nsw>
 ; CHECK-NEXT:     Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-EMPTY:
 ; CHECK-NEXT:    symbolic max exit count for latch: %N
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is ((zext i32 %N to i64) umin (1 + (zext i32 %M to i64))<nuw><nsw>)
 ; CHECK-NEXT:   Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-NEXT:  Loop %loop: Predicated constant max backedge-taken count is i64 4294967295
 ; CHECK-NEXT:   Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-NEXT:  Loop %loop: Predicated symbolic max backedge-taken count is ((zext i32 %N to i64) umin (1 + (zext i32 %M to i64))<nuw><nsw>)
 ; CHECK-NEXT:   Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
+; CHECK-NEXT:  Loop %loop: Trip multiple is 1
 ;
 entry:
   br label %loop
@@ -198,26 +194,22 @@ define void @le_from_zero_no_nuw(i32 %M, i32 %N) {
 ; CHECK-NEXT:    exit count for loop: ***COULDNOTCOMPUTE***
 ; CHECK-NEXT:    predicated exit count for loop: (1 + (zext i32 %M to i64))<nuw><nsw>
 ; CHECK-NEXT:     Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-EMPTY:
 ; CHECK-NEXT:    exit count for latch: %N
 ; CHECK-NEXT:  Loop %loop: constant max backedge-taken count is i32 -1
 ; CHECK-NEXT:  Loop %loop: symbolic max backedge-taken count is %N
 ; CHECK-NEXT:    symbolic max exit count for loop: ***COULDNOTCOMPUTE***
 ; CHECK-NEXT:    predicated symbolic max exit count for loop: (1 + (zext i32 %M to i64))<nuw><nsw>
 ; CHECK-NEXT:     Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-EMPTY:
 ; CHECK-NEXT:    symbolic max exit count for latch: %N
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is ((zext i32 %N to i64) umin (1 + (zext i32 %M to i64))<nuw><nsw>)
 ; CHECK-NEXT:   Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-NEXT:  Loop %loop: Predicated constant max backedge-taken count is i64 4294967295
 ; CHECK-NEXT:   Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
 ; CHECK-NEXT:  Loop %loop: Predicated symbolic max backedge-taken count is ((zext i32 %N to i64) umin (1 + (zext i32 %M to i64))<nuw><nsw>)
 ; CHECK-NEXT:   Predicates:
-; CHECK-NEXT:      {0,+,1}<%loop> Added Flags: <nusw>
+; CHECK-NEXT:  Loop %loop: Trip multiple is 1
 ;
 entry:
   br label %loop