[InstCombine] use m_APInt to allow icmp (shr exact X, Y), 0 folds for splat constant vectors

llvm-svn: 279472

2 files changed, 15 insertions, 18 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index c32821c73f8..68e55c4a66c 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -2061,25 +2061,24 @@ Instruction *InstCombiner::foldICmpShlConstant(ICmpInst &Cmp, Instruction *Shl,
   return nullptr;
 }
 
+/// Fold icmp ({al}shr X, Y), C.
 Instruction *InstCombiner::foldICmpShrConstant(ICmpInst &ICI, Instruction *LHSI,
                                                const APInt *RHSV) {
-  // FIXME: This check restricts all folds under here to scalar types.
-  ConstantInt *RHS = dyn_cast<ConstantInt>(ICI.getOperand(1));
-  if (!RHS)
-    return nullptr;
+  // An exact shr only shifts out zero bits, so:
+  // icmp eq/ne (shr X, Y), 0 --> icmp eq/ne X, 0
+  CmpInst::Predicate Pred = ICI.getPredicate();
+  BinaryOperator *BO = cast<BinaryOperator>(LHSI);
+  if (ICI.isEquality() && BO->isExact() && BO->hasOneUse() && *RHSV == 0)
+    return new ICmpInst(Pred, BO->getOperand(0), ICI.getOperand(1));
 
+  // FIXME: This check restricts all folds under here to scalar types.
   // Handle equality comparisons of shift-by-constant.
-  BinaryOperator *BO = cast<BinaryOperator>(LHSI);
-  if (ConstantInt *ShAmt = dyn_cast<ConstantInt>(LHSI->getOperand(1))) {
-    if (Instruction *Res = foldICmpShrConstConst(ICI, BO, ShAmt))
-      return Res;
-  }
+  ConstantInt *ShAmt = dyn_cast<ConstantInt>(LHSI->getOperand(1));
+  if (!ShAmt)
+    return nullptr;
 
-  // Handle exact shr's.
-  if (ICI.isEquality() && BO->isExact() && BO->hasOneUse()) {
-    if (RHSV->isMinValue())
-      return new ICmpInst(ICI.getPredicate(), BO->getOperand(0), RHS);
-  }
+  if (Instruction *Res = foldICmpShrConstConst(ICI, BO, ShAmt))
+    return Res;
 
   return nullptr;
 }
diff --git a/llvm/test/Transforms/InstCombine/icmp.ll b/llvm/test/Transforms/InstCombine/icmp.ll
index 884ddbb2a72..780529c9de6 100644
--- a/llvm/test/Transforms/InstCombine/icmp.ll
+++ b/llvm/test/Transforms/InstCombine/icmp.ll
@@ -490,8 +490,7 @@ define i1 @test39(i32 %X, i32 %Y) {
 
 define <2 x i1> @test39vec(<2 x i32> %X, <2 x i32> %Y) {
 ; CHECK-LABEL: @test39vec(
-; CHECK-NEXT:    [[A:%.*]] = ashr exact <2 x i32> %X, %Y
-; CHECK-NEXT:    [[B:%.*]] = icmp eq <2 x i32> [[A]], zeroinitializer
+; CHECK-NEXT:    [[B:%.*]] = icmp eq <2 x i32> %X, zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[B]]
 ;
   %A = ashr exact <2 x i32> %X, %Y
@@ -511,8 +510,7 @@ define i1 @test40(i32 %X, i32 %Y) {
 
 define <2 x i1> @test40vec(<2 x i32> %X, <2 x i32> %Y) {
 ; CHECK-LABEL: @test40vec(
-; CHECK-NEXT:    [[A:%.*]] = lshr exact <2 x i32> %X, %Y
-; CHECK-NEXT:    [[B:%.*]] = icmp ne <2 x i32> [[A]], zeroinitializer
+; CHECK-NEXT:    [[B:%.*]] = icmp ne <2 x i32> %X, zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[B]]
 ;
   %A = lshr exact <2 x i32> %X, %Y