[SCEV] Compute ranges for ashr recurrences

[reames]

Summary

Straight forward extension to the recently added infrastructure which was pioneered with shl. This was originally posted as part of D99687, but split off for ease of review.

(I also decided to exclude the unknown start sign case explicitly for simplicity of understanding.)

[nikic]

LGTM

llvm/lib/Analysis/ScalarEvolution.cpp

5695

Outdated TODO

5741

As a general note on this function: You initialize CR to the full set, then always intersect you result with it, and return the intersection... you may as well just directly return your result range.

[reames]

llvm/lib/Analysis/ScalarEvolution.cpp

5741

Done in 0cc3e10f

llvm/lib/Analysis/ScalarEvolution.cpp

getRangeForUnknownRecurrence(const SCEVUnknown *U) {
   if (!L->contains(BO->getParent()))
     // NOTE: This bailout should be an assert instead.  However, asserting
     // the condition here exposes a case where LoopFusion is querying SCEV
     // with malformed loop information during the midst of the transform.
     // There doesn't appear to be an obvious fix, so for the moment bailout
     // until the caller issue can be fixed.  PR49566 tracks the bug.
     return CR;
 
-  // TODO: Extend to other opcodes such as ashr, mul, and div
+  // TODO: Extend to other opcodes such as mul, and div
   switch (BO->getOpcode()) {
   default:
     return CR;
+  case Instruction::AShr:
   case Instruction::LShr:
   case Instruction::Shl:
     break;
   };
 
   if (BO->getOperand(0) != P)
     // TODO: Handle the power function forms some day.
     return CR;
getRangeForUnknownRecurrence(const SCEVUnknown *U) {
   bool Overflow = false;
   auto TotalShift = MaxShiftAmt.umul_ov(TCAP, Overflow);
   if (Overflow)
     return CR;
 
   switch (BO->getOpcode()) {
   default:
     llvm_unreachable("filtered out above");
+  case Instruction::AShr: {
+    // For each ashr, three cases:
+    //   shift = 0 => unchanged value
+    //   saturation => 0 or -1
+    //   other => a value closer to zero (of the same sign)
+    // Thus, the end value is closer to zero than the start.
+    auto KnownEnd = KnownBits::ashr(KnownStart,
+                                    KnownBits::makeConstant(TotalShift));
+    if (KnownStart.isNonNegative()) {
+      // Analogous to lshr (simply not yet canonicalized)
+      auto R = ConstantRange::getNonEmpty(KnownEnd.getMinValue(),
+                                          KnownStart.getMaxValue() + 1);
+      CR = CR.intersectWith(R);
+    } else if (KnownStart.isNegative()) {
+      // End >=u Start && End <=s Start
+      auto R = ConstantRange::getNonEmpty(KnownStart.getMinValue(),
+                                          KnownEnd.getMaxValue() + 1);
+      CR = CR.intersectWith(R);
+    }
+    break;
+  }
   case Instruction::LShr: {
     // For each lshr, three cases:
     //   shift = 0 => unchanged value
     //   saturation => 0
     //   other => a smaller positive number
     // Thus, the low end of the unsigned range is the last value produced.
     auto KnownEnd = KnownBits::lshr(KnownStart,
                                     KnownBits::makeConstant(TotalShift));

llvm/test/Analysis/ScalarEvolution/shift-recurrences.ll

 ; by using the trip count information on the loop.
 
 define void @test_ashr_tc_positive() {
 ; CHECK-LABEL: 'test_ashr_tc_positive'
 ; CHECK-NEXT:  Classifying expressions for: @test_ashr_tc_positive
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<%loop> U: [0,5) S: [0,5) Exits: 4 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.ashr = phi i64 [ 1023, %entry ], [ %iv.ashr.next, %loop ]
-; CHECK-NEXT:    --> %iv.ashr U: [0,1024) S: [0,1024) Exits: 63 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.ashr U: [63,1024) S: [63,1024) Exits: 63 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
 ; CHECK-NEXT:    --> {1,+,1}<%loop> U: [1,6) S: [1,6) Exits: 5 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.ashr.next = ashr i64 %iv.ashr, 1
 ; CHECK-NEXT:    --> %iv.ashr.next U: [0,512) S: [0,512) Exits: 31 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_ashr_tc_positive
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: max backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is 4

 }
 
 define void @test_ashr_tc_negative() {
 ; CHECK-LABEL: 'test_ashr_tc_negative'
 ; CHECK-NEXT:  Classifying expressions for: @test_ashr_tc_negative
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<%loop> U: [0,5) S: [0,5) Exits: 4 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.ashr = phi i8 [ -128, %entry ], [ %iv.ashr.next, %loop ]
-; CHECK-NEXT:    --> %iv.ashr U: [-128,0) S: [-128,0) Exits: -8 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.ashr U: [-128,-7) S: [-128,-7) Exits: -8 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
 ; CHECK-NEXT:    --> {1,+,1}<%loop> U: [1,6) S: [1,6) Exits: 5 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.ashr.next = ashr i8 %iv.ashr, 1
 ; CHECK-NEXT:    --> %iv.ashr.next U: [-64,0) S: [-64,0) Exits: -4 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_ashr_tc_negative
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: max backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is 4

 }
 
 define void @test_ashr_zero_shift() {
 ; CHECK-LABEL: 'test_ashr_zero_shift'
 ; CHECK-NEXT:  Classifying expressions for: @test_ashr_zero_shift
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<%loop> U: [0,5) S: [0,5) Exits: 4 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.ashr = phi i64 [ 1023, %entry ], [ %iv.ashr.next, %loop ]
-; CHECK-NEXT:    --> %iv.ashr U: [0,1024) S: [0,1024) Exits: 1023 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.ashr U: [1023,1024) S: [1023,1024) Exits: 1023 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
 ; CHECK-NEXT:    --> {1,+,1}<%loop> U: [1,6) S: [1,6) Exits: 5 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.ashr.next = ashr i64 %iv.ashr, 0
-; CHECK-NEXT:    --> %iv.ashr U: [0,1024) S: [0,1024) Exits: 1023 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.ashr U: [1023,1024) S: [1023,1024) Exits: 1023 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_ashr_zero_shift
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: max backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is 4
 ; CHECK-NEXT:   Predicates:
 ; CHECK:       Loop %loop: Trip multiple is 5
 ;
 entry: