[SCEV] Compute ranges for ashr recurrences

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.)

Differential Revision: https://reviews.llvm.org/D101181

2 files changed, 27 insertions, 5 deletions

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 4de5abfdc3f1..c2bf309e0591 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -5692,10 +5692,11 @@ getRangeForUnknownRecurrence(const SCEVUnknown *U) {
     // 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;
@@ -5725,6 +5726,27 @@ getRangeForUnknownRecurrence(const SCEVUnknown *U) {
   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
diff --git a/llvm/test/Analysis/ScalarEvolution/shift-recurrences.ll b/llvm/test/Analysis/ScalarEvolution/shift-recurrences.ll
index 2ed80c7f0a43..568e014605cf 100644
--- a/llvm/test/Analysis/ScalarEvolution/shift-recurrences.ll
+++ b/llvm/test/Analysis/ScalarEvolution/shift-recurrences.ll
@@ -503,7 +503,7 @@ define void @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
@@ -534,7 +534,7 @@ define void @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
@@ -599,11 +599,11 @@ define void @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