2 files changed, 118 insertions, 2 deletions

diff --git a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
index ec8b0d42626..4ebafd8ea78 100644
--- a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
@@ -25,6 +25,7 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -80,6 +81,7 @@ namespace {
                               bool IsSigned);
     bool eliminateSDiv(BinaryOperator *SDiv);
     bool strengthenOverflowingOperation(BinaryOperator *OBO, Value *IVOperand);
+    bool strengthenRightShift(BinaryOperator *BO, Value *IVOperand);
   };
 }
 
@@ -583,6 +585,35 @@ bool SimplifyIndvar::strengthenOverflowingOperation(BinaryOperator *BO,
   return Changed;
 }
 
+/// Annotate the Shr in (X << IVOperand) >> C as exact using the
+/// information from the IV's range. Returns true if anything changed, false
+/// otherwise.
+bool SimplifyIndvar::strengthenRightShift(BinaryOperator *BO,
+                                          Value *IVOperand) {
+  using namespace llvm::PatternMatch;
+
+  if (BO->getOpcode() == Instruction::Shl) {
+    bool Changed = false;
+    ConstantRange IVRange = SE->getUnsignedRange(SE->getSCEV(IVOperand));
+    for (auto *U : BO->users()) {
+      const APInt *C;
+      if (match(U,
+                m_AShr(m_Shl(m_Value(), m_Specific(IVOperand)), m_APInt(C))) ||
+          match(U,
+                m_LShr(m_Shl(m_Value(), m_Specific(IVOperand)), m_APInt(C)))) {
+        BinaryOperator *Shr = cast<BinaryOperator>(U);
+        if (!Shr->isExact() && IVRange.getUnsignedMin().uge(*C)) {
+          Shr->setIsExact(true);
+          Changed = true;
+        }
+      }
+    }
+    return Changed;
+  }
+
+  return false;
+}
+
 /// Add all uses of Def to the current IV's worklist.
 static void pushIVUsers(
   Instruction *Def,
@@ -675,8 +706,9 @@ void SimplifyIndvar::simplifyUsers(PHINode *CurrIV, IVVisitor *V) {
     }
 
     if (BinaryOperator *BO = dyn_cast<BinaryOperator>(UseOper.first)) {
-      if (isa<OverflowingBinaryOperator>(BO) &&
-          strengthenOverflowingOperation(BO, IVOperand)) {
+      if ((isa<OverflowingBinaryOperator>(BO) &&
+           strengthenOverflowingOperation(BO, IVOperand)) ||
+          (isa<ShlOperator>(BO) && strengthenRightShift(BO, IVOperand))) {
         // re-queue uses of the now modified binary operator and fall
         // through to the checks that remain.
         pushIVUsers(IVOperand, Simplified, SimpleIVUsers);
diff --git a/llvm/test/Transforms/IndVarSimplify/strengthen-overflow.ll b/llvm/test/Transforms/IndVarSimplify/strengthen-overflow.ll
index 2bafe96e1cc..6e0538e04d6 100644
--- a/llvm/test/Transforms/IndVarSimplify/strengthen-overflow.ll
+++ b/llvm/test/Transforms/IndVarSimplify/strengthen-overflow.ll
@@ -104,5 +104,89 @@ define i32 @test.unsigned.add.1(i32* %array, i32 %length, i32 %init) {
   ret i32 42
 }
 
+define hidden void @test.shl.exact.equal() {
+; CHECK-LABEL: @test.shl.exact.equal
+entry:
+  br label %for.body
+
+for.body:
+; CHECK-LABEL: for.body
+  %k.021 = phi i32 [ 1, %entry ], [ %inc, %for.body ]
+  %shl = shl i32 1, %k.021
+  %shr1 = ashr i32 %shl, 1
+; CHECK: %shr1 = ashr exact i32 %shl, 1
+  %shr2 = lshr i32 %shl, 1
+; CHECK: %shr2 = lshr exact i32 %shl, 1
+  %inc = add nuw nsw i32 %k.021, 1
+  %exitcond = icmp eq i32 %inc, 9
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+  ret void
+}
+
+define hidden void @test.shl.exact.greater() {
+; CHECK-LABEL: @test.shl.exact.greater
+entry:
+  br label %for.body
+
+for.body:
+; CHECK-LABEL: for.body
+  %k.021 = phi i32 [ 3, %entry ], [ %inc, %for.body ]
+  %shl = shl i32 1, %k.021
+  %shr1 = ashr i32 %shl, 2
+; CHECK: %shr1 = ashr exact i32 %shl, 2
+  %shr2 = lshr i32 %shl, 2
+; CHECK: %shr2 = lshr exact i32 %shl, 2
+  %inc = add nuw nsw i32 %k.021, 1
+  %exitcond = icmp eq i32 %inc, 9
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+  ret void
+}
+
+define hidden void @test.shl.exact.unbound(i32 %arg) {
+; CHECK-LABEL: @test.shl.exact.unbound
+entry:
+  br label %for.body
+
+for.body:
+; CHECK-LABEL: for.body
+  %k.021 = phi i32 [ 2, %entry ], [ %inc, %for.body ]
+  %shl = shl i32 1, %k.021
+  %shr1 = ashr i32 %shl, 2
+; CHECK: %shr1 = ashr exact i32 %shl, 2
+  %shr2 = lshr i32 %shl, 2
+; CHECK: %shr2 = lshr exact i32 %shl, 2
+  %inc = add nuw nsw i32 %k.021, 1
+  %exitcond = icmp eq i32 %inc, %arg
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+  ret void
+}
+
+define hidden void @test.shl.nonexact() {
+; CHECK-LABEL: @test.shl.nonexact
+entry:
+  br label %for.body
+
+for.body:
+; CHECK-LABEL: for.body
+  %k.021 = phi i32 [ 2, %entry ], [ %inc, %for.body ]
+  %shl = shl i32 1, %k.021
+  %shr1 = ashr i32 %shl, 3
+; CHECK: %shr1 = ashr i32 %shl, 3
+  %shr2 = lshr i32 %shl, 3
+; CHECK: %shr2 = lshr i32 %shl, 3
+  %inc = add nuw nsw i32 %k.021, 1
+  %exitcond = icmp eq i32 %inc, 9
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+  ret void
+}
+
 !0 = !{i32 0, i32 2}
 !1 = !{i32 0, i32 42}