2 files changed, 35 insertions, 2 deletions
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index bb4220d6164..1171149180b 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InstructionSimplify.h"
@@ -1004,9 +1005,18 @@ static void computeKnownBitsFromShiftOperator(Operator *I,
   // calculation. Reusing the APInts here to prevent unnecessary allocations.
   KnownZero.clearAllBits(), KnownOne.clearAllBits();
 
+  // If we know the shifter operand is nonzero, we can sometimes infer more
+  // known bits. However this is expensive to compute, so be lazy about it and
+  // only compute it when absolutely necessary.
+  Optional<bool> ShifterOperandIsNonZero;
+
   // Early exit if we can't constrain any well-defined shift amount.
-  if (!(ShiftAmtKZ & (BitWidth-1)) && !(ShiftAmtKO & (BitWidth-1)))
-    return;
+  if (!(ShiftAmtKZ & (BitWidth - 1)) && !(ShiftAmtKO & (BitWidth - 1))) {
+    ShifterOperandIsNonZero =
+        isKnownNonZero(I->getOperand(1), DL, Depth + 1, Q);
+    if (!*ShifterOperandIsNonZero)
+      return;
+  }
 
   computeKnownBits(I->getOperand(0), KnownZero2, KnownOne2, DL, Depth + 1, Q);
 
@@ -1018,6 +1028,16 @@ static void computeKnownBitsFromShiftOperator(Operator *I,
       continue;
     if ((ShiftAmt | ShiftAmtKO) != ShiftAmt)
       continue;
+    // If we know the shifter is nonzero, we may be able to infer more known
+    // bits. This check is sunk down as far as possible to avoid the expensive
+    // call to isKnownNonZero if the cheaper checks above fail.
+    if (ShiftAmt == 0) {
+      if (!ShifterOperandIsNonZero.hasValue())
+        ShifterOperandIsNonZero =
+            isKnownNonZero(I->getOperand(1), DL, Depth + 1, Q);
+      if (*ShifterOperandIsNonZero)
+        continue;
+    }
 
     KnownZero &= KZF(KnownZero2, ShiftAmt);
     KnownOne  &= KOF(KnownOne2, ShiftAmt);
diff --git a/llvm/test/Analysis/ValueTracking/knownnonzero-shift.ll b/llvm/test/Analysis/ValueTracking/knownnonzero-shift.ll
new file mode 100644
index 00000000000..e59d19cc2e2
--- /dev/null
+++ b/llvm/test/Analysis/ValueTracking/knownnonzero-shift.ll
@@ -0,0 +1,13 @@
+; RUN: opt -instsimplify -S < %s | FileCheck %s
+
+; CHECK-LABEL: @test
+define i1 @test(i8 %p, i8* %pq) {
+  %q = load i8, i8* %pq, !range !0 ; %q is known nonzero; no known bits
+  %1 = shl i8 %p, %q              ; because %q is nonzero, %1[0] is known to be zero.
+  %2 = and i8 %1, 1
+  %x = icmp eq i8 %2, 0
+  ; CHECK: ret i1 true
+  ret i1 %x
+}
+
+!0 = !{ i8 1, i8 5 }