[InstCombine] Teach select01 helper of foldSelectIntoOp to handle vector splats

We were handling some vectors in foldSelectIntoOp, but not if the operand of the bin op was any kind of vector constant. This patch fixes it to treat vector splats the same as scalars.

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

llvm-svn: 311940

3 files changed, 55 insertions, 9 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 5712fa4f09c..08aba886575 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -220,16 +220,15 @@ Instruction *InstCombiner::foldSelectOpOp(SelectInst &SI, Instruction *TI,
 }
 
 static bool isSelect01(Constant *C1, Constant *C2) {
-  ConstantInt *C1I = dyn_cast<ConstantInt>(C1);
-  if (!C1I)
+  const APInt *C1I, *C2I;
+  if (!match(C1, m_APInt(C1I)))
     return false;
-  ConstantInt *C2I = dyn_cast<ConstantInt>(C2);
-  if (!C2I)
+  if (!match(C2, m_APInt(C2I)))
     return false;
-  if (!C1I->isZero() && !C2I->isZero()) // One side must be zero.
+  if (!C1I->isNullValue() && !C2I->isNullValue()) // One side must be zero.
     return false;
-  return C1I->isOne() || C1I->isMinusOne() ||
-         C2I->isOne() || C2I->isMinusOne();
+  return C1I->isOneValue() || C1I->isAllOnesValue() ||
+         C2I->isOneValue() || C2I->isAllOnesValue();
 }
 
 /// Try to fold the select into one of the operands to allow further
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index 413efb8a4af..580e0998e0e 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -230,6 +230,17 @@ define i32 @test12(i1 %cond, i32 %a) {
 ; CHECK: ret i32 %c
 }
 
+define <2 x i32> @test12vec(<2 x i1> %cond, <2 x i32> %a) {
+; CHECK-LABEL: @test12vec(
+; CHECK-NEXT:    [[B:%.*]] = zext <2 x i1> [[COND:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[C:%.*]] = or <2 x i32> [[B]], [[A:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[C]]
+;
+  %b = or <2 x i32> %a, <i32 1, i32 1>
+  %c = select <2 x i1> %cond, <2 x i32> %b, <2 x i32> %a
+  ret <2 x i32> %c
+}
+
 define i32 @test12a(i1 %cond, i32 %a) {
         %b = ashr i32 %a, 1
         %c = select i1 %cond, i32 %b, i32 %a
@@ -240,6 +251,17 @@ define i32 @test12a(i1 %cond, i32 %a) {
 ; CHECK: ret i32 %c
 }
 
+define <2 x i32> @test12avec(<2 x i1> %cond, <2 x i32> %a) {
+; CHECK-LABEL: @test12avec(
+; CHECK-NEXT:    [[B:%.*]] = zext <2 x i1> [[COND:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[C:%.*]] = ashr <2 x i32> [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret <2 x i32> [[C]]
+;
+  %b = ashr <2 x i32> %a, <i32 1, i32 1>
+  %c = select <2 x i1> %cond, <2 x i32> %b, <2 x i32> %a
+  ret <2 x i32> %c
+}
+
 define i32 @test12b(i1 %cond, i32 %a) {
 ; CHECK-LABEL: @test12b(
 ; CHECK-NEXT:    [[NOT_COND:%.*]] = xor i1 %cond, true
@@ -252,6 +274,18 @@ define i32 @test12b(i1 %cond, i32 %a) {
   ret i32 %d
 }
 
+define <2 x i32> @test12bvec(<2 x i1> %cond, <2 x i32> %a) {
+; CHECK-LABEL: @test12bvec(
+; CHECK-NEXT:    [[NOT_COND:%.*]] = xor <2 x i1> [[COND:%.*]], <i1 true, i1 true>
+; CHECK-NEXT:    [[B:%.*]] = zext <2 x i1> [[NOT_COND]] to <2 x i32>
+; CHECK-NEXT:    [[D:%.*]] = ashr <2 x i32> [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret <2 x i32> [[D]]
+;
+  %b = ashr <2 x i32> %a, <i32 1, i32 1>
+  %d = select <2 x i1> %cond, <2 x i32> %a, <2 x i32> %b
+  ret <2 x i32> %d
+}
+
 define i32 @test13(i32 %a, i32 %b) {
         %C = icmp eq i32 %a, %b
         %V = select i1 %C, i32 %a, i32 %b
@@ -772,6 +806,19 @@ define i32 @test42(i32 %x, i32 %y) {
 ; CHECK-NEXT: ret i32 %c
 }
 
+define <2 x i32> @test42vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @test42vec(
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq <2 x i32> [[X:%.*]], zeroinitializer
+; CHECK-NEXT:    [[B:%.*]] = sext <2 x i1> [[COND]] to <2 x i32>
+; CHECK-NEXT:    [[C:%.*]] = add <2 x i32> [[B]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[C]]
+;
+  %b = add <2 x i32> %y, <i32 -1, i32 -1>
+  %cond = icmp eq <2 x i32> %x, zeroinitializer
+  %c = select <2 x i1> %cond, <2 x i32> %b, <2 x i32> %y
+  ret <2 x i32> %c
+}
+
 ; PR8994
 
 ; This select instruction can't be eliminated because trying to do so would
diff --git a/llvm/test/Transforms/LoopVectorize/phi-cost.ll b/llvm/test/Transforms/LoopVectorize/phi-cost.ll
index 5ccea66c76a..057c55d55d3 100644
--- a/llvm/test/Transforms/LoopVectorize/phi-cost.ll
+++ b/llvm/test/Transforms/LoopVectorize/phi-cost.ll
@@ -10,8 +10,8 @@ target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[INDEX_NEXT:%.*]], %vector.body ]
 ; CHECK:         [[WIDE_LOAD:%.*]] = load <2 x i32>, <2 x i32>* {{.*}}
 ; CHECK:         [[TMP5:%.*]] = icmp sgt <2 x i32> [[WIDE_LOAD]], zeroinitializer
-; CHECK-NEXT:    [[TMP6:%.*]] = add <2 x i32> [[WIDE_LOAD]], <i32 1, i32 1>
-; CHECK-NEXT:    [[PREDPHI:%.*]] = select <2 x i1> [[TMP5]], <2 x i32> [[TMP6]], <2 x i32> [[WIDE_LOAD]]
+; CHECK-NEXT:    [[TMP6:%.*]] = zext <2 x i1> [[TMP5]] to <2 x i32>
+; CHECK-NEXT:    [[PREDPHI:%.*]] = add <2 x i32> [[WIDE_LOAD]], [[TMP6]]
 ; CHECK:         store <2 x i32> [[PREDPHI]], <2 x i32>* {{.*}}
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add i64 [[INDEX]], 2
 ;